[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,48 @@
+# 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:
+    """
+    Abstract base class for parsing agent responses. Implement this class to customize how
+    agent outputs are processed and transformed.
+
+    This class allows developers to define custom parsing logic for agent responses,
+    which can be useful for:
+    - Extracting specific information from the response
+    - Formatting or structuring the output in a specific way
+    - Validating or sanitizing the agent's response
+
+    To use this class:
+    1. Create a subclass of OutputParser
+    2. Implement the `parse` method
+    3. Pass your parser instance to the Agent's constructor
+
+    Example:
+        class MyCustomParser(OutputParser):
+            def parse(self, output_message: CompletionMessage) -> CompletionMessage:
+                # Add your custom parsing logic here
+                return processed_message
+
+    Methods:
+        parse(output_message: CompletionMessage) -> CompletionMessage:
+            Abstract method that must be implemented by subclasses to process
+            the agent's response.
+
+            Args:
+                output_message (CompletionMessage): The response message from agent turn
+
+            Returns:
+                CompletionMessage: The processed/transformed response message
+    """
+
+    @abstractmethod
+    def parse(self, output_message: CompletionMessage) -> CompletionMessage:
+        raise NotImplementedError