traia-iatp 0.1.1__py3-none-any.whl

traia_iatp/special_agencies/registry_search_agency.py

@@ -0,0 +1,392 @@
+"""
+Registry Search Utility Agency
+
+A special utility agency that provides search capabilities for finding other utility agencies.
+This agency exposes MongoDB registry search functionality via the A2A protocol.
+"""
+
+import os
+import json
+import logging
+from typing import List, Dict, Any, Optional
+from fastapi import FastAPI, HTTPException
+from a2a import A2AServer, ToolResult
+import uvicorn
+from datetime import datetime
+
+from ..registry.mongodb_registry import UtilityAgencyRegistry
+from ..utils import get_now_in_utc
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Create FastAPI app
+app = FastAPI(title="Registry Search Utility Agency")
+
+# Create A2A server
+a2a_server = A2AServer(
+    name="Registry Search Agency",
+    description="Searches for utility agencies in the IATP registry",
+    version="1.0.0"
+)
+
+# MongoDB registry instance
+_registry: Optional[UtilityAgencyRegistry] = None
+
+
+def get_registry() -> UtilityAgencyRegistry:
+    """Get or create MongoDB registry instance."""
+    global _registry
+    if _registry is None:
+        mongodb_uri = os.getenv("MONGODB_URI")
+        if not mongodb_uri:
+            raise ValueError("MONGODB_URI environment variable is required")
+        _registry = UtilityAgencyRegistry(mongodb_uri)
+    return _registry
+
+
+@a2a_server.tool(
+    name="search_utility_agencies",
+    description="Search for utility agencies by query, tags, or capabilities"
+)
+async def search_utility_agencies(
+    query: Optional[str] = None,
+    tags: Optional[List[str]] = None,
+    capabilities: Optional[List[str]] = None,
+    limit: int = 10
+) -> ToolResult:
+    """Search for utility agencies in the registry."""
+    try:
+        registry = get_registry()
+        
+        # Query agencies
+        agencies = await registry.query_agencies(
+            query=query,
+            tags=tags,
+            capabilities=capabilities,
+            active_only=True,
+            limit=limit
+        )
+        
+        # Convert to simple dict format
+        results = []
+        for agency in agencies:
+            results.append({
+                "agency_id": agency.agency_id,
+                "name": agency.name,
+                "description": agency.description,
+                "endpoint": str(agency.endpoint),
+                "capabilities": agency.capabilities,
+                "tags": agency.tags,
+                "is_active": agency.is_active
+            })
+        
+        return ToolResult(
+            output=results,
+            metadata={
+                "count": len(results),
+                "query": query,
+                "tags": tags,
+                "capabilities": capabilities
+            }
+        )
+    except Exception as e:
+        logger.error(f"Error searching agencies: {e}")
+        raise
+
+
+@a2a_server.tool(
+    name="get_agency_by_id",
+    description="Get detailed information about a specific utility agency"
+)
+async def get_agency_by_id(agency_id: str) -> ToolResult:
+    """Get a specific agency by ID."""
+    try:
+        registry = get_registry()
+        
+        agency = await registry.get_agency_by_id(agency_id)
+        if not agency:
+            return ToolResult(
+                output=None,
+                metadata={"error": f"Agency {agency_id} not found"}
+            )
+        
+        result = {
+            "agency_id": agency.agency_id,
+            "name": agency.name,
+            "description": agency.description,
+            "endpoint": str(agency.endpoint),
+            "capabilities": agency.capabilities,
+            "tags": agency.tags,
+            "metadata": agency.metadata,
+            "registered_at": agency.registered_at.isoformat() if agency.registered_at else None,
+            "last_health_check": agency.last_health_check.isoformat() if agency.last_health_check else None,
+            "is_active": agency.is_active
+        }
+        
+        return ToolResult(
+            output=result,
+            metadata={"agency_id": agency_id}
+        )
+    except Exception as e:
+        logger.error(f"Error getting agency {agency_id}: {e}")
+        raise
+
+
+@a2a_server.tool(
+    name="get_registry_statistics",
+    description="Get statistics about the utility agency registry"
+)
+async def get_registry_statistics() -> ToolResult:
+    """Get registry statistics."""
+    try:
+        registry = get_registry()
+        
+        stats = await registry.get_statistics()
+        
+        return ToolResult(
+            output=stats,
+            metadata={"timestamp": get_now_in_utc().isoformat()}
+        )
+    except Exception as e:
+        logger.error(f"Error getting statistics: {e}")
+        raise
+
+
+@a2a_server.tool(
+    name="find_agencies_by_capability",
+    description="Find all agencies that have a specific capability"
+)
+async def find_agencies_by_capability(capability: str, limit: int = 20) -> ToolResult:
+    """Find agencies by specific capability."""
+    try:
+        registry = get_registry()
+        
+        agencies = await registry.query_agencies(
+            capabilities=[capability],
+            active_only=True,
+            limit=limit
+        )
+        
+        results = []
+        for agency in agencies:
+            results.append({
+                "agency_id": agency.agency_id,
+                "name": agency.name,
+                "description": agency.description,
+                "endpoint": str(agency.endpoint),
+                "all_capabilities": agency.capabilities
+            })
+        
+        return ToolResult(
+            output=results,
+            metadata={
+                "capability": capability,
+                "count": len(results)
+            }
+        )
+    except Exception as e:
+        logger.error(f"Error finding agencies by capability: {e}")
+        raise
+
+
+# FastAPI routes
+@app.get("/health")
+async def health_check():
+    """Health check endpoint."""
+    try:
+        # Test MongoDB connection
+        registry = get_registry()
+        stats = await registry.get_statistics()
+        return {
+            "status": "healthy",
+            "agency": "Registry Search Agency",
+            "total_agencies": stats.get("total_agencies", 0),
+            "active_agencies": stats.get("active_agencies", 0)
+        }
+    except Exception as e:
+        return {
+            "status": "unhealthy",
+            "error": str(e)
+        }
+
+
+@app.get("/info")
+async def get_info():
+    """Get agency information."""
+    return {
+        "name": "Registry Search Agency",
+        "description": "Searches for utility agencies in the IATP registry",
+        "capabilities": [
+            "search_utility_agencies",
+            "get_agency_by_id",
+            "get_registry_statistics",
+            "find_agencies_by_capability"
+        ],
+        "type": "special_agency",
+        "version": "1.0.0"
+    }
+
+
+@app.post("/a2a")
+async def a2a_endpoint(request: dict):
+    """Handle A2A protocol requests."""
+    try:
+        # Process the A2A request
+        action = request.get("action")
+        parameters = request.get("parameters", {})
+        context = request.get("context", {})
+        
+        # Map action to tool
+        tool_mapping = {
+            "search_utility_agencies": search_utility_agencies,
+            "get_agency_by_id": get_agency_by_id,
+            "get_registry_statistics": get_registry_statistics,
+            "find_agencies_by_capability": find_agencies_by_capability
+        }
+        
+        if action not in tool_mapping:
+            raise HTTPException(status_code=400, detail=f"Unknown action: {action}")
+        
+        # Execute the tool
+        tool_func = tool_mapping[action]
+        result = await tool_func(**parameters)
+        
+        return {
+            "result": result.output,
+            "status": "success",
+            "metadata": result.metadata
+        }
+        
+    except Exception as e:
+        logger.error(f"A2A request error: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.on_event("shutdown")
+async def shutdown_event():
+    """Clean up on shutdown."""
+    global _registry
+    if _registry:
+        _registry.close()
+        _registry = None
+
+
+def create_registry_search_agency_structure(output_dir: str = "special_agencies") -> str:
+    """Create the folder structure for the registry search agency."""
+    from pathlib import Path
+    
+    agency_path = Path(output_dir) / "registry_search_agency"
+    agency_path.mkdir(parents=True, exist_ok=True)
+    
+    # Create Dockerfile
+    dockerfile_content = """FROM python:3.12-slim
+
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y gcc && rm -rf /var/lib/apt/lists/*
+
+# Copy requirements
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy agency files
+COPY . .
+
+# Expose port
+EXPOSE 8000
+
+# Run the server
+CMD ["python", "-m", "uvicorn", "server:app", "--host", "0.0.0.0", "--port", "8000"]
+"""
+    
+    with open(agency_path / "Dockerfile", "w") as f:
+        f.write(dockerfile_content)
+    
+    # Create requirements.txt
+    requirements = """fastapi>=0.100.0
+uvicorn>=0.20.0
+a2a>=0.1.0
+pymongo>=4.0.0
+python-dotenv>=1.0.0
+"""
+    
+    with open(agency_path / "requirements.txt", "w") as f:
+        f.write(requirements)
+    
+    # Copy this file as server.py
+    import shutil
+    shutil.copy(__file__, agency_path / "server.py")
+    
+    # Create README.md
+    readme_content = """# Registry Search Utility Agency
+
+A special utility agency that provides search capabilities for finding other utility agencies in the IATP registry.
+
+## Features
+- Search agencies by text query
+- Filter by tags and capabilities
+- Get detailed agency information
+- View registry statistics
+
+## Environment Variables
+- `MONGODB_URI`: MongoDB connection string (required)
+
+## Running the Agency
+
+### Using Docker
+```bash
+docker build -t registry-search-agency .
+docker run -p 8000:8000 -e MONGODB_URI="your-mongodb-uri" registry-search-agency
+```
+
+### Without Docker
+```bash
+export MONGODB_URI="your-mongodb-uri"
+pip install -r requirements.txt
+python -m uvicorn server:app --host 0.0.0.0 --port 8000
+```
+
+## A2A Tools Available
+- `search_utility_agencies`: Search for agencies
+- `get_agency_by_id`: Get specific agency details
+- `get_registry_statistics`: Get registry stats
+- `find_agencies_by_capability`: Find agencies by capability
+
+## Usage Example
+
+```python
+from iatp import create_utility_agency_tools
+
+# This agency will be discoverable like any other
+tools = create_utility_agency_tools(
+    query="registry search"
+)
+
+# Use the tool to find other agencies
+result = tools[0]._run(
+    action="search_utility_agencies",
+    parameters={"query": "data processing", "limit": 5}
+)
+```
+"""
+    
+    with open(agency_path / "README.md", "w") as f:
+        f.write(readme_content)
+    
+    # Create .env.example
+    env_example = """# MongoDB connection string (required)
+MONGODB_URI=mongodb+srv://user:pass@cluster.mongodb.net/iatp
+"""
+    
+    with open(agency_path / ".env.example", "w") as f:
+        f.write(env_example)
+    
+    logger.info(f"Created registry search agency structure at: {agency_path}")
+    return str(agency_path)
+
+
+if __name__ == "__main__":
+    # Running directly - start the server
+    uvicorn.run(app, host="0.0.0.0", port=8000) 

traia_iatp/utils/__init__.py

@@ -0,0 +1,10 @@
+"""Traia IATP utility functions module."""
+
+from .general import get_now_in_utc, get_logger, is_empty, not_empty
+
+__all__ = [
+    "get_now_in_utc",
+    "get_logger", 
+    "is_empty",
+    "not_empty",
+] 

traia_iatp/utils/docker_utils.py

@@ -0,0 +1,251 @@
+"""Docker utility functions for running generated utility agents locally."""
+
+import logging
+import subprocess
+import os
+from pathlib import Path
+from typing import Dict, Any, Optional
+import docker
+import httpx
+import asyncio
+
+logger = logging.getLogger(__name__)
+
+
+class LocalDockerRunner:
+    """Run generated utility agents locally using their Docker configuration."""
+    
+    def __init__(self):
+        """Initialize Docker client."""
+        try:
+            self.client = docker.from_env()
+        except Exception as e:
+            logger.error(f"Failed to initialize Docker client: {e}")
+            raise RuntimeError("Docker is not available. Please ensure Docker is installed and running.")
+    
+    async def run_agent_docker(
+        self, 
+        agent_path: Path,
+        port: int = 8000,
+        environment: Optional[Dict[str, str]] = None,
+        detached: bool = True
+    ) -> Dict[str, Any]:
+        """Run a generated utility agent using its Docker configuration.
+        
+        Args:
+            agent_path: Path to the generated agent directory
+            port: Port to expose (default: 8000)
+            environment: Additional environment variables
+            detached: Whether to run in detached mode
+            
+        Returns:
+            Dict with deployment information
+        """
+        agent_path = Path(agent_path)
+        
+        # Verify the agent directory has required files
+        required_files = ["Dockerfile", "docker-compose.yml", "pyproject.toml"]
+        for file in required_files:
+            if not (agent_path / file).exists():
+                raise FileNotFoundError(f"Required file {file} not found in {agent_path}")
+        
+        # Create .env file if it doesn't exist
+        env_file = agent_path / ".env"
+        if not env_file.exists():
+            self._create_env_file(agent_path, port, environment)
+        
+        # Use docker-compose to run the agent
+        compose_file = agent_path / "docker-compose.yml"
+        
+        try:
+            # Build the image
+            logger.info(f"Building Docker image for agent at {agent_path}")
+            subprocess.run(
+                ["docker-compose", "-f", str(compose_file), "build"],
+                cwd=str(agent_path),
+                check=True,
+                capture_output=True,
+                text=True
+            )
+            
+            # Run the container
+            logger.info(f"Starting container on port {port}")
+            cmd = ["docker-compose", "-f", str(compose_file), "up"]
+            if detached:
+                cmd.append("-d")
+                
+            result = subprocess.run(
+                cmd,
+                cwd=str(agent_path),
+                check=True,
+                capture_output=True,
+                text=True,
+                env={**os.environ, "PORT": str(port)}
+            )
+            
+            # Get container info
+            container_name = self._get_container_name(agent_path)
+            
+            # Wait for the service to be ready
+            if detached:
+                await self._wait_for_service(port)
+            
+            return {
+                "success": True,
+                "container_name": container_name,
+                "port": port,
+                "base_url": f"http://localhost:{port}",
+                "iatp_endpoint": f"http://localhost:{port}/a2a",
+                "health_endpoint": f"http://localhost:{port}/health",
+                "logs_command": f"docker logs -f {container_name}",
+                "stop_command": f"docker-compose -f {compose_file} down"
+            }
+            
+        except subprocess.CalledProcessError as e:
+            logger.error(f"Docker command failed: {e.stderr}")
+            raise RuntimeError(f"Failed to run agent: {e.stderr}")
+        except Exception as e:
+            logger.error(f"Error running agent: {e}")
+            raise
+    
+    async def stop_agent_docker(self, agent_path: Path) -> bool:
+        """Stop a running agent.
+        
+        Args:
+            agent_path: Path to the agent directory
+            
+        Returns:
+            True if stopped successfully
+        """
+        agent_path = Path(agent_path)
+        compose_file = agent_path / "docker-compose.yml"
+        
+        if not compose_file.exists():
+            raise FileNotFoundError(f"docker-compose.yml not found in {agent_path}")
+        
+        try:
+            subprocess.run(
+                ["docker-compose", "-f", str(compose_file), "down"],
+                cwd=str(agent_path),
+                check=True,
+                capture_output=True,
+                text=True
+            )
+            logger.info(f"Stopped agent at {agent_path}")
+            return True
+        except subprocess.CalledProcessError as e:
+            logger.error(f"Failed to stop agent: {e.stderr}")
+            return False
+    
+    def get_agent_logs(self, agent_path: Path, tail: int = 100) -> str:
+        """Get logs from a running agent.
+        
+        Args:
+            agent_path: Path to the agent directory
+            tail: Number of lines to return
+            
+        Returns:
+            Log output as string
+        """
+        container_name = self._get_container_name(agent_path)
+        
+        try:
+            container = self.client.containers.get(container_name)
+            return container.logs(tail=tail).decode('utf-8')
+        except docker.errors.NotFound:
+            return f"Container {container_name} not found"
+        except Exception as e:
+            return f"Error getting logs: {e}"
+    
+    def _create_env_file(self, agent_path: Path, port: int, environment: Optional[Dict[str, str]] = None):
+        """Create a .env file for the agent."""
+        env_content = [
+            f"PORT={port}",
+            "HOST=0.0.0.0",
+            "LOG_LEVEL=INFO",
+            ""
+        ]
+        
+        # Add custom environment variables
+        if environment:
+            for key, value in environment.items():
+                env_content.append(f"{key}={value}")
+        
+        env_file = agent_path / ".env"
+        env_file.write_text("\n".join(env_content))
+        logger.info(f"Created .env file at {env_file}")
+    
+    def _get_container_name(self, agent_path: Path) -> str:
+        """Get the container name from the agent directory name."""
+        return f"{agent_path.name}-local"
+    
+    async def _wait_for_service(self, port: int, timeout: int = 30):
+        """Wait for the service to be ready."""
+        start_time = asyncio.get_event_loop().time()
+        
+        async with httpx.AsyncClient() as client:
+            while asyncio.get_event_loop().time() - start_time < timeout:
+                try:
+                    response = await client.get(f"http://localhost:{port}/health")
+                    if response.status_code == 200:
+                        logger.info(f"Service is ready on port {port}")
+                        return
+                except Exception:
+                    pass
+                await asyncio.sleep(1)
+        
+        logger.warning(f"Service did not become ready within {timeout} seconds")
+
+
+def run_generated_agent_locally(
+    agent_path: str,
+    port: int = 8000,
+    environment: Optional[Dict[str, str]] = None
+) -> Dict[str, Any]:
+    """Convenience function to run a generated agent locally.
+    
+    This is a synchronous wrapper around the async functionality.
+    
+    Args:
+        agent_path: Path to the generated agent directory
+        port: Port to expose
+        environment: Additional environment variables
+        
+    Returns:
+        Dict with deployment information
+    """
+    runner = LocalDockerRunner()
+    return asyncio.run(runner.run_agent_docker(
+        Path(agent_path),
+        port=port,
+        environment=environment
+    ))
+
+
+def use_run_local_docker_script(agent_path: str) -> subprocess.CompletedProcess:
+    """Run the agent using its generated run_local_docker.sh script.
+    
+    This is the simplest way to run a generated agent locally,
+    as it uses the script generated from templates.
+    
+    Args:
+        agent_path: Path to the generated agent directory
+        
+    Returns:
+        CompletedProcess result
+    """
+    agent_path = Path(agent_path)
+    script_path = agent_path / "run_local_docker.sh"
+    
+    if not script_path.exists():
+        raise FileNotFoundError(f"run_local_docker.sh not found in {agent_path}")
+    
+    # Make sure the script is executable
+    script_path.chmod(0o755)
+    
+    # Run the script
+    return subprocess.run(
+        ["./run_local_docker.sh"],
+        cwd=str(agent_path),
+        check=True
+    )