agent-mcp 0.1.3__py3-none-any.whl → 0.1.4__py3-none-any.whl

agent_mcp/heterogeneous_group_chat.py

@@ -11,6 +11,122 @@ from typing import List, Dict, Any, Optional, Union, Sequence
 from .mcp_transport import HTTPTransport
 from .enhanced_mcp_agent import EnhancedMCPAgent
 from .mcp_agent import MCPAgent
+import re
+import string
+import logging
+from typing import Dict, Any, Optional, List
+import os
+import time
+
+logger = logging.getLogger(__name__)
+
+class CoordinatorAgent(EnhancedMCPAgent):
+    def __init__(self, group_chat, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.group_chat = group_chat  # Reference to HeterogeneousGroupChat instance
+
+    async def handle_incoming_message(self, message: Dict):
+
+        # First, call the super method to handle default processing
+        await super().handle_incoming_message(message)
+
+        # Delegate to the group chat's custom handler
+        await self.group_chat._handle_coordinator_message(message, message.get('message_id'))
+
+class ContextAgent(EnhancedMCPAgent):
+    """Agent that maintains and provides access to task context and results."""
+    
+    def __init__(self, name: str, **kwargs):
+        super().__init__(
+            name=name,
+            transport=None,  # No transport needed for internal use
+            server_mode=False,
+            client_mode=False,
+            **kwargs
+        )
+        self.task_results = {}  # Store task results
+        self.task_descriptions = {}  # Store task descriptions
+        self._pending_tasks = {}  # Track pending tasks with their completion events
+        print(f"[DEBUG] __init__: self ID: {id(self)}, _pending_tasks ID: {id(self._pending_tasks)}")
+
+        logger.info(f"[{self.name}] Initialized as context agent")
+
+    async def query_context(self, query: str) -> Dict[str, Any]:
+        """
+        Query the context agent for information.
+        
+        Args:
+            query: Natural language query about tasks, results, or context.
+                  Can also be a request to generate email content.
+            
+        Returns:
+            Dict with 'answer' and 'supporting_data' keys
+        """ 
+        # Regular context query
+        return await self.generate_response(query)
+        
+    async def generate_response(self, query: str) -> Dict[str, Any]:
+        """Answer general questions about tasks and results."""
+        context = {
+            "task_descriptions": self.task_descriptions,
+            "task_results": {
+                tid: str(r)[:500] for tid, r in self.task_results.items()
+            }
+        }
+        
+        messages = [{
+            "role": "system",
+            "content": f"""You are a context assistant that generates responses, results or content based on task as query. 
+            You will be given a query, task description and expected to generate a content, response, as result or output
+            
+            Available context, tasks and their results: {context}
+            
+            """
+        }, {
+            "role": "user",
+            "content": f"""Generate a response, result or content based on these instructions:
+            {query}
+            
+            """
+        }]
+        
+        try:
+            response = await self.a_generate_reply(messages)
+
+             # Ensure we have a valid response
+            if not response:
+                raise ValueError("Empty response from LLM")
+                
+            # Handle both string and dictionary responses
+            if isinstance(response, str):
+                content = response
+            elif isinstance(response, dict):
+                content = response.get("content", "")
+                if not content and "message" in response:
+                    content = response.get("message", {}).get("content", "")
+            else:
+                content = str(response)
+                
+            print(f"Generated response: {content}")  # Log first 200 chars
+            return {
+                "answer": content,
+                "supporting_data": context
+            }
+        except Exception as e:
+            logger.error(f"Error querying context: {e}")
+            return {
+                "answer": f"Error: {str(e)}",
+                "supporting_data": {}
+            }
+            
+
+    async def update_task(self, task_id: str, task_data: Dict, result: Optional[Any] = None):
+        """Update task information and results."""
+        self.task_descriptions[task_id] = task_data.get("description", "No description")
+        if result is not None:
+            self.task_results[task_id] = result
+        logger.debug(f"Updated context for task {task_id}")
+ 
 
 class HeterogeneousGroupChat:
     """
@@ -44,12 +160,38 @@ class HeterogeneousGroupChat:
         # Initialize directly on the group chat instance first
         self.task_results: Dict[str, Any] = {} 
         self.task_dependencies: Dict[str, Dict] = {}
+        self.dependency_results: Dict[str, Any] = {}  # Initialize dependency results
+        
+        # Store coordinator config for later use
+        self.coordinator_config = coordinator_config or {}
+        
+        # Initialize context agent with default LLM config
+        default_llm_config = {
+            "config_list": [{
+                "model": "gpt-3.5-turbo",
+                "api_key": os.getenv("OPENAI_API_KEY"),  # Will be set when coordinator is created
+                "api_type": "openai"
+            }]
+        }
+        
+        # Use provided config or default
+        llm_config = self.coordinator_config.get("llm_config", default_llm_config).copy()
+        
+        # Ensure config_list exists and has at least one config
+        if not llm_config.get("config_list"):
+            llm_config["config_list"] = default_llm_config["config_list"]
+        
+        # Create context agent
+        self.context_agent = ContextAgent(
+            name=f"{self.name}_context",
+            llm_config=llm_config
+        )
 
     def _get_agent_url(self, agent_name: str) -> str:
         """Get the URL for an agent on the deployed server"""
         return f"{self.server_url}/agents/{agent_name}"
         
-    def create_coordinator(self, api_key: str) -> EnhancedMCPAgent:
+    def create_coordinator(self, api_key: str) -> CoordinatorAgent:
         """Create the coordinator agent for the group chat"""
         # Avoid creating coordinator if it already exists
         if self.coordinator:
@@ -79,6 +221,7 @@ class HeterogeneousGroupChat:
             },
         }
         
+        
         # --- Merge Default and User Config --- 
         # User config takes precedence
         final_config = default_config.copy() # Start with defaults
@@ -102,12 +245,24 @@ class HeterogeneousGroupChat:
              final_config["llm_config"]["config_list"][0]["api_key"] = api_key
 
 
+        # Update context agent's LLM config to match coordinator's
+        if hasattr(self, 'context_agent') and self.context_agent:
+            # Get the final config that will be used by the coordinator
+            context_llm_config = final_config.get('llm_config', {})
+            # Update the context agent's config
+            if hasattr(self.context_agent, 'llm_config'):
+                self.context_agent.llm_config = context_llm_config
+                logger.info(f"Updated context agent's LLM config to match coordinator")
+        
         # --- Create Coordinator Agent --- 
         print(f"Creating coordinator with config: {final_config}") # Debug: Log final config
-        self.coordinator = EnhancedMCPAgent(**final_config)
+        self.coordinator = CoordinatorAgent(self, **final_config)
         
         # --- Set Message Handler ---
-        self.coordinator.transport.set_message_handler(self._handle_coordinator_message)
+        #self.coordinator.transport.set_message_handler(self._handle_coordinator_message)
+        # Use a lambda to explicitly capture the correct 'self' (the HeterogeneousGroupChat instance)
+        #self.coordinator.transport.set_message_handler(lambda msg, msg_id: self._handle_coordinator_message(msg, msg_id))
+
         return self.coordinator
         
     def add_agents(self, agents: Union[MCPAgent, Sequence[MCPAgent]]) -> List[MCPAgent]:
@@ -218,27 +373,186 @@ class HeterogeneousGroupChat:
             # Don't run the agent if registration fails - it won't be able to communicate
             return None # Indicate failure
         
-    async def submit_task(self, task: Dict[str, Any]) -> None:
-        """Submit a task to the group chat."""
-        print(f"***** [{self.name}] ENTERING submit_task *****", flush=True) # Ensure entry is logged
-        if not self.coordinator:
-            raise ValueError("Group chat not connected. Call connect() first.")
+ 
+    async def query_context(self, query: str) -> Dict[str, Any]:
+        """
+        Query the context agent for information about tasks and results.
+        
+        Args:
+            query: Natural language query about tasks, results, or context
             
-        self.task_results = {} # Reset results for new task submission
+        Returns:
+            Dict with 'answer' and 'supporting_data' keys
+        """
+        return await self.context_agent.query_context(query)
+
+    def _inject_dependency_results(self, step: Dict, dependency_results: Dict) -> Dict:
+        """Injects dependency results into a step's content.
+
+        If the step contains string.Template style placeholders (e.g., ${task_id}),
+        it substitutes them with the corresponding results.
+
+        If no placeholders are found, it assumes the agent needs the raw results
+        and adds them to the step's content under the key 'dependency_data'.
+        """
+        if not step:
+            return step
+
+        # Check if any part of the step contains a placeholder
+        logger.info(f"No placeholders detected in step {step.get('task_id', 'N/A')}. Adding raw dependency data.")
+        dependency_data = {}
+        for dep_task_id in step.get("depends_on", []):
+            result_value = dependency_results.get(dep_task_id)
+            if result_value is None:
+                logger.warning(f"No result found for dependency '{dep_task_id}' when preparing raw data for step '{step.get('task_id', 'N/A')}'")
+                extracted_value = None # Or some placeholder? 
+            elif isinstance(result_value, dict):
+                 # Prioritize 'output', then 'result', then string representation
+                if 'output' in result_value: # Check presence first
+                    extracted_value = result_value['output']
+                elif 'result' in result_value: # Check presence first
+                    extracted_value = result_value['result']
+                else:
+                    logger.warning(f"Raw dependency '{dep_task_id}': Neither 'output' nor 'result' key found in dict result. Using full dict.")
+                    extracted_value = result_value # Pass the whole dict
+            else:
+                extracted_value = result_value # Pass strings, numbers, lists as-is
+
+            dependency_data[dep_task_id] = extracted_value
+            
+        # Ensure 'content' exists and add the data
+        if "content" not in step:
+            step["content"] = {}
+        if not isinstance(step["content"], dict):
+            logger.warning(f"Step {step.get('task_id', 'N/A')} content is not a dict, cannot add dependency_data. Content: {step['content']}")
+        else:
+            step["content"]["dependency_data"] = dependency_data
+            
+        return step
+
+    async def submit_task(self, task: Dict[str, Any], inject_at_submit_time: bool = False) -> None:
+        """
+        Submit a group task. If inject_at_submit_time is True, inject dependency results into each step now.
+        If False, inject at the last possible moment (just before sending to agents).
+        """
+        # Reset state for new task submission
+        self.task_results = {}
+        self.context_agent.task_results = {}
+        self.context_agent.task_descriptions = {}
+        self._pending_tasks = {}  # Track pending tasks with their completion events to ensure its always new for each submission
+        print(f"[DEBUG] __init__: self ID: {id(self)}, _pending_tasks ID: {id(self._pending_tasks)}")
+        steps = task.get("steps", [])
+        self.task_dependencies = {step["task_id"]: step for step in steps} # Store task dependencies
+
+        self._inject_at_submit_time = inject_at_submit_time
+        if inject_at_submit_time:
+            steps = [self._inject_dependency_results(step, self.task_results) for step in steps]
+        self._pending_steps = steps  # Store for later use
+        await self._submit_steps(steps)
+
+    async def _submit_steps(self, steps):
+        for step in steps:
+            try:
+                # Only inject here if not already injected at submit time
+                if not getattr(self, '_inject_at_submit_time', True):
+                    step = self._inject_dependency_results(step, self.task_results)
+                
+                # Create and store event for this task
+                task_id = step['task_id']
+                future = asyncio.Future()
+                self._pending_tasks[task_id] = future
+                print(f"[DEBUG] _submit_steps: Added task {task_id}. Current _pending_tasks: {list(self._pending_tasks.keys())} (ID: {id(self._pending_tasks)})")
+                print(f"[DEBUG] Added task {task_id} to pending_tasks")
+                # Submit the task
+                await self._send_step_to_agent(step)
+                
+                # Wait for task completion with timeout
+                try:
+                    await asyncio.wait_for(future, timeout=60)
+                    print(f"[DEBUG] Task {task_id} completed with result: {future.result()}")
+                except asyncio.TimeoutError:
+                    logger.error(f"Task {task_id} timed out")
+                    raise TimeoutError(f"Task {task_id} timed out")
+                finally:
+                    # Clean up
+                    self._pending_tasks.pop(task_id, None)
+                    
+            except Exception as e:
+                logger.error(f"Error in step {step.get('task_id')}: {e}")
+                raise
+
+    async def _generate_content(self, description: str, content: Dict) -> Dict:
+        """Generate content using the ContextAgent based on the task description.
+        
+        Args:
+            description: The task description
+            content: Existing content to be augmented with generated content
+            
+        Returns:
+            Dict: Content with generated fields merged in
+        """
+        # Ensure we have a dictionary to work with
+        if not isinstance(content, dict):
+            content = {}
+            
+        try:
+            # Generate content using the context agent
+            generated = await self.context_agent.generate_response(description)
+            
+            print(f"Generated content: {generated}")
+            if generated and isinstance(generated, dict) and "answer" in generated:
+                try:
+                    # Try to parse as JSON first
+                    generated_content = json.loads(generated["answer"])
+                    # If it's a dictionary, merge it intelligently
+                    if isinstance(generated_content, dict):
+                        # Merge with existing content, with generated content taking precedence
+                        content = {**content, **generated_content}
+                    # If it's not a dictionary, store it as generated_content
+                    else:
+                        content["content"] = generated_content # this is the generated content that will be sent to the agent
+                except (json.JSONDecodeError, TypeError):
+                    # If not JSON, store the raw answer
+                    content["content"] = generated["answer"] # this is the generated content that will be sent to the agent
+        
+        except Exception as e:
+            logger.error(f"Error generating content: {e}")
+            print(f"Error generating content: {e}")
+            content["error"] = str(e)
+        
+        return content
+
+    async def _send_step_to_agent(self, step):
+        # 1-line dependency enforcement: skip if any dependency missing or empty
+        #if any(not self.dependency_results.get(dep) for dep in step.get("depends_on", [])):
+        #    print(f"Skipping {step['task_id']} because dependencies not satisfied: {step.get('depends_on', [])}")
+        #    return
+        step = self._inject_dependency_results(step, self.task_results)
+        
+        # Update context agent with task info
+        task_id = step.get("task_id", str(id(step)))
+        await self.context_agent.update_task(task_id, step)
+        
         print("\n=== Submitting task to group ===")
 
         # Ensure task is in the correct format
-        if not isinstance(task, dict) or 'type' not in task:
-            task = {'type': 'task', 'content': task}
+        if not isinstance(step, dict):
+            raise ValueError("Invalid task format. Expected a dictionary.")
+
+        # Handle content generation if requested
+        if step.get("content", {}).get("generate_content", False):
+            description = step.get("description", "")
+            if description:
+                # Make a copy to avoid modifying the original
+                content = step.get("content", {}).copy()
+                # Remove the flag before generation
+                content.pop("generate_content", None)
+                # Generate and merge content
+                step["content"] = await self._generate_content(description, content)
 
         # Store task dependencies from the input task definition
         # We need a dictionary where keys are the step task_ids
-        if isinstance(task['content'], dict) and all(isinstance(v, dict) for v in task['content'].values()):
-            # If task is already a dict mapping task_ids to task info
-            self.task_dependencies = task['content']
-        else:
-            # If task has a steps list, convert it to a dict
-            self.task_dependencies = {step["task_id"]: step for step in task['content'].get("steps", [])}
+        self.task_dependencies[step["task_id"]] = step
         print(f"Parsed Step Dependencies: {self.task_dependencies}")
 
         # Also store in coordinator instance if it exists
@@ -259,22 +573,21 @@ class HeterogeneousGroupChat:
 
         # Assign tasks to agents based on the structure
         # Submit tasks to their respective agents
-        for task_id, task_info in self.task_dependencies.items():
-            print(f"[DEBUG - {self.name}] Loop Iteration: Processing task_id '{task_id}' for agent '{task_info['agent']}'", flush=True)
-            agent_name = task_info["agent"]
-            # Create message with all necessary fields including content
-            message = {
-                "type": "task",
-                "task_id": task_id,
-                "description": task_info["description"],
-                "content": task_info.get("content", {}),  # Include task content
-                "depends_on": task_info.get("depends_on", []),  # Include dependencies
-                "reply_to": f"{self.server_url}/message/{self.coordinator.name}" # Full URL for reply
-            }
-            print(f"Sending task to {agent_name}")
-            print(f"Task message: {message}")
-            # Use coordinator's transport to send task to agent
-            await coordinator_transport.send_message(agent_name, message)
+        agent_name = step["agent"]
+        # Create message with all necessary fields including content
+        message = {
+            "type": "task",
+            "task_id": step["task_id"],
+            "description": step["description"],
+            "sender": self.coordinator.name,
+            "content": step.get("content", {}),  # Include task content
+            "depends_on": step.get("depends_on", []),  # Include dependencies
+            "reply_to": f"{self.server_url}/message/{self.coordinator.name}" # Full URL for reply
+        }
+        print(f"Sending task to {agent_name}")
+        print(f"Task message: {message}")
+        # Use coordinator's transport to send task to agent
+        await coordinator_transport.send_message(agent_name, message)
             
         print("Task submitted. Waiting for completion...")
         
@@ -317,11 +630,12 @@ class HeterogeneousGroupChat:
             
     async def _handle_coordinator_message(self, message: Dict, message_id: str):
         """Handles messages received by the coordinator's transport."""
+        
         if not self.coordinator: # Ensure coordinator exists
-            print("[Coordinator Handler] Error: Coordinator not initialized.")
+            print(f"[Coordinator Handler] Error: Coordinator not initialized. Self ID: {id(self)}")
             return
             
-        print(f"\n[Coordinator {self.coordinator.name}] Received message: {message}")
+        print(f"\n[Coordinator {self.coordinator.name}] Received message: {message}. Self ID: {id(self)}")        
         
         # Handle messages wrapped in 'body' field
         if isinstance(message, dict) and 'body' in message:
@@ -339,19 +653,66 @@ class HeterogeneousGroupChat:
         msg_type = message.get("type")
         task_id = message.get("task_id")
         
-        print(f"[Coordinator {self.coordinator.name}] Processing message type '{msg_type}' for task {task_id}")
+        print(f"[Coordinator {self.coordinator.name}] Processing message type '{msg_type}' for task {task_id}. Current _pending_tasks in handler: {list(self._pending_tasks.keys())} (ID: {id(self._pending_tasks)})")        
         
         if msg_type in ["result", "task_result"]:  # Handle both result types
-            result_content = message.get("result") or message.get("description")  # Try both fields
+            # First try direct fields, then try parsing content.text if it exists
+            result_content = None
+            
+            # Try direct fields first
+            result_content = message.get("result") or message.get("description")
+            
+            # If not found, try to parse from content.text
+            if result_content is None and "content" in message and isinstance(message["content"], dict):
+                content_text = message["content"].get("text")
+                if content_text:
+                    try:
+                        content_data = json.loads(content_text)
+                        result_content = content_data.get("result") or content_data.get("description")
+                        # Update task_id from content if not set
+                        if not task_id and "task_id" in content_data:
+                            task_id = content_data["task_id"]
+                    except (json.JSONDecodeError, AttributeError, TypeError) as e:
+                        print(f"[Coordinator {self.coordinator.name}] Error parsing content.text: {e}")
+            
             if task_id and result_content is not None:
                 print(f"[Coordinator {self.coordinator.name}] Storing result for task {task_id}")
                 # Store result in both the group chat and coordinator
                 self.task_results[task_id] = result_content
+                self.dependency_results[task_id] = result_content  # Required for template resolution
+                if "dependency_results" not in self.coordinator.task_results:
+                    self.coordinator.task_results["dependency_results"] = {}
+                self.coordinator.task_results["dependency_results"][task_id] = result_content
                 self.coordinator.task_results[task_id] = result_content
-                print(f"[Coordinator {self.coordinator.name}] Stored result: {result_content[:100]}...")
+                print(f"[Coordinator {self.coordinator.name}] Stored result for task {task_id}")
+                print(f"[Coordinator {self.coordinator.name}] Stored result: {result_content}...")
                 print(f"[Coordinator {self.coordinator.name}] Current task results: {list(self.task_results.keys())}")
                 print(f"[Coordinator {self.coordinator.name}] Current dependencies: {self.task_dependencies}")
                 
+                # Signal task completion if anyone is waiting
+                if not hasattr(self, '_pending_tasks'):
+                    self._pending_tasks = {}
+                print(f"[Coordinator {self.coordinator.name}] Current pending tasks: {list(self._pending_tasks.keys())}")
+                print(f"[DEBUG] Checking if task {task_id} is in pending_tasks: {task_id in self._pending_tasks}")
+                
+                if task_id in self._pending_tasks:
+                    print(f"[Coordinator {self.coordinator.name}] Signaling completion for task {task_id}")
+                    future = self._pending_tasks[task_id]
+                    if not future.done():
+                        #future.set_result(result_content)
+                        asyncio.get_event_loop().call_soon_threadsafe(lambda: future.set_result(result_content) if not future.done() else None)
+                        print(f"[DEBUG] Set result for task {task_id}")
+                        await asyncio.sleep(0)
+                        print(f"[Coordinator {self.coordinator.name}] Completed task {task_id}")
+                        #asyncio.get_running_loop().call_soon(future.set_result, result_content)
+                        #asyncio.get_event_loop().call_soon_threadsafe(lambda: future.set_result(result_content) if not future.done() else None)
+                    # Clean up the task after signaling
+                    #if task_id in self._pending_tasks:
+                        #del self._pending_tasks[task_id]
+                    print(f"[Coordinator {self.coordinator.name}] Completed task {task_id}")
+                else:
+                    print(f"[Coordinator {self.coordinator.name}] Task {task_id} not found in pending tasks")
+                
                 # Acknowledge the message
                 try:
                     if message_id:  # Only acknowledge if we have a message ID
@@ -359,6 +720,7 @@ class HeterogeneousGroupChat:
                         print(f"[Coordinator {self.coordinator.name}] Acknowledged message {message_id}")
                 except Exception as e:
                     print(f"[Coordinator {self.coordinator.name}] Error acknowledging message {message_id}: {e}")
+                return
             else:
                 print(f"[Coordinator {self.coordinator.name}] Received invalid result message (missing task_id or result): {message}")
         elif msg_type == "get_result":  # Handle get result request
@@ -393,6 +755,14 @@ class HeterogeneousGroupChat:
             except Exception as e:
                 print(f"[Coordinator {self.coordinator.name}] Error acknowledging message {message_id}: {e}")
 
+    @property
+    def group_state(self) -> dict:
+        """
+        Returns a merged dictionary of all task results (group and coordinator).
+        Agents can use this to access the shared group chat history/results.
+        """
+        return {**self.task_results, **(self.coordinator.task_results if self.coordinator else {})}
+
     async def shutdown(self):
         """Gracefully disconnect all agents and cancel their tasks."""
         print(f"Initiating shutdown for {len(self._agent_tasks)} agent tasks...")
@@ -422,3 +792,7 @@ class HeterogeneousGroupChat:
             print("All agent transports disconnected.")
             
         print("Shutdown complete.")
+
+    # === Minimal free-flow chat: send a message to any agent ===
+    async def send_chat_message(self, agent_name, message):
+        await self.coordinator.transport.send_message(agent_name, {"type": "message", "content": message})