nits from review

Author: sydney-runkle

libs/langchain_v1/langchain/agents/middleware/types.py

@@ -230,10 +230,13 @@ def override(self, **overrides: Unpack[_ModelRequestOverrides]) -> ModelRequest:
                 )
                 ```
 
-            !!! example "Override system prompt (backward compatible)"
+            !!! example "Override multiple attributes"
 
                 ```python
-                new_request = request.override(system_prompt="New instructions")
+                new_request = request.override(
+                    model=ChatOpenAI(model="gpt-4o"),
+                    system_message=SystemMessage(content="New instructions"),
+                )
                 ```
         """
         # Handle system_prompt/system_message conversion
@@ -695,7 +698,7 @@ def __call__(
         ...
 
 
-class _CallableReturningPromptString(Protocol[StateT_contra, ContextT]):  # type: ignore[misc]
+class _CallableReturningSystemMessage(Protocol[StateT_contra, ContextT]):  # type: ignore[misc]
     """Callable that returns a prompt string or SystemMessage given `ModelRequest`."""
 
     def __call__(
@@ -1382,24 +1385,24 @@ def wrapped(
 
 @overload
 def dynamic_prompt(
-    func: _CallableReturningPromptString[StateT, ContextT],
+    func: _CallableReturningSystemMessage[StateT, ContextT],
 ) -> AgentMiddleware[StateT, ContextT]: ...
 
 
 @overload
 def dynamic_prompt(
     func: None = None,
 ) -> Callable[
-    [_CallableReturningPromptString[StateT, ContextT]],
+    [_CallableReturningSystemMessage[StateT, ContextT]],
     AgentMiddleware[StateT, ContextT],
 ]: ...
 
 
 def dynamic_prompt(
-    func: _CallableReturningPromptString[StateT, ContextT] | None = None,
+    func: _CallableReturningSystemMessage[StateT, ContextT] | None = None,
 ) -> (
     Callable[
-        [_CallableReturningPromptString[StateT, ContextT]],
+        [_CallableReturningSystemMessage[StateT, ContextT]],
         AgentMiddleware[StateT, ContextT],
     ]
     | AgentMiddleware[StateT, ContextT]
@@ -1453,7 +1456,7 @@ def context_aware_prompt(request: ModelRequest) -> str:
     """
 
     def decorator(
-        func: _CallableReturningPromptString[StateT, ContextT],
+        func: _CallableReturningSystemMessage[StateT, ContextT],
     ) -> AgentMiddleware[StateT, ContextT]:
         is_async = iscoroutinefunction(func)
 
@@ -1488,7 +1491,7 @@ def wrapped(
             request: ModelRequest,
             handler: Callable[[ModelRequest], ModelResponse],
         ) -> ModelCallResult:
-            prompt = func(request)
+            prompt = cast("Callable[[ModelRequest], SystemMessage | str]", func)(request)
             if isinstance(prompt, SystemMessage):
                 request = request.override(system_message=prompt)
             else:
@@ -1501,7 +1504,7 @@ async def async_wrapped_from_sync(
             handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
         ) -> ModelCallResult:
             # Delegate to sync function
-            prompt = func(request)
+            prompt = cast("Callable[[ModelRequest], SystemMessage | str]", func)(request)
             if isinstance(prompt, SystemMessage):
                 request = request.override(system_message=prompt)
             else: