hive/backend/app/api/agents.py

"""
Hive API - Agent Management Endpoints

This module provides comprehensive API endpoints for managing Ollama-based AI agents
in the Hive 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
from ..core.unified_coordinator import Agent, AgentType
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


@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 Hive 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 Hive 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 Hive 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 Hive 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
    hive_coordinator = getattr(request.app.state, 'hive_coordinator', None)
    if not hive_coordinator:
        # Fallback to global coordinator if app state not available
        from ..main import unified_coordinator
        hive_coordinator = unified_coordinator
        
    if not hive_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
        hive_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 Hive 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 Hive 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
    hive_coordinator = getattr(request.app.state, 'hive_coordinator', None)
    if not hive_coordinator:
        from ..main import unified_coordinator
        hive_coordinator = unified_coordinator
        
    if not hive_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
            hive_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)}"
        )