massgen 0.0.3__py3-none-any.whl → 0.1.0__py3-none-any.whl

massgen/chat_agent.py

@@ -1,16 +1,21 @@
+# -*- coding: utf-8 -*-
 """
 Common chat interface for MassGen agents.
 
 Defines the standard interface that both individual agents and the orchestrator implement,
 allowing seamless interaction regardless of whether you're talking to a single agent
 or a coordinated multi-agent system.
+
+# TODO: Consider how to best handle stateful vs stateless backends in this interface.
 """
 
 import uuid
 from abc import ABC, abstractmethod
-from typing import Dict, List, Optional, Any, AsyncGenerator
+from typing import Any, AsyncGenerator, Dict, List, Optional
 
 from .backend.base import LLMBackend, StreamChunk
+from .stream_chunk import ChunkType
+from .utils import CoordinationStage
 
 
 class ChatAgent(ABC):
@@ -32,6 +37,7 @@ class ChatAgent(ABC):
         tools: List[Dict[str, Any]] = None,
         reset_chat: bool = False,
         clear_history: bool = False,
+        current_stage: CoordinationStage = None,
     ) -> AsyncGenerator[StreamChunk, None]:
         """
         Enhanced chat interface supporting tool calls and responses.
@@ -45,11 +51,11 @@ class ChatAgent(ABC):
             tools: Optional tools to provide to the agent
             reset_chat: If True, reset the agent's conversation history to the provided messages
             clear_history: If True, clear history but keep system message before processing messages
+            current_stage: Optional current coordination stage for orchestrator use
 
         Yields:
             StreamChunk: Streaming response chunks
         """
-        pass
 
     async def chat_simple(self, user_message: str) -> AsyncGenerator[StreamChunk, None]:
         """
@@ -68,12 +74,22 @@ class ChatAgent(ABC):
     @abstractmethod
     def get_status(self) -> Dict[str, Any]:
         """Get current agent status and state."""
-        pass
 
     @abstractmethod
-    def reset(self) -> None:
+    async def reset(self) -> None:
         """Reset agent state for new conversation."""
-        pass
+
+    @abstractmethod
+    def get_configurable_system_message(self) -> Optional[str]:
+        """
+        Get the user-configurable part of the system message.
+
+        Returns the domain expertise, role definition, or custom instructions
+        that were configured for this agent, without backend-specific details.
+
+        Returns:
+            The configurable system message if available, None otherwise
+        """
 
     # Common conversation management
     def get_conversation_history(self) -> List[Dict[str, Any]]:
@@ -133,13 +149,27 @@ class SingleAgent(ChatAgent):
 
         # Add system message to history if provided
         if self.system_message:
-            self.conversation_history.append(
-                {"role": "system", "content": self.system_message}
-            )
+            self.conversation_history.append({"role": "system", "content": self.system_message})
 
-    async def _process_stream(
-        self, backend_stream, tools: List[Dict[str, Any]] = None
-    ) -> AsyncGenerator[StreamChunk, None]:
+    @staticmethod
+    def _get_chunk_type_value(chunk) -> str:
+        """
+        Extract chunk type as string, handling both legacy and typed chunks.
+
+        Args:
+            chunk: StreamChunk, TextStreamChunk, or MultimodalStreamChunk
+
+        Returns:
+            String representation of chunk type (e.g., "content", "tool_calls")
+        """
+        chunk_type = chunk.type
+
+        if isinstance(chunk_type, ChunkType):
+            return chunk_type.value
+
+        return str(chunk_type)
+
+    async def _process_stream(self, backend_stream, tools: List[Dict[str, Any]] = None) -> AsyncGenerator[StreamChunk, None]:
         """Common streaming logic for processing backend responses."""
         assistant_response = ""
         tool_calls = []
@@ -147,50 +177,41 @@ class SingleAgent(ChatAgent):
 
         try:
             async for chunk in backend_stream:
-                if chunk.type == "content":
+                chunk_type = self._get_chunk_type_value(chunk)
+                if chunk_type == "content":
                     assistant_response += chunk.content
                     yield chunk
-                elif chunk.type == "tool_calls":
+                elif chunk_type == "tool_calls":
                     chunk_tool_calls = getattr(chunk, "tool_calls", []) or []
                     tool_calls.extend(chunk_tool_calls)
                     yield chunk
-                elif chunk.type == "complete_message":
+                elif chunk_type == "complete_message":
                     # Backend provided the complete message structure
                     complete_message = chunk.complete_message
                     # Don't yield this - it's for internal use
-                elif chunk.type == "complete_response":
+                elif chunk_type == "complete_response":
                     # Backend provided the raw Responses API response
                     if chunk.response:
                         complete_message = chunk.response
 
                         # Extract and yield tool calls for orchestrator processing
-                        if (
-                            isinstance(chunk.response, dict)
-                            and "output" in chunk.response
-                        ):
+                        if isinstance(chunk.response, dict) and "output" in chunk.response:
                             response_tool_calls = []
                             for output_item in chunk.response["output"]:
                                 if output_item.get("type") == "function_call":
                                     response_tool_calls.append(output_item)
-                                    tool_calls.append(
-                                        output_item
-                                    )  # Also store for fallback
+                                    tool_calls.append(output_item)  # Also store for fallback
 
                             # Yield tool calls so orchestrator can process them
                             if response_tool_calls:
-                                yield StreamChunk(
-                                    type="tool_calls", tool_calls=response_tool_calls
-                                )
+                                yield StreamChunk(type="tool_calls", tool_calls=response_tool_calls)
                     # Complete response is for internal use - don't yield it
-                elif chunk.type == "done":
+                elif chunk_type == "done":
                     # Add complete response to history
                     if complete_message:
                         # For Responses API: complete_message is the response object with 'output' array
                         # Each item in output should be added to conversation history individually
-                        if (
-                            isinstance(complete_message, dict)
-                            and "output" in complete_message
-                        ):
+                        if isinstance(complete_message, dict) and "output" in complete_message:
                             self.conversation_history.extend(complete_message["output"])
                         else:
                             # Fallback if it's already in message format
@@ -220,23 +241,41 @@ class SingleAgent(ChatAgent):
         tools: List[Dict[str, Any]] = None,
         reset_chat: bool = False,
         clear_history: bool = False,
+        current_stage: CoordinationStage = None,
     ) -> AsyncGenerator[StreamChunk, None]:
+        # print("Agent: ", self.agent_id)
+        # for message in messages:
+        #     print(f"Message: {message}\n")
+        # print("Messages End. \n")
         """Process messages through single backend with tool support."""
         if clear_history:
             # Clear history but keep system message if it exists
-            system_messages = [
-                msg for msg in self.conversation_history if msg.get("role") == "system"
-            ]
+            system_messages = [msg for msg in self.conversation_history if msg.get("role") == "system"]
             self.conversation_history = system_messages.copy()
+            # Clear backend history while maintaining session
+            if self.backend.is_stateful():
+                await self.backend.clear_history()
 
         if reset_chat:
             # Reset conversation history to the provided messages
             self.conversation_history = messages.copy()
+            # Reset backend state completely
+            if self.backend.is_stateful():
+                await self.backend.reset_state()
             backend_messages = self.conversation_history.copy()
         else:
             # Regular conversation - append new messages to agent's history
             self.conversation_history.extend(messages)
-            backend_messages = self.conversation_history.copy()
+            # Handle stateful vs stateless backends differently
+            if self.backend.is_stateful():
+                # Stateful: only send new messages, backend maintains context
+                backend_messages = messages.copy()
+            else:
+                # Stateless: send full conversation history
+                backend_messages = self.conversation_history.copy()
+
+        if current_stage:
+            self.backend.set_stage(current_stage)
 
         # Create backend stream and process it
         backend_stream = self.backend.stream_with_tools(
@@ -244,11 +283,16 @@ class SingleAgent(ChatAgent):
             tools=tools,  # Use provided tools (for MassGen workflow)
             agent_id=self.agent_id,
             session_id=self.session_id,
+            **self._get_backend_params(),
         )
 
         async for chunk in self._process_stream(backend_stream, tools):
             yield chunk
 
+    def _get_backend_params(self) -> Dict[str, Any]:
+        """Get additional backend parameters. Override in subclasses."""
+        return {}
+
     def get_status(self) -> Dict[str, Any]:
         """Get current agent status."""
         return {
@@ -259,15 +303,21 @@ class SingleAgent(ChatAgent):
             "conversation_length": len(self.conversation_history),
         }
 
-    def reset(self) -> None:
+    async def reset(self) -> None:
         """Reset conversation for new chat."""
         self.conversation_history.clear()
 
+        # Reset stateful backend if needed
+        if self.backend.is_stateful():
+            await self.backend.reset_state()
+
         # Re-add system message if it exists
         if self.system_message:
-            self.conversation_history.append(
-                {"role": "system", "content": self.system_message}
-            )
+            self.conversation_history.append({"role": "system", "content": self.system_message})
+
+    def get_configurable_system_message(self) -> Optional[str]:
+        """Get the user-configurable part of the system message."""
+        return self.system_message
 
     def set_model(self, model: str) -> None:
         """Set the model for this agent."""
@@ -278,16 +328,11 @@ class SingleAgent(ChatAgent):
         self.system_message = system_message
 
         # Remove old system message if exists
-        if (
-            self.conversation_history
-            and self.conversation_history[0].get("role") == "system"
-        ):
+        if self.conversation_history and self.conversation_history[0].get("role") == "system":
             self.conversation_history.pop(0)
 
         # Add new system message at the beginning
-        self.conversation_history.insert(
-            0, {"role": "system", "content": system_message}
-        )
+        self.conversation_history.insert(0, {"role": "system", "content": system_message})
 
 
 class ConfigurableAgent(SingleAgent):
@@ -297,6 +342,13 @@ class ConfigurableAgent(SingleAgent):
     This bridges the gap between SingleAgent and the MassGen system by supporting
     all the advanced configuration options (web search, code execution, etc.)
     while maintaining the simple chat interface.
+
+    TODO: Consider merging with SingleAgent. The main difference is:
+    - SingleAgent: backend parameters passed directly to constructor/methods
+    - ConfigurableAgent: backend parameters come from AgentConfig object
+
+    Could be unified by making SingleAgent accept an optional config parameter
+    and using _get_backend_params() pattern for all parameter sources.
     """
 
     def __init__(
@@ -323,42 +375,9 @@ class ConfigurableAgent(SingleAgent):
 
         # ConfigurableAgent relies on backend_params for model configuration
 
-    async def chat(
-        self,
-        messages: List[Dict[str, Any]],
-        tools: List[Dict[str, Any]] = None,
-        reset_chat: bool = False,
-        clear_history: bool = False,
-    ) -> AsyncGenerator[StreamChunk, None]:
-        """Process messages with full AgentConfig capabilities."""
-        if clear_history:
-            # Clear history but keep system message if it exists
-            system_messages = [
-                msg for msg in self.conversation_history if msg.get("role") == "system"
-            ]
-            self.conversation_history = system_messages.copy()
-
-        if reset_chat:
-            # Reset conversation history to the provided messages
-            self.conversation_history = messages.copy()
-            backend_messages = self.conversation_history.copy()
-        else:
-            # Regular conversation - append new messages to agent's history
-            self.conversation_history.extend(messages)
-            backend_messages = self.conversation_history.copy()
-
-        # Create backend stream with config parameters and process it
-        backend_params = self.config.get_backend_params()
-        backend_stream = self.backend.stream_with_tools(
-            messages=backend_messages,
-            tools=tools,  # Use provided tools (for MassGen workflow)
-            agent_id=self.agent_id,
-            session_id=self.session_id,
-            **backend_params,
-        )
-
-        async for chunk in self._process_stream(backend_stream, tools):
-            yield chunk
+    def _get_backend_params(self) -> Dict[str, Any]:
+        """Get backend parameters from config."""
+        return self.config.get_backend_params()
 
     def get_status(self) -> Dict[str, Any]:
         """Get current agent status with config details."""
@@ -368,26 +387,43 @@ class ConfigurableAgent(SingleAgent):
                 "agent_type": "configurable",
                 "config": self.config.to_dict(),
                 "capabilities": {
-                    "web_search": self.config.backend_params.get(
-                        "enable_web_search", False
-                    ),
-                    "code_execution": self.config.backend_params.get(
-                        "enable_code_interpreter", False
-                    ),
+                    "web_search": self.config.backend_params.get("enable_web_search", False),
+                    "code_execution": self.config.backend_params.get("enable_code_interpreter", False),
                 },
-            }
+            },
         )
         return status
 
+    def get_configurable_system_message(self) -> Optional[str]:
+        """Get the user-configurable part of the system message for ConfigurableAgent."""
+        # Try multiple sources in order of preference
+
+        # First check if backend has system prompt configuration
+        if self.config and self.config.backend_params:
+            backend_params = self.config.backend_params
+
+            # For Claude Code: prefer system_prompt (complete override)
+            if "system_prompt" in backend_params:
+                return backend_params["system_prompt"]
+
+            # Then append_system_prompt (additive)
+            if "append_system_prompt" in backend_params:
+                return backend_params["append_system_prompt"]
+
+        # Fall back to custom_system_instruction (deprecated but still supported)
+        if self.config and self.config.custom_system_instruction:
+            return self.config.custom_system_instruction
+
+        # Finally fall back to parent class implementation
+        return super().get_configurable_system_message()
+
 
 # =============================================================================
 # CONVENIENCE FUNCTIONS
 # =============================================================================
 
 
-def create_simple_agent(
-    backend: LLMBackend, system_message: str = None, agent_id: str = None
-) -> SingleAgent:
+def create_simple_agent(backend: LLMBackend, system_message: str = None, agent_id: str = None) -> SingleAgent:
     """Create a simple single agent."""
     # Use MassGen evaluation system message if no custom system message provided
     if system_message is None:
@@ -395,24 +431,10 @@ def create_simple_agent(
 
         templates = MessageTemplates()
         system_message = templates.evaluation_system_message()
-
-        import pdb
-
-        pdb.set_trace()
-        print(system_message)
-
-    import pdb
-
-    pdb.set_trace()
-    print(system_message)
-    return SingleAgent(
-        backend=backend, agent_id=agent_id, system_message=system_message
-    )
+    return SingleAgent(backend=backend, agent_id=agent_id, system_message=system_message)
 
 
-def create_expert_agent(
-    domain: str, backend: LLMBackend, model: str = "gpt-4o-mini"
-) -> ConfigurableAgent:
+def create_expert_agent(domain: str, backend: LLMBackend, model: str = "gpt-4o-mini") -> ConfigurableAgent:
     """Create an expert agent for a specific domain."""
     from .agent_config import AgentConfig
 
@@ -420,9 +442,7 @@ def create_expert_agent(
     return ConfigurableAgent(config=config, backend=backend)
 
 
-def create_research_agent(
-    backend: LLMBackend, model: str = "gpt-4o-mini"
-) -> ConfigurableAgent:
+def create_research_agent(backend: LLMBackend, model: str = "gpt-4o-mini") -> ConfigurableAgent:
     """Create a research agent with web search capabilities."""
     from .agent_config import AgentConfig
 
@@ -430,9 +450,7 @@ def create_research_agent(
     return ConfigurableAgent(config=config, backend=backend)
 
 
-def create_computational_agent(
-    backend: LLMBackend, model: str = "gpt-4o-mini"
-) -> ConfigurableAgent:
+def create_computational_agent(backend: LLMBackend, model: str = "gpt-4o-mini") -> ConfigurableAgent:
     """Create a computational agent with code execution."""
     from .agent_config import AgentConfig