[RFC] Client Agent SDK OutputParser

src/llama_stack_client/lib/agents/agent.py

@@ -12,7 +12,9 @@
 from llama_stack_client.types.agents.turn_create_params import Document, Toolgroup
 from llama_stack_client.types.agents.turn_create_response import AgentTurnResponseStreamChunk
 
+
 from .client_tool import ClientTool
+from .output_parser import OutputParser
 
 DEFAULT_MAX_ITER = 10
 
@@ -24,13 +26,15 @@ def __init__(
         agent_config: AgentConfig,
         client_tools: Tuple[ClientTool] = (),
         memory_bank_id: Optional[str] = None,
+        output_parser: Optional[OutputParser] = None,
     ):
         self.client = client
         self.agent_config = agent_config
         self.agent_id = self._create_agent(agent_config)
         self.client_tools = {t.get_name(): t for t in client_tools}
         self.sessions = []
         self.memory_bank_id = memory_bank_id
+        self.output_parser = output_parser
 
     def _create_agent(self, agent_config: AgentConfig) -> int:
         agentic_system_create_response = self.client.agents.create(
@@ -54,6 +58,11 @@ def _has_tool_call(self, chunk: AgentTurnResponseStreamChunk) -> bool:
         message = chunk.event.payload.turn.output_message
         if message.stop_reason == "out_of_tokens":
             return False
+
+        if self.output_parser:
+            parsed_message = self.output_parser.parse(message)
+            message = parsed_message
+
         return len(message.tool_calls) > 0
 
     def _run_tool(self, chunk: AgentTurnResponseStreamChunk) -> ToolResponseMessage:
@@ -64,7 +73,7 @@ def _run_tool(self, chunk: AgentTurnResponseStreamChunk) -> ToolResponseMessage:
                 call_id=tool_call.call_id,
                 tool_name=tool_call.tool_name,
                 content=f"Unknown tool `{tool_call.tool_name}` was called.",
-                role="ipython",
+                role="tool",
             )
         tool = self.client_tools[tool_call.tool_name]
         result_messages = tool.run([message])

src/llama_stack_client/lib/agents/output_parser.py

@@ -0,0 +1,22 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from abc import abstractmethod
+
+from llama_stack_client.types.agents.turn import CompletionMessage
+
+
+class OutputParser:
+    """
+    Developers can define their own response output parser to parse the response from the agent turn.
+
+    Developers need to implement the `parse` method to parse the response from the agent turn.
+    The return value should be a CompletionMessage object with parsed result popoulated in content and tool_calls.
+    """
+
+    @abstractmethod
+    def parse(self, output_message: CompletionMessage) -> CompletionMessage:
+        raise NotImplementedError