""" WHOOSH API - Agent Management Endpoints This module provides comprehensive API endpoints for managing Ollama-based AI agents in the WHOOSH distributed orchestration platform. It handles agent registration, status monitoring, and lifecycle management. Key Features: - Agent registration and validation - Real-time status monitoring - Comprehensive error handling - Detailed API documentation - Authentication and authorization """ from fastapi import APIRouter, HTTPException, Request, Depends, status from typing import List, Dict, Any import time import logging from ..models.agent import Agent from ..models.responses import ( AgentListResponse, AgentRegistrationResponse, AgentRegistrationRequest, ErrorResponse, AgentModel ) from ..core.auth_deps import get_current_user_context router = APIRouter() from app.core.database import SessionLocal from app.models.agent import Agent as ORMAgent from ..services.agent_service import AgentType logger = logging.getLogger(__name__) @router.get( "/agents", response_model=AgentListResponse, status_code=status.HTTP_200_OK, summary="List all registered agents", description=""" Retrieve a comprehensive list of all registered agents in the WHOOSH cluster. This endpoint returns detailed information about each agent including: - Agent identification and endpoint information - Current status and utilization metrics - Specialization and capacity limits - Health and heartbeat information **Use Cases:** - Monitor cluster capacity and agent health - Identify available agents for task assignment - Track agent utilization and performance - Debug agent connectivity issues **Response Notes:** - Agents are returned in registration order - Status reflects real-time agent availability - Utilization is calculated as current_tasks / max_concurrent """, responses={ 200: {"description": "List of agents retrieved successfully"}, 500: {"model": ErrorResponse, "description": "Internal server error"} } ) async def get_agents( request: Request, current_user: Dict[str, Any] = Depends(get_current_user_context) ) -> AgentListResponse: """ Get all registered agents with detailed status information. Returns: AgentListResponse: Comprehensive list of all registered agents Raises: HTTPException: If database query fails """ try: with SessionLocal() as db: db_agents = db.query(ORMAgent).all() agents_list = [] for db_agent in db_agents: agent_model = AgentModel( id=db_agent.id, endpoint=db_agent.endpoint, model=db_agent.model, specialty=db_agent.specialty, max_concurrent=db_agent.max_concurrent, current_tasks=db_agent.current_tasks, status="available" if db_agent.current_tasks < db_agent.max_concurrent else "busy", utilization=db_agent.current_tasks / db_agent.max_concurrent if db_agent.max_concurrent > 0 else 0.0 ) agents_list.append(agent_model) return AgentListResponse( agents=agents_list, total=len(agents_list), message=f"Retrieved {len(agents_list)} registered agents" ) except Exception as e: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to retrieve agents: {str(e)}" ) @router.post( "/agents", response_model=AgentRegistrationResponse, status_code=status.HTTP_201_CREATED, summary="Register a new Ollama agent", description=""" Register a new Ollama-based AI agent with the WHOOSH cluster. This endpoint allows you to add new Ollama agents to the distributed AI network. The agent will be validated for connectivity and model availability before registration. **Agent Registration Process:** 1. Validate agent connectivity and model availability 2. Add agent to the coordinator's active agent pool 3. Store agent configuration in the database 4. Perform initial health check 5. Return registration confirmation with agent details **Supported Agent Specializations:** - `kernel_dev`: Linux kernel development and debugging - `pytorch_dev`: PyTorch model development and optimization - `profiler`: Performance profiling and optimization - `docs_writer`: Documentation generation and technical writing - `tester`: Automated testing and quality assurance - `general_ai`: General-purpose AI assistance - `reasoning`: Complex reasoning and problem-solving tasks **Requirements:** - Agent endpoint must be accessible from the WHOOSH cluster - Specified model must be available on the target Ollama instance - Agent ID must be unique across the cluster """, responses={ 201: {"description": "Agent registered successfully"}, 400: {"model": ErrorResponse, "description": "Invalid agent configuration"}, 409: {"model": ErrorResponse, "description": "Agent ID already exists"}, 503: {"model": ErrorResponse, "description": "Agent endpoint unreachable"} } ) async def register_agent( agent_data: AgentRegistrationRequest, request: Request, current_user: Dict[str, Any] = Depends(get_current_user_context) ) -> AgentRegistrationResponse: """ Register a new Ollama agent in the WHOOSH cluster. Args: agent_data: Agent configuration and registration details request: FastAPI request object for accessing app state current_user: Current authenticated user context Returns: AgentRegistrationResponse: Registration confirmation with agent details Raises: HTTPException: If registration fails due to validation or connectivity issues """ # Access coordinator through the dependency injection whoosh_coordinator = getattr(request.app.state, 'whoosh_coordinator', None) if not whoosh_coordinator: # Fallback to global coordinator if app state not available from ..main import unified_coordinator whoosh_coordinator = unified_coordinator if not whoosh_coordinator: raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Coordinator service unavailable" ) try: # Check if agent ID already exists with SessionLocal() as db: existing_agent = db.query(ORMAgent).filter(ORMAgent.id == agent_data.id).first() if existing_agent: raise HTTPException( status_code=status.HTTP_409_CONFLICT, detail=f"Agent with ID '{agent_data.id}' already exists" ) # Create agent instance agent = Agent( id=agent_data.id, endpoint=agent_data.endpoint, model=agent_data.model, specialty=AgentType(agent_data.specialty.value), max_concurrent=agent_data.max_concurrent, ) # Add agent to coordinator whoosh_coordinator.add_agent(agent) return AgentRegistrationResponse( agent_id=agent.id, endpoint=agent.endpoint, message=f"Agent '{agent.id}' registered successfully with specialty '{agent_data.specialty}'" ) except ValueError as e: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"Invalid agent configuration: {str(e)}" ) except Exception as e: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to register agent: {str(e)}" ) @router.get( "/agents/{agent_id}", response_model=AgentModel, status_code=status.HTTP_200_OK, summary="Get specific agent details", description=""" Retrieve detailed information about a specific agent by its ID. This endpoint provides comprehensive status information for a single agent, including real-time metrics, health status, and configuration details. **Returned Information:** - Agent identification and configuration - Current status and utilization - Recent activity and performance metrics - Health check results and connectivity status **Use Cases:** - Monitor specific agent performance - Debug agent connectivity issues - Verify agent configuration - Check agent availability for task assignment """, responses={ 200: {"description": "Agent details retrieved successfully"}, 404: {"model": ErrorResponse, "description": "Agent not found"}, 500: {"model": ErrorResponse, "description": "Internal server error"} } ) async def get_agent( agent_id: str, request: Request, current_user: Dict[str, Any] = Depends(get_current_user_context) ) -> AgentModel: """ Get detailed information about a specific agent. Args: agent_id: Unique identifier of the agent to retrieve request: FastAPI request object current_user: Current authenticated user context Returns: AgentModel: Detailed agent information and status Raises: HTTPException: If agent not found or query fails """ try: with SessionLocal() as db: db_agent = db.query(ORMAgent).filter(ORMAgent.id == agent_id).first() if not db_agent: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"Agent with ID '{agent_id}' not found" ) agent_model = AgentModel( id=db_agent.id, endpoint=db_agent.endpoint, model=db_agent.model, specialty=db_agent.specialty, max_concurrent=db_agent.max_concurrent, current_tasks=db_agent.current_tasks, status="available" if db_agent.current_tasks < db_agent.max_concurrent else "busy", utilization=db_agent.current_tasks / db_agent.max_concurrent if db_agent.max_concurrent > 0 else 0.0 ) return agent_model except HTTPException: raise except Exception as e: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to retrieve agent: {str(e)}" ) @router.delete( "/agents/{agent_id}", status_code=status.HTTP_204_NO_CONTENT, summary="Unregister an agent", description=""" Remove an agent from the WHOOSH cluster. This endpoint safely removes an agent from the cluster by: 1. Checking for active tasks and optionally waiting for completion 2. Removing the agent from the coordinator's active pool 3. Cleaning up database records 4. Confirming successful removal **Safety Measures:** - Active tasks are checked before removal - Graceful shutdown procedures are followed - Database consistency is maintained - Error handling for cleanup failures **Use Cases:** - Remove offline or problematic agents - Scale down cluster capacity - Perform maintenance on agent nodes - Clean up test or temporary agents """, responses={ 204: {"description": "Agent unregistered successfully"}, 404: {"model": ErrorResponse, "description": "Agent not found"}, 409: {"model": ErrorResponse, "description": "Agent has active tasks"}, 500: {"model": ErrorResponse, "description": "Internal server error"} } ) async def unregister_agent( agent_id: str, request: Request, force: bool = False, current_user: Dict[str, Any] = Depends(get_current_user_context) ): """ Unregister an agent from the WHOOSH cluster. Args: agent_id: Unique identifier of the agent to remove request: FastAPI request object force: Whether to force removal even with active tasks current_user: Current authenticated user context Raises: HTTPException: If agent not found, has active tasks, or removal fails """ # Access coordinator whoosh_coordinator = getattr(request.app.state, 'whoosh_coordinator', None) if not whoosh_coordinator: from ..main import unified_coordinator whoosh_coordinator = unified_coordinator if not whoosh_coordinator: raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Coordinator service unavailable" ) try: with SessionLocal() as db: db_agent = db.query(ORMAgent).filter(ORMAgent.id == agent_id).first() if not db_agent: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"Agent with ID '{agent_id}' not found" ) # Check for active tasks unless forced if not force and db_agent.current_tasks > 0: raise HTTPException( status_code=status.HTTP_409_CONFLICT, detail=f"Agent '{agent_id}' has {db_agent.current_tasks} active tasks. Use force=true to override." ) # Remove from coordinator whoosh_coordinator.remove_agent(agent_id) # Remove from database db.delete(db_agent) db.commit() except HTTPException: raise except Exception as e: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to unregister agent: {str(e)}" ) @router.post( "/agents/heartbeat", status_code=status.HTTP_200_OK, summary="Agent heartbeat update", description=""" Update agent status and maintain registration through periodic heartbeat. This endpoint allows agents to: - Confirm they are still online and responsive - Update their current status and metrics - Report any capability or configuration changes - Maintain their registration in the cluster Agents should call this endpoint every 30-60 seconds to maintain their active status in the WHOOSH cluster. """, responses={ 200: {"description": "Heartbeat received successfully"}, 404: {"model": ErrorResponse, "description": "Agent not registered"}, 400: {"model": ErrorResponse, "description": "Invalid heartbeat data"} } ) async def agent_heartbeat( heartbeat_data: Dict[str, Any], request: Request ): """ Process agent heartbeat to maintain registration. Args: heartbeat_data: Agent status and metrics data request: FastAPI request object Returns: Success confirmation and any coordinator updates """ agent_id = heartbeat_data.get("agent_id") if not agent_id: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Missing agent_id in heartbeat data" ) # Access coordinator whoosh_coordinator = getattr(request.app.state, 'whoosh_coordinator', None) if not whoosh_coordinator: from ..main import unified_coordinator whoosh_coordinator = unified_coordinator if not whoosh_coordinator: raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Coordinator service unavailable" ) try: # Update agent heartbeat timestamp agent_service = whoosh_coordinator.agent_service if agent_service: agent_service.update_agent_heartbeat(agent_id) # Update current tasks if provided - use raw SQL to avoid role column if "current_tasks" in heartbeat_data: current_tasks = heartbeat_data["current_tasks"] try: with SessionLocal() as db: from sqlalchemy import text db.execute(text( "UPDATE agents SET current_tasks = :current_tasks, last_seen = NOW() WHERE id = :agent_id" ), { "current_tasks": current_tasks, "agent_id": agent_id }) db.commit() except Exception as e: logger.warning(f"Could not update agent tasks: {e}") return { "status": "success", "message": f"Heartbeat received from agent '{agent_id}'", "timestamp": time.time() } except Exception as e: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to process heartbeat: {str(e)}" ) @router.post( "/agents/auto-register", response_model=AgentRegistrationResponse, status_code=status.HTTP_201_CREATED, summary="Automatic agent registration", description=""" Register an agent automatically with capability detection. This endpoint is designed for Bzzz agents running as systemd services to automatically register themselves with the WHOOSH coordinator. Features: - Automatic capability detection based on available models - Network discovery support - Retry-friendly for service startup scenarios - Health validation before registration """, responses={ 201: {"description": "Agent auto-registered successfully"}, 400: {"model": ErrorResponse, "description": "Invalid agent configuration"}, 409: {"model": ErrorResponse, "description": "Agent already registered"}, 503: {"model": ErrorResponse, "description": "Agent endpoint unreachable"} } ) async def auto_register_agent( agent_data: Dict[str, Any], request: Request ) -> AgentRegistrationResponse: """ Automatically register a Bzzz agent with the WHOOSH coordinator. Args: agent_data: Agent configuration including endpoint, models, etc. request: FastAPI request object Returns: AgentRegistrationResponse: Registration confirmation """ # Extract required fields agent_id = agent_data.get("agent_id") endpoint = agent_data.get("endpoint") hostname = agent_data.get("hostname") if not agent_id or not endpoint: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Missing required fields: agent_id, endpoint" ) # Access coordinator whoosh_coordinator = getattr(request.app.state, 'whoosh_coordinator', None) if not whoosh_coordinator: from ..main import unified_coordinator whoosh_coordinator = unified_coordinator if not whoosh_coordinator: raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Coordinator service unavailable" ) try: # Check if agent already exists - use basic query to avoid role column try: with SessionLocal() as db: from sqlalchemy import text existing_agent = db.execute(text( "SELECT id, endpoint FROM agents WHERE id = :agent_id LIMIT 1" ), {"agent_id": agent_id}).fetchone() if existing_agent: # Update existing agent db.execute(text( "UPDATE agents SET endpoint = :endpoint, last_seen = NOW() WHERE id = :agent_id" ), {"endpoint": endpoint, "agent_id": agent_id}) db.commit() return AgentRegistrationResponse( agent_id=agent_id, endpoint=endpoint, message=f"Agent '{agent_id}' registration updated successfully" ) except Exception as e: logger.warning(f"Could not check existing agent: {e}") # Detect capabilities and models models = agent_data.get("models", []) if not models: # Try to detect models from endpoint try: import aiohttp async with aiohttp.ClientSession() as session: async with session.get(f"{endpoint}/api/tags", timeout=aiohttp.ClientTimeout(total=5)) as response: if response.status == 200: tags_data = await response.json() models = [model["name"] for model in tags_data.get("models", [])] except Exception as e: logger.warning(f"Could not detect models for {agent_id}: {e}") # Determine specialty based on models or hostname specialty = AgentType.GENERAL_AI # Default if "codellama" in str(models).lower() or "code" in hostname.lower(): specialty = AgentType.KERNEL_DEV elif "gemma" in str(models).lower(): specialty = AgentType.PYTORCH_DEV elif any(model for model in models if "llama" in model.lower()): specialty = AgentType.GENERAL_AI # Insert agent directly into database try: with SessionLocal() as db: from sqlalchemy import text # Insert new agent using raw SQL to avoid role column issues db.execute(text(""" INSERT INTO agents (id, name, endpoint, model, specialty, max_concurrent, current_tasks, status, created_at, last_seen) VALUES (:agent_id, :name, :endpoint, :model, :specialty, :max_concurrent, 0, 'active', NOW(), NOW()) ON CONFLICT (id) DO UPDATE SET endpoint = EXCLUDED.endpoint, model = EXCLUDED.model, specialty = EXCLUDED.specialty, max_concurrent = EXCLUDED.max_concurrent, last_seen = NOW() """), { "agent_id": agent_id, "name": agent_id, # Use agent_id as name "endpoint": endpoint, "model": models[0] if models else "unknown", "specialty": specialty.value, "max_concurrent": agent_data.get("max_concurrent", 2) }) db.commit() return AgentRegistrationResponse( agent_id=agent_id, endpoint=endpoint, message=f"Agent '{agent_id}' auto-registered successfully with specialty '{specialty.value}'" ) except Exception as e: logger.error(f"Database insert failed: {e}") raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to register agent in database: {str(e)}" ) except Exception as e: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to auto-register agent: {str(e)}" )