agent_orchestrator.py

#!/usr/bin/env python3
"""
agent orchestrator - demonstrates agent-to-agent communication using fasterpc
"""

import asyncio
import time
import json
from typing import Dict, List, Any
from fasterpc import rpc_methods_base, websocket_rpc_client, logging_config, LoggingModes

# configure logging
logging_config.set_mode(LoggingModes.UVICORN)

class agent_orchestrator:
    def __init__(self):
        self.agent_endpoints = {
            "coordination": "ws://localhost:9000/ws",
            "research": "ws://localhost:9001/ws",
            "analysis": "ws://localhost:9002/ws",
            "writing": "ws://localhost:9003/ws",
            "storage": "ws://localhost:9004/ws"
        }
        self.workflow_results = []
        
    async def test_agent_connection(self, agent_type: str):
        """test connection to a specific agent"""
        endpoint = self.agent_endpoints.get(agent_type)
        if not endpoint:
            print(f"   ❌ {agent_type} agent endpoint not found")
            return False
            
        try:
            async with websocket_rpc_client(endpoint, rpc_methods_base()) as client:
                if hasattr(client.other, 'get_agent_info'):
                    await client.other.get_agent_info()
                    print(f"   ✅ {agent_type} agent is available")
                    return True
                else:
                    print(f"   ⚠️  {agent_type} agent connected but no get_agent_info method")
                    return True
        except Exception as e:
            print(f"   ❌ failed to connect to {agent_type} agent: {e}")
            return False
    
    async def test_agent_status(self):
        """test agent status checking"""
        print("\n🔍 testing agent status...")
        
        coordination_endpoint = self.agent_endpoints.get("coordination")
        if not coordination_endpoint:
            print("   ❌ coordination agent endpoint not found")
            return
        
        try:
            async with websocket_rpc_client(coordination_endpoint, rpc_methods_base()) as coordination:
                # check status of each agent
                for agent_type in ["research", "analysis", "writing", "storage"]:
                    try:
                        status = await coordination.other.get_agent_status(agent_type=agent_type)
                        print(f"   📊 {agent_type} agent status: {status.result['status']}")
                    except Exception as e:
                        print(f"   ❌ error checking {agent_type} agent status: {e}")
        except Exception as e:
            print(f"   ❌ error connecting to coordination agent: {e}")
    
    async def test_research_workflow(self):
        """test a complete research workflow"""
        print("\n🔬 testing research workflow...")
        
        # get endpoints for all required agents
        research_endpoint = self.agent_endpoints.get("research")
        analysis_endpoint = self.agent_endpoints.get("analysis")
        writing_endpoint = self.agent_endpoints.get("writing")
        storage_endpoint = self.agent_endpoints.get("storage")
        
        if not all([research_endpoint, analysis_endpoint, writing_endpoint, storage_endpoint]):
            print("   ❌ not all required agent endpoints found")
            return
        
        try:
            # step 1: research agent searches for information
            print("   1️⃣ research agent searching for 'python websockets'...")
            async with websocket_rpc_client(research_endpoint, rpc_methods_base()) as research:
                search_result = await research.other.search_topic(topic="python websockets")
                print(f"      found {len(search_result.result['sources'])} sources")
            
            # step 2: analysis agent analyzes the search results
            print("   2️⃣ analysis agent analyzing search results...")
            async with websocket_rpc_client(analysis_endpoint, rpc_methods_base()) as analysis:
                analysis_result = await analysis.other.analyze_data(data=search_result.result)
                print(f"      found {len(analysis_result.result['patterns_found'])} patterns")
            
            # step 3: writing agent creates a summary
            print("   3️⃣ writing agent creating summary...")
            summary_data = {
                "topic": search_result.result["topic"],
                "sources": search_result.result["sources"],
                "insights": analysis_result.result.get("insights", [])
            }
            async with websocket_rpc_client(writing_endpoint, rpc_methods_base()) as writing:
                summary_result = await writing.other.write_summary(data=summary_data)
                print(f"      created summary with {summary_result.result['word_count']} words")
            
            # step 4: storage agent saves the results
            print("   4️⃣ storage agent saving results...")
            storage_data = {
                "search_results": search_result.result,
                "analysis_results": analysis_result.result,
                "summary": summary_result.result,
                "workflow_timestamp": time.time()
            }
            async with websocket_rpc_client(storage_endpoint, rpc_methods_base()) as storage:
                save_result = await storage.other.save_data(key="research_workflow_001", data=storage_data)
                print(f"      saved data with key: {save_result.result['key']}")
            
            # record workflow result
            workflow_result = {
                "workflow_type": "research",
                "steps_completed": 4,
                "final_result": save_result.result,
                "timestamp": time.time()
            }
            self.workflow_results.append(workflow_result)
            
            print("   ✅ research workflow completed successfully!")
            
        except Exception as e:
            print(f"   ❌ research workflow failed: {e}")
    
    async def test_agent_to_agent_communication(self):
        """test direct agent-to-agent communication"""
        print("\n🤝 testing agent-to-agent communication...")
        
        research_endpoint = self.agent_endpoints.get("research")
        analysis_endpoint = self.agent_endpoints.get("analysis")
        
        if not all([research_endpoint, analysis_endpoint]):
            print("   ❌ required agent endpoints not found")
            return
        
        try:
            # research agent searches for information
            print("   🔍 research agent searching...")
            async with websocket_rpc_client(research_endpoint, rpc_methods_base()) as research:
                search_result = await research.other.search_topic(topic="artificial intelligence")
            
            # analysis agent directly processes research results
            print("   📊 analysis agent processing research results...")
            async with websocket_rpc_client(analysis_endpoint, rpc_methods_base()) as analysis:
                analysis_result = await analysis.other.generate_insights(
                    text=f"research on {search_result.result['topic']} found {len(search_result.result['sources'])} sources"
                )
            
            print(f"   ✅ agent-to-agent communication successful!")
            print(f"      research found {len(search_result.result['sources'])} sources")
            print(f"      analysis extracted {len(analysis_result.result['key_terms'])} key terms")
            
        except Exception as e:
            print(f"   ❌ agent-to-agent communication failed: {e}")
    
    async def test_complex_workflow(self):
        """test a complex multi-agent workflow"""
        print("\n🔄 testing complex workflow...")
        
        coordination_endpoint = self.agent_endpoints.get("coordination")
        if not coordination_endpoint:
            print("   ❌ coordination agent endpoint not found")
            return
        
        # define a complex workflow
        workflow = {
            "name": "comprehensive_analysis_workflow",
            "steps": [
                {
                    "action": "research_topic",
                    "agent_type": "research",
                    "parameters": {"topic": "machine learning"}
                },
                {
                    "action": "analyze_sources",
                    "agent_type": "analysis",
                    "parameters": {"analysis_type": "source_comparison"}
                },
                {
                    "action": "generate_report",
                    "agent_type": "writing",
                    "parameters": {"report_type": "technical"}
                },
                {
                    "action": "store_results",
                    "agent_type": "storage",
                    "parameters": {"storage_key": "ml_analysis_001"}
                }
            ]
        }
        
        try:
            print("   🎯 coordinating complex workflow...")
            async with websocket_rpc_client(coordination_endpoint, rpc_methods_base()) as coordination:
                workflow_result = await coordination.other.coordinate_workflow(workflow=workflow)
            
            print(f"   ✅ complex workflow completed!")
            print(f"      workflow id: {workflow_result.result['workflow_id']}")
            print(f"      status: {workflow_result.result['status']}")
            print(f"      duration: {workflow_result.result['duration']:.2f}s")
            
            # record workflow result
            self.workflow_results.append(workflow_result.result)
            
        except Exception as e:
            print(f"   ❌ complex workflow failed: {e}")
    
    async def test_performance_benchmark(self):
        """test performance of agent communication"""
        print("\n⚡ testing performance benchmark...")
        
        research_endpoint = self.agent_endpoints.get("research")
        if not research_endpoint:
            print("   ❌ research agent endpoint not found")
            return
        
        # benchmark multiple concurrent calls
        start_time = time.time()
        iterations = 10  # reduced for demo
        
        print(f"   🏃 running {iterations} concurrent research calls...")
        
        async def single_call(i):
            try:
                async with websocket_rpc_client(research_endpoint, rpc_methods_base()) as research:
                    return await research.other.search_topic(topic=f"test_topic_{i}")
            except Exception as e:
                return e
        
        tasks = [single_call(i) for i in range(iterations)]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        end_time = time.time()
        
        duration = end_time - start_time
        calls_per_second = iterations / duration
        success_count = sum(1 for r in results if not isinstance(r, Exception))
        
        print(f"   📊 performance results:")
        print(f"      completed {iterations} calls in {duration:.2f}s")
        print(f"      rate: {calls_per_second:.2f} calls/second")
        print(f"      success rate: {success_count}/{iterations} ({success_count/iterations*100:.1f}%)")
    
    async def get_all_agent_stats(self):
        """get statistics from all agents"""
        print("\n📊 collecting agent statistics...")
        
        stats = {}
        
        for agent_type, endpoint in self.agent_endpoints.items():
            try:
                async with websocket_rpc_client(endpoint, rpc_methods_base()) as client:
                    if hasattr(client.other, 'get_agent_info'):
                        agent_info = await client.other.get_agent_info()
                        stats[agent_type] = agent_info.result
                        print(f"   📈 {agent_type} agent: {agent_info.result.get('agent_id', 'unknown')}")
            except Exception as e:
                print(f"   ❌ error getting {agent_type} agent stats: {e}")
        
        return stats
    
    async def run_demo(self):
        """run the complete agent orchestration demo"""
        print("🚀 starting agent orchestration demo...")
        print("=" * 50)
        
        try:
            # test connections to all agents
            print("🔌 testing agent connections...")
            for agent_type in self.agent_endpoints.keys():
                await self.test_agent_connection(agent_type)
            
            # run various tests
            await self.test_agent_status()
            await self.test_research_workflow()
            await self.test_agent_to_agent_communication()
            await self.test_complex_workflow()
            await self.test_performance_benchmark()
            
            # collect statistics
            stats = await self.get_all_agent_stats()
            
            # summary
            print("\n📋 demo summary:")
            print(f"   available agents: {len(self.agent_endpoints)}")
            print(f"   completed workflows: {len(self.workflow_results)}")
            print(f"   total agent stats collected: {len(stats)}")
            
            print("\n✅ agent orchestration demo completed successfully!")
            
        except Exception as e:
            print(f"\n❌ demo failed: {e}")
            raise

async def main():
    """main function"""
    orchestrator = agent_orchestrator()
    await orchestrator.run_demo()

if __name__ == "__main__":
    print("🤖 agent orchestrator demo")
    print("make sure all agent servers are running:")
    print("  - coordination agent: python coordination_agent.py")
    print("  - research agent: python research_agent.py")
    print("  - analysis agent: python analysis_agent.py")
    print("  - writing agent: python writing_agent.py")
    print("  - storage agent: python storage_agent.py")
    print()
    
    asyncio.run(main())