from collections.abc import Awaitable, Callable

from langchain.agents.middleware import AgentMiddleware, ModelRequest, ModelResponse
from langchain.messages import HumanMessage, ToolMessage


class RequireResearchMiddleware(AgentMiddleware):
    def __init__(self, tool_name: str = "research") -> None:
        self.tool_name = tool_name

    @staticmethod
    def _has_current_tool(messages: list, tool_name: str) -> bool:
        for message in reversed(messages):
            if isinstance(message, ToolMessage) and message.name == tool_name:
                return True
            if isinstance(message, HumanMessage):
                return False
        return False

    async def awrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
    ) -> ModelResponse:
        has_tool_result = self._has_current_tool(request.messages, self.tool_name)
        tool_choice = "none" if has_tool_result else "required"
        return await handler(request.override(tool_choice=tool_choice))
