Changing AgentState.response -> AgentState.messages

Author: splunk-dtaborski

splunklib/ai/engines/langchain.py

@@ -279,9 +279,13 @@ async def awrap_tool_call(
                     assert resp.artifact is None, "artifact is already populated"
 
                     if resp.name.startswith(AGENT_PREFIX):
-                        resp.artifact = SubagentFailureResult(str(resp.content))  # pyright: ignore[reportUnknownArgumentType]
+                        resp.artifact = SubagentFailureResult(
+                            str(resp.content)
+                        )  # pyright: ignore[reportUnknownArgumentType]
                     else:
-                        resp.artifact = ToolFailureResult(str(resp.content))  # pyright: ignore[reportUnknownArgumentType]
+                        resp.artifact = ToolFailureResult(
+                            str(resp.content)
+                        )  # pyright: ignore[reportUnknownArgumentType]
 
                 return resp
 
@@ -967,9 +971,9 @@ async def _sdk_handler(request: ToolRequest) -> ToolResponse:
         lc_request = _convert_tool_request_to_lc(request, original_request)
         result = await handler(lc_request)
         sdk_result = _convert_tool_message_from_lc(result)
-        assert isinstance(sdk_result, ToolMessage), (
-            "Expected tool response from tool middleware handler"
-        )
+        assert isinstance(
+            sdk_result, ToolMessage
+        ), "Expected tool response from tool middleware handler"
         return ToolResponse(sdk_result.result)
 
     return _sdk_handler
@@ -987,9 +991,9 @@ async def _sdk_handler(
         lc_request = _convert_subagent_request_to_lc(request, original_request)
         result = await handler(lc_request)
         sdk_result = _convert_tool_message_from_lc(result)
-        assert isinstance(sdk_result, SubagentMessage), (
-            "Expected subagent response from subagent middleware handler"
-        )
+        assert isinstance(
+            sdk_result, SubagentMessage
+        ), "Expected subagent response from subagent middleware handler"
         return SubagentResponse(sdk_result.result)
 
     return _sdk_handler
@@ -1182,16 +1186,18 @@ def _convert_tool_message_from_lc(
             )
         case LC_ToolMessage():
             # If this is reached, we likely passed an invalid tool name to LangChain.
-            assert message.name is not None, (
-                "LangChain responded with a nameless tool call"
-            )
+            assert (
+                message.name is not None
+            ), "LangChain responded with a nameless tool call"
 
             if message.name.startswith(TOOL_STRATEGY_TOOL_PREFIX):
                 return StructuredOutputMessage(
                     name=message.name.removeprefix(TOOL_STRATEGY_TOOL_PREFIX),
                     call_id=message.tool_call_id,
                     status=message.status,
-                    content=str(message.content),  # pyright: ignore[reportUnknownArgumentType]
+                    content=str(
+                        message.content
+                    ),  # pyright: ignore[reportUnknownArgumentType]
                 )
 
             assert isinstance(message.artifact, ToolResult) or isinstance(
@@ -1266,7 +1272,7 @@ def _convert_model_result_from_lc(model_response: LC_ModelCallResult) -> ModelRe
 
 
 def _convert_agent_state_to_lc(state: AgentState) -> LC_AgentState[Any]:
-    messages = [_map_message_to_langchain(m) for m in state.response.messages]
+    messages = [_map_message_to_langchain(m) for m in state.messages]
     return LC_AgentState(messages=messages)
 
 
@@ -1351,7 +1357,9 @@ async def _tool_call(
         except ToolException as e:
             raise LC_ToolException(*e.args) from e
         except LC_ToolException:
-            assert False, (  # noqa: PT015
+            assert (
+                False
+            ), (  # noqa: PT015
                 "ToolException from LangChain should not be raised in tool.func"
             )
 
@@ -1454,6 +1462,7 @@ async def _run(  # pyright: ignore[reportRedeclaration]
                 content: str, thread_id: str
             ) -> tuple[OutputT | str, SubagentStructuredResult | SubagentTextResult]:
                 return await invoke_agent(HumanMessage(content=content), thread_id)
+
         else:
 
             async def _run(  # pyright: ignore[reportRedeclaration]
@@ -1627,14 +1636,9 @@ def _convert_agent_state_from_langchain(
     messages = state["messages"]
     total_tokens_counter = _get_approximate_token_counter(model)
     total_tokens = total_tokens_counter(messages)
-
-    response = AgentResponse[Any | None](
-        messages=[_map_message_from_langchain(m) for m in state["messages"]],
-        structured_output=state.get("structured_response"),
-    )
-
+    messages = [_map_message_from_langchain(m) for m in state["messages"]]
     return AgentState(
-        response=response,
+        messages=messages,
         total_steps=len(messages),
         token_count=total_tokens,
     )
@@ -1646,7 +1650,9 @@ def _get_approximate_token_counter(model: BaseChatModel) -> LC_TokenCounter:
     # NOTE: This is adapted from the backend provider library
     # 3.3 was estimated in an offline experiment, comparing with Claude's token-counting
     # API: https://platform.claude.com/docs/en/build-with-claude/token-counting
-    if model._llm_type == ANTHROPIC_CHAT_MODEL_TYPE:  # pyright: ignore[reportPrivateUsage]
+    if (
+        model._llm_type == ANTHROPIC_CHAT_MODEL_TYPE
+    ):  # pyright: ignore[reportPrivateUsage]
         return partial(count_tokens_approximately, chars_per_token=3.3)
     return count_tokens_approximately
 

splunklib/ai/middleware.py

@@ -12,7 +12,7 @@
 # License for the specific language governing permissions and limitations
 # under the License.
 
-from collections.abc import Awaitable, Callable
+from collections.abc import Sequence, Awaitable, Callable
 from dataclasses import dataclass
 from typing import Any, override
 
@@ -35,7 +35,7 @@ class AgentState:
     """AgentState is available through certain middlewares and contains information about the current state of an agent execution."""
 
     # holds messages exchanged so far in the conversation
-    response: AgentResponse[Any | None]
+    messages: Sequence[BaseMessage]
     # steps taken so far in the conversation
     total_steps: int
     # tokens used so far in the conversation
@@ -96,7 +96,7 @@ def __post_init__(self) -> None:
 
 @dataclass(frozen=True)
 class AgentRequest:
-    messages: list[BaseMessage]
+    messages: Sequence[BaseMessage]
 
 
 AgentMiddlewareHandler = Callable[[AgentRequest], Awaitable[AgentResponse[Any | None]]]

tests/integration/ai/test_agent.py

@@ -65,9 +65,9 @@ async def test_agent_with_openai_round_trip(self):
             )
 
             response = result.final_message.content.strip().lower().replace(".", "")
-            assert result.structured_output is None, (
-                "The structured output should not be populated"
-            )
+            assert (
+                result.structured_output is None
+            ), "The structured output should not be populated"
             assert "stefan" in response
 
     @pytest.mark.asyncio
@@ -160,9 +160,9 @@ class Person(BaseModel):
 
             # check if the last message contains the response in natural language
             assert response.name in last_message, "Name field not found in the message"
-            assert str(response.age) in last_message, (
-                "Age field not found in the message"
-            )
+            assert (
+                str(response.age) in last_message
+            ), "Age field not found in the message"
 
     async def test_agent_uses_subagent(self):
         pytest.importorskip("langchain_openai")
@@ -215,9 +215,9 @@ class NicknameGeneratorInput(BaseModel):
             subagent_message = next(
                 filter(lambda m: m.role == "subagent", result.messages), None
             )
-            assert isinstance(subagent_message, SubagentMessage), (
-                "Invalid subagent message"
-            )
+            assert isinstance(
+                subagent_message, SubagentMessage
+            ), "Invalid subagent message"
             assert subagent_message, "No subagent message found in response"
 
             response = result.final_message.content
@@ -366,12 +366,12 @@ class SupervisorOutput(BaseModel):
                 )
 
                 response = result.structured_output
-                assert type(response) == SupervisorOutput, (
-                    "Response is not of type Team"
-                )
-                assert len(response.member_descriptions) == 3, (
-                    "Team does not have 3 members"
-                )
+                assert (
+                    type(response) == SupervisorOutput
+                ), "Response is not of type Team"
+                assert (
+                    len(response.member_descriptions) == 3
+                ), "Team does not have 3 members"
 
     @pytest.mark.asyncio
     async def test_duplicated_subagent_name(self) -> None:
@@ -520,9 +520,9 @@ async def _subagent_call_middleware(
 
             # Override the arguments, such that are invalid.
             resp = await handler(replace(request, call=replace(request.call, args={})))
-            assert isinstance(resp.result, SubagentFailureResult), (
-                "subagent call did not fail"
-            )
+            assert isinstance(
+                resp.result, SubagentFailureResult
+            ), "subagent call did not fail"
 
             after_subagent_call = True
             return resp
@@ -532,7 +532,7 @@ async def _model_call_middleware(
             req: ModelRequest, _handler: ModelMiddlewareHandler
         ) -> ModelResponse:
             if after_subagent_call:
-                msgs = req.state.response.messages
+                msgs = req.state.messages
                 assert isinstance(msgs[-1], SubagentMessage)
                 assert isinstance(msgs[-1].result, SubagentFailureResult)
 

tests/integration/ai/test_conversation_store.py

@@ -66,9 +66,9 @@ async def _model_middleware(
 
             if after_first_call:
                 # Previous messages included.
-                assert len(request.state.response.messages) == 3
+                assert len(request.state.messages) == 3
             else:
-                assert len(request.state.response.messages) == 1
+                assert len(request.state.messages) == 1
             return await handler(request)
 
         @agent_middleware
@@ -166,7 +166,7 @@ async def _model_middleware(
             nonlocal model_middleware_called
             model_middleware_called = True
 
-            assert len(request.state.response.messages) == 1
+            assert len(request.state.messages) == 1
             return await handler(request)
 
         async with Agent(
@@ -186,9 +186,9 @@ async def _model_middleware(
                 thread_id="2",
             )
             response = result.final_message.content
-            assert "Mike" not in response, (
-                "Agent remembered the name from a different thread_id"
-            )
+            assert (
+                "Mike" not in response
+            ), "Agent remembered the name from a different thread_id"
 
         assert model_middleware_called
 
@@ -276,9 +276,9 @@ async def _model_middleware(
             nonlocal after_first_call
 
             if after_first_call:
-                assert len(request.state.response.messages) == 3
+                assert len(request.state.messages) == 3
             else:
-                assert len(request.state.response.messages) == 1
+                assert len(request.state.messages) == 1
 
             after_first_call = True
             return await handler(request)
@@ -347,9 +347,9 @@ async def _model_middleware(
             nonlocal after_first_call
 
             if after_first_call:
-                assert len(request.state.response.messages) == 3
+                assert len(request.state.messages) == 3
             else:
-                assert len(request.state.response.messages) == 1
+                assert len(request.state.messages) == 1
 
             after_first_call = True
             return await handler(request)

tests/integration/ai/test_hooks.py

@@ -30,7 +30,13 @@
     before_model,
 )
 from splunklib.ai.messages import AIMessage, AgentResponse, HumanMessage
-from splunklib.ai.middleware import AgentRequest, ModelMiddlewareHandler, ModelRequest, ModelResponse, model_middleware
+from splunklib.ai.middleware import (
+    AgentRequest,
+    ModelMiddlewareHandler,
+    ModelRequest,
+    ModelResponse,
+    model_middleware,
+)
 from tests.ai_testlib import AITestCase
 
 
@@ -47,15 +53,15 @@ def test_hook_before(req: ModelRequest) -> None:
             hook_calls += 1
 
             assert req.system_message.startswith("Your name is stefan")
-            assert len(req.state.response.messages) == 1
+            assert len(req.state.messages) == 1
 
         @before_model
         async def test_async_hook_before(req: ModelRequest) -> None:
             nonlocal hook_calls
             hook_calls += 1
 
             assert req.system_message.startswith("Your name is stefan")
-            assert len(req.state.response.messages) == 1
+            assert len(req.state.messages) == 1
 
         @after_model
         def test_hook_after(resp: ModelResponse) -> None:
@@ -197,10 +203,12 @@ async def test_agent_loop_stop_conditions_conversation_limit(self) -> None:
             with pytest.raises(
                 StepsLimitExceededException, match="Steps limit of 2 exceeded"
             ):
-                _ = await agent.invoke([
-                    HumanMessage(content="hi, my name is Chris"),
-                    HumanMessage(content="What is my name?"),
-                ])
+                _ = await agent.invoke(
+                    [
+                        HumanMessage(content="hi, my name is Chris"),
+                        HumanMessage(content="What is my name?"),
+                    ]
+                )
 
     @pytest.mark.asyncio
     async def test_agent_loop_stop_conditions_conversation_limit_with_checkpointer(
@@ -220,13 +228,17 @@ async def test_agent_loop_stop_conditions_conversation_limit_with_checkpointer(
             with pytest.raises(
                 StepsLimitExceededException, match="Steps limit of 2 exceeded"
             ):
-                _ = await agent.invoke([
-                    HumanMessage(content="What is my name?"),
-                    HumanMessage(content="Are you sure?"),
-                ])
+                _ = await agent.invoke(
+                    [
+                        HumanMessage(content="What is my name?"),
+                        HumanMessage(content="Are you sure?"),
+                    ]
+                )
 
     @pytest.mark.asyncio
-    async def test_agent_loop_stop_conditions_steps_accumulate_across_invokes(self) -> None:
+    async def test_agent_loop_stop_conditions_steps_accumulate_across_invokes(
+        self,
+    ) -> None:
         pytest.importorskip("langchain_openai")
 
         step_limit = StepLimitMiddleware(2)

tests/integration/ai/test_middleware.py

@@ -78,7 +78,7 @@ async def test_middleware(
             assert call.args == {"city": "Krakow"}
 
             state = request.state
-            assert len(state.response.messages) == 2
+            assert len(state.messages) == 2
 
             response = await handler(request)
             assert isinstance(response.result, ToolResult)
@@ -500,9 +500,9 @@ async def test_middleware(
             )
             assert subagent_message, "SubagentMessage not found in messages"
             assert isinstance(subagent_message.result, SubagentTextResult)
-            assert subagent_message.result.content == "Chris-superstar", (
-                "Invalid response from subagent"
-            )
+            assert (
+                subagent_message.result.content == "Chris-superstar"
+            ), "Invalid response from subagent"
             assert middleware_called, "Middleware was not called"
 
     @pytest.mark.asyncio
@@ -699,10 +699,7 @@ async def mutating_middleware(
         ) -> ModelResponse:
             new_state = replace(
                 request.state,
-                response=replace(
-                    request.state.response,
-                    messages=[HumanMessage(content="What is the capital of France?")],
-                ),
+                messages=[HumanMessage(content="What is the capital of France?")],
             )
             return await handler(replace(request, state=new_state))