from langchain.agents.middleware import wrap_model_call, AgentState

class MyState(AgentState):
    pass

@wrap_model_call(state_schema=MyState)
async def add_cache_control(request, handler):
    return await handler(request)

---

error[invalid-argument-type]: Argument is incorrect
  --> src/graph.py:101:1
    |
101 | @wrap_model_call(state_schema=MyState)
    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Expected `_CallableReturningModelResponse[MyState, None]`,
    |                                         found `def add_cache_control(request, handler) -> CoroutineType[Any, Any, Unknown]`

---

class _CallableReturningModelResponse(Protocol[StateT_contra, ContextT]):
    def __call__(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], ModelResponse],
    ) -> ModelCallResult: ...

---

class _AsyncCallableReturningModelResponse(Protocol[StateT_contra, ContextT]):
    async def __call__(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
    ) -> ModelCallResult: ...

@overload
def wrap_model_call(
    func: _AsyncCallableReturningModelResponse[StateT, ContextT],
) -> AgentMiddleware[StateT, ContextT]: ...

@overload
def wrap_model_call(
    func: None = None,
    *,
    state_schema: type[StateT] | None = None,
    tools: list[BaseTool] | None = None,
    name: str | None = None,
) -> Callable[
    [_CallableReturningModelResponse[StateT, ContextT] | _AsyncCallableReturningModelResponse[StateT, ContextT]],
    AgentMiddleware[StateT, ContextT],
]: ...

---

langchain==1.0.0a14
Python 3.12
ty 0.0.1a12 (Astral)