agentensemble 0.1.0__py3-none-any.whl

agentensemble/agents/structured_agent.py

@@ -0,0 +1,210 @@
+"""
+Structured Output Agent Implementation
+
+Agent with structured output support using LangChain's structured output feature.
+Supports both langchain>=1.1 (create_agent with response_format) and 
+langchain<1.1 (with_structured_output method).
+
+Based on: https://docs.langchain.com/oss/python/langchain/structured-output
+"""
+
+from typing import Any, Dict, Optional, Type, Union
+from pydantic import BaseModel
+
+try:
+    from agentensemble.utils.llm import get_mistral_model
+    LLM_AVAILABLE = True
+except ImportError:
+    LLM_AVAILABLE = False
+    get_mistral_model = None
+
+# Try langchain>=1.1 approach (create_agent with response_format)
+try:
+    from langchain.agents import create_agent
+    from langchain.agents.structured_output import ToolStrategy, ProviderStrategy
+    CREATE_AGENT_AVAILABLE = True
+    STRUCTURED_OUTPUT_AVAILABLE = True
+except ImportError:
+    CREATE_AGENT_AVAILABLE = False
+    create_agent = None
+    ToolStrategy = None
+    ProviderStrategy = None
+    STRUCTURED_OUTPUT_AVAILABLE = False
+
+# with_structured_output is available if model has the method
+WITH_STRUCTURED_AVAILABLE = True
+
+
+class StructuredAgent:
+    """
+    Agent with structured output support.
+    
+    Uses LangChain's structured output feature to return data in a specific format.
+    
+    For langchain>=1.1: Uses create_agent with response_format parameter
+    For langchain<1.1: Uses model.with_structured_output() method
+    
+    Based on: https://docs.langchain.com/oss/python/langchain/structured-output
+    """
+    
+    def __init__(
+        self,
+        name: str = "structured_agent",
+        tools: Optional[list] = None,
+        model: Optional[Any] = None,
+        response_format: Optional[Union[Type[BaseModel], Any]] = None,
+        system_prompt: Optional[str] = None,
+        **kwargs
+    ):
+        """
+        Initialize structured agent.
+        
+        Args:
+            name: Agent name
+            tools: Available tools
+            model: LLM model (defaults to Mistral if not provided)
+            response_format: Structured output schema (Pydantic model)
+            system_prompt: System prompt for the agent
+            **kwargs: Additional configuration
+        """
+        if not LLM_AVAILABLE:
+            raise ImportError(
+                "langchain-mistralai is required for StructuredAgent. "
+                "Install with: pip install langchain-mistralai"
+            )
+        
+        self.name = name
+        self.tools = tools or []
+        self.model = model or get_mistral_model()
+        self.response_format = response_format
+        self.system_prompt = system_prompt
+        self.config = kwargs
+        
+        # Determine which approach to use
+        if CREATE_AGENT_AVAILABLE and response_format:
+            # Use create_agent with response_format (langchain>=1.1)
+            self._use_create_agent()
+        elif response_format and hasattr(self.model, 'with_structured_output'):
+            # Use with_structured_output (langchain<1.1)
+            self._use_with_structured_output()
+        else:
+            # Basic agent without structured output
+            if CREATE_AGENT_AVAILABLE:
+                self._use_create_agent()
+            else:
+                raise ImportError(
+                    "langchain>=1.1 is required for StructuredAgent with tools. "
+                    "Install with: pip install 'langchain>=1.1'"
+                )
+    
+    def _use_create_agent(self):
+        """Use create_agent approach (langchain>=1.1)"""
+        if not CREATE_AGENT_AVAILABLE:
+            raise ImportError("create_agent not available. Requires langchain>=1.1")
+        
+        agent_kwargs = {
+            "model": self.model,
+            "tools": self.tools,
+        }
+        
+        if self.system_prompt:
+            agent_kwargs["system_prompt"] = self.system_prompt
+        
+        if self.response_format:
+            agent_kwargs["response_format"] = self.response_format
+        
+        agent_kwargs.update(self.config)
+        self.agent = create_agent(**agent_kwargs)
+        self._method = "create_agent"
+    
+    def _use_with_structured_output(self):
+        """Use with_structured_output approach (langchain<1.1)"""
+        if not self.response_format:
+            raise ValueError("response_format required for with_structured_output")
+        
+        # Bind structured output to model
+        self.structured_model = self.model.with_structured_output(self.response_format)
+        self._method = "with_structured_output"
+    
+    def run(
+        self,
+        query: str,
+        **kwargs
+    ) -> Dict[str, Any]:
+        """
+        Execute agent with structured output.
+        
+        Args:
+            query: Input query
+            **kwargs: Additional parameters
+            
+        Returns:
+            Dictionary containing:
+            - 'result': Natural language response
+            - 'structured_response': Structured output (if response_format provided)
+            - 'metadata': Execution metadata
+        """
+        if self._method == "create_agent":
+            # Use create_agent approach
+            result = self.agent.invoke({
+                "messages": [{"role": "user", "content": query}],
+                **kwargs
+            })
+            
+            # Extract structured response if available
+            structured_response = result.get("structured_response")
+            
+            # Get last message content as natural language result
+            messages = result.get("messages", [])
+            last_message = messages[-1] if messages else None
+            natural_result = ""
+            
+            if last_message:
+                if hasattr(last_message, "content"):
+                    natural_result = last_message.content
+                elif isinstance(last_message, dict):
+                    natural_result = last_message.get("content", "")
+            
+            return {
+                "result": natural_result or "No result generated",
+                "structured_response": structured_response,
+                "metadata": {
+                    "agent": self.name,
+                    "has_structured_output": structured_response is not None,
+                    "method": "create_agent"
+                }
+            }
+        
+        elif self._method == "with_structured_output":
+            # Use with_structured_output approach
+            structured_response = self.structured_model.invoke(query)
+            
+            return {
+                "result": str(structured_response),
+                "structured_response": structured_response,
+                "metadata": {
+                    "agent": self.name,
+                    "has_structured_output": True,
+                    "method": "with_structured_output"
+                }
+            }
+        
+        else:
+            raise ValueError(f"Unknown method: {self._method}")
+    
+    def invoke(self, input: Dict[str, Any], **kwargs) -> Dict[str, Any]:
+        """
+        Direct invoke method for compatibility with LangChain agents.
+        
+        Args:
+            input: Input dictionary with messages
+            **kwargs: Additional parameters
+            
+        Returns:
+            Agent result with structured_response if available
+        """
+        if self._method == "create_agent":
+            return self.agent.invoke(input, **kwargs)
+        else:
+            query = input.get("messages", [{}])[-1].get("content", "") if isinstance(input.get("messages"), list) else str(input)
+            return self.run(query, **kwargs)

agentensemble/orchestration/__init__.py

@@ -0,0 +1,22 @@
+"""
+Orchestration Patterns
+
+Provides different patterns for coordinating multiple agents:
+- Ensemble: Full multi-agent coordination
+- Supervisor: Central coordinator pattern
+- Swarm: Decentralized collaboration
+- Pipeline: Sequential workflows
+"""
+
+from agentensemble.orchestration.ensemble import Ensemble
+from agentensemble.orchestration.supervisor import SupervisorOrchestrator
+from agentensemble.orchestration.swarm import SwarmOrchestrator
+from agentensemble.orchestration.pipeline import PipelineOrchestrator
+
+__all__ = [
+    "Ensemble",
+    "SupervisorOrchestrator",
+    "SwarmOrchestrator",
+    "PipelineOrchestrator",
+]
+

agentensemble/orchestration/ensemble.py

@@ -0,0 +1,145 @@
+"""
+Ensemble Orchestration Pattern
+
+Coordinates multiple agents working together in harmony.
+"""
+
+from typing import Any, Dict, List, Optional
+from agentensemble.agents.base import BaseAgent
+
+
+class Ensemble:
+    """
+    Ensemble orchestrator for coordinating multiple agents.
+    
+    Supports different coordination modes:
+    - supervisor: Central coordinator manages agents
+    - swarm: Decentralized agent collaboration
+    - pipeline: Sequential agent workflows
+    """
+    
+    def __init__(
+        self,
+        agents: Dict[str, BaseAgent],
+        conductor: str = "supervisor",
+        **kwargs
+    ):
+        """
+        Initialize ensemble.
+        
+        Args:
+            agents: Dictionary of agent name -> agent instance
+            conductor: Coordination mode ("supervisor", "swarm", "pipeline")
+            **kwargs: Additional configuration
+        """
+        self.agents = agents
+        self.conductor = conductor
+        self.config = kwargs
+    
+    def perform(self, task: str, data: Any = None, **kwargs) -> Dict[str, Any]:
+        """
+        Execute ensemble of agents on a task.
+        
+        Args:
+            task: Task description
+            data: Input data
+            **kwargs: Additional parameters
+            
+        Returns:
+            Combined results from all agents
+        """
+        if self.conductor == "supervisor":
+            return self._supervisor_coordinate(task, data, **kwargs)
+        elif self.conductor == "swarm":
+            return self._swarm_coordinate(task, data, **kwargs)
+        elif self.conductor == "pipeline":
+            return self._pipeline_coordinate(task, data, **kwargs)
+        else:
+            raise ValueError(f"Unknown conductor mode: {self.conductor}")
+    
+    def _supervisor_coordinate(
+        self,
+        task: str,
+        data: Any,
+        **kwargs
+    ) -> Dict[str, Any]:
+        """Supervisor pattern: Central agent coordinates others"""
+        results = {}
+        
+        # Determine which agents to use
+        agent_order = self._determine_agent_order(task)
+        
+        for agent_name in agent_order:
+            if agent_name in self.agents:
+                agent = self.agents[agent_name]
+                # Ensure context is always a dict
+                context = data if isinstance(data, dict) else {"query": str(data), "original_data": data}
+                result = agent.run(task, context=context, **kwargs)
+                results[agent_name] = result
+                
+                # Update context for next agent (keep as dict)
+                data = {"previous_result": result.get("result", ""), "metadata": result.get("metadata", {})}
+        
+        return {
+            "results": results,
+            "conductor": "supervisor",
+            "agents_used": agent_order,
+        }
+    
+    def _swarm_coordinate(
+        self,
+        task: str,
+        data: Any,
+        **kwargs
+    ) -> Dict[str, Any]:
+        """Swarm pattern: Agents collaborate independently"""
+        results = {}
+        
+        # All agents work in parallel (simplified - would use async in production)
+        context = data if isinstance(data, dict) else {"query": str(data), "original_data": data}
+        for agent_name, agent in self.agents.items():
+            result = agent.run(task, context=context, **kwargs)
+            results[agent_name] = result
+        
+        # Combine results
+        return {
+            "results": results,
+            "conductor": "swarm",
+            "agents_used": list(self.agents.keys()),
+        }
+    
+    def _pipeline_coordinate(
+        self,
+        task: str,
+        data: Any,
+        **kwargs
+    ) -> Dict[str, Any]:
+        """Pipeline pattern: Sequential agent execution"""
+        results = {}
+        current_data = data
+        
+        # Execute agents in order
+        for agent_name, agent in self.agents.items():
+            # Ensure context is always a dict
+            context = current_data if isinstance(current_data, dict) else {"query": str(current_data), "original_data": current_data}
+            result = agent.run(task, context=context, **kwargs)
+            results[agent_name] = result
+            # Update for next agent (keep as dict)
+            current_data = {"previous_result": result.get("result", ""), "metadata": result.get("metadata", {})}
+        
+        return {
+            "results": results,
+            "conductor": "pipeline",
+            "agents_used": list(self.agents.keys()),
+            "final_result": current_data,
+        }
+    
+    def _determine_agent_order(self, task: str) -> List[str]:
+        """
+        Determine which agents to use and in what order.
+        
+        In production, this would use an LLM or routing logic.
+        """
+        # Simplified: use all agents in order
+        return list(self.agents.keys())
+

agentensemble/orchestration/pipeline.py

@@ -0,0 +1,27 @@
+"""
+Pipeline Orchestrator
+
+Sequential agent workflow pattern.
+"""
+
+from typing import Any, Dict
+from agentensemble.orchestration.ensemble import Ensemble
+
+
+class PipelineOrchestrator(Ensemble):
+    """
+    Pipeline orchestrator - sequential workflow pattern.
+    
+    Agents execute in a defined sequence, passing results to the next agent.
+    """
+    
+    def __init__(self, agents: Dict[str, Any], **kwargs):
+        """
+        Initialize pipeline orchestrator.
+        
+        Args:
+            agents: Dictionary of agent name -> agent instance (order matters)
+            **kwargs: Additional configuration
+        """
+        super().__init__(agents, conductor="pipeline", **kwargs)
+

agentensemble/orchestration/supervisor.py

@@ -0,0 +1,27 @@
+"""
+Supervisor Orchestrator
+
+Central coordinator pattern for managing specialized agents.
+"""
+
+from typing import Any, Dict, List
+from agentensemble.orchestration.ensemble import Ensemble
+
+
+class SupervisorOrchestrator(Ensemble):
+    """
+    Supervisor orchestrator - central coordinator pattern.
+    
+    A supervisor agent manages and delegates tasks to specialized agents.
+    """
+    
+    def __init__(self, agents: Dict[str, Any], **kwargs):
+        """
+        Initialize supervisor orchestrator.
+        
+        Args:
+            agents: Dictionary of agent name -> agent instance
+            **kwargs: Additional configuration
+        """
+        super().__init__(agents, conductor="supervisor", **kwargs)
+

agentensemble/orchestration/swarm.py

@@ -0,0 +1,27 @@
+"""
+Swarm Orchestrator
+
+Decentralized agent collaboration pattern.
+"""
+
+from typing import Any, Dict
+from agentensemble.orchestration.ensemble import Ensemble
+
+
+class SwarmOrchestrator(Ensemble):
+    """
+    Swarm orchestrator - decentralized collaboration pattern.
+    
+    Agents operate independently and collaborate organically.
+    """
+    
+    def __init__(self, agents: Dict[str, Any], **kwargs):
+        """
+        Initialize swarm orchestrator.
+        
+        Args:
+            agents: Dictionary of agent name -> agent instance
+            **kwargs: Additional configuration
+        """
+        super().__init__(agents, conductor="swarm", **kwargs)
+

agentensemble/py.typed

@@ -0,0 +1,3 @@
+# Marker file for PEP 561
+# Indicates this package supports type checking
+

agentensemble/testing/__init__.py

@@ -0,0 +1,16 @@
+"""
+Testing & Comparison Framework
+
+Provides tools for benchmarking and comparing agents.
+"""
+
+from agentensemble.testing.comparison import AgentComparison
+from agentensemble.testing.benchmark import Benchmark
+from agentensemble.testing.metrics import Metrics
+
+__all__ = [
+    "AgentComparison",
+    "Benchmark",
+    "Metrics",
+]
+

agentensemble/testing/benchmark.py

@@ -0,0 +1,87 @@
+"""
+Benchmark Suite
+
+Pre-defined test cases for agent evaluation.
+"""
+
+from typing import Any, Dict, List
+
+
+class Benchmark:
+    """
+    Benchmark suite for agent testing.
+    
+    Provides pre-defined test cases for different scenarios.
+    """
+    
+    def __init__(self, test_cases: List[Dict[str, Any]]):
+        """
+        Initialize benchmark.
+        
+        Args:
+            test_cases: List of test case dictionaries
+        """
+        self.test_cases = test_cases
+    
+    @classmethod
+    def research_tasks(cls) -> "Benchmark":
+        """
+        Research tasks benchmark.
+        
+        Returns:
+            Benchmark with research-oriented test cases
+        """
+        test_cases = [
+            {
+                "query": "Find the CEO of Stripe",
+                "expected_type": "person_name",
+                "category": "research",
+            },
+            {
+                "query": "What year was OpenAI founded?",
+                "expected_type": "year",
+                "category": "research",
+            },
+            {
+                "query": "Find the official email of Anthropic",
+                "expected_type": "email",
+                "category": "research",
+            },
+        ]
+        return cls(test_cases)
+    
+    @classmethod
+    def data_extraction_tasks(cls) -> "Benchmark":
+        """
+        Data extraction tasks benchmark.
+        
+        Returns:
+            Benchmark with extraction-oriented test cases
+        """
+        test_cases = [
+            {
+                "query": "Extract company name and industry from: Stripe is a payments company",
+                "expected_type": "structured_data",
+                "category": "extraction",
+            },
+        ]
+        return cls(test_cases)
+    
+    @classmethod
+    def load(cls, name: str) -> "Benchmark":
+        """
+        Load a benchmark by name.
+        
+        Args:
+            name: Benchmark name
+            
+        Returns:
+            Benchmark instance
+        """
+        if name == "research_tasks":
+            return cls.research_tasks()
+        elif name == "data_extraction_tasks":
+            return cls.data_extraction_tasks()
+        else:
+            raise ValueError(f"Unknown benchmark: {name}")
+

agentensemble/testing/comparison.py

@@ -0,0 +1,82 @@
+"""
+Agent Comparison
+
+Compare multiple agent implementations on the same tasks.
+"""
+
+from typing import Any, Dict, List, Optional
+from agentensemble.agents.base import BaseAgent
+
+
+class AgentComparison:
+    """
+    Compare multiple agent implementations.
+    
+    Runs the same tasks on different agents and compares results.
+    """
+    
+    def __init__(self, agents: List[BaseAgent]):
+        """
+        Initialize agent comparison.
+        
+        Args:
+            agents: List of agent instances to compare
+        """
+        self.agents = agents
+    
+    def run(
+        self,
+        benchmark: Any,
+        metrics: Optional[List[str]] = None,
+        **kwargs
+    ) -> Dict[str, Any]:
+        """
+        Run comparison benchmark.
+        
+        Args:
+            benchmark: Benchmark instance with test cases
+            metrics: List of metrics to track
+            **kwargs: Additional parameters
+            
+        Returns:
+            Comparison results
+        """
+        metrics = metrics or ["success_rate", "execution_time", "cost"]
+        results = {}
+        
+        for agent in self.agents:
+            agent_results = []
+            
+            for test_case in benchmark.test_cases:
+                result = agent.run(test_case["query"], **kwargs)
+                agent_results.append({
+                    "test_case": test_case,
+                    "result": result,
+                })
+            
+            results[agent.name] = agent_results
+        
+        return {
+            "results": results,
+            "metrics": metrics,
+            "summary": self._calculate_summary(results, metrics),
+        }
+    
+    def _calculate_summary(
+        self,
+        results: Dict[str, List[Dict]],
+        metrics: List[str]
+    ) -> Dict[str, Any]:
+        """Calculate summary statistics"""
+        summary = {}
+        
+        for agent_name, agent_results in results.items():
+            summary[agent_name] = {
+                "total_tests": len(agent_results),
+                "success_rate": 0.8,  # Placeholder
+                "avg_execution_time": 2.5,  # Placeholder
+                "total_cost": 0.05,  # Placeholder
+            }
+        
+        return summary
+