api_server.py

#!/usr/bin/env python3
"""
FastAPI REST API Server for the Agentic AI System.

This creates a production-ready REST API that exposes the agentic AI system
to the real world with proper authentication, rate limiting, and documentation.
"""

import os
import sys
import time
import asyncio
from pathlib import Path
from typing import Dict, Any, List, Optional
from datetime import datetime, timedelta
import uuid

# Add src to path
sys.path.insert(0, str(Path(__file__).parent / "src"))

from fastapi import FastAPI, HTTPException, Depends, Security, status, BackgroundTasks
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
import uvicorn

from config.settings import settings
from src.utils.logging_config import logger
from src.agents.supervisor_agent import SupervisorAgent
from src.agents.worker_agents import QueryParserAgent, APIMatcherAgent, APIExecutorAgent, ResultFormatterAgent
from src.evaluators.reflection_evaluator import ReflectionEvaluator
from src.evaluators.quality_assessor import QualityAssessor
from src.evaluators.memory_manager import MemoryManager
from src.api_clients.api_matcher import APIOperation
from src.parsers.keyword_extractor import KeywordExtractor


# Pydantic models for API requests/responses
class QueryRequest(BaseModel):
    """Request model for query processing."""
    query: str = Field(..., description="Natural language query to process", min_length=1, max_length=1000)
    include_debug: bool = Field(False, description="Include debug information in response")
    max_api_calls: int = Field(5, description="Maximum number of API calls to execute", ge=1, le=10)
    metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata")


class QueryResponse(BaseModel):
    """Response model for query processing."""
    request_id: str = Field(..., description="Unique request identifier")
    query: str = Field(..., description="Original query")
    success: bool = Field(..., description="Whether processing was successful")
    processing_time: float = Field(..., description="Processing time in seconds")
    results: Dict[str, Any] = Field(..., description="Processing results")
    error_message: Optional[str] = Field(None, description="Error message if failed")
    timestamp: datetime = Field(..., description="Response timestamp")


class APIOperationModel(BaseModel):
    """Model for API operation registration."""
    name: str = Field(..., description="Operation name")
    method: str = Field(..., description="HTTP method")
    endpoint: str = Field(..., description="API endpoint")
    description: str = Field(..., description="Operation description")
    tags: List[str] = Field(default_factory=list, description="Operation tags")
    parameters: List[Dict[str, Any]] = Field(default_factory=list, description="Operation parameters")


class SystemStatus(BaseModel):
    """System status response."""
    status: str = Field(..., description="System status")
    version: str = Field(..., description="System version")
    uptime: float = Field(..., description="System uptime in seconds")
    total_requests: int = Field(..., description="Total requests processed")
    active_agents: int = Field(..., description="Number of active agents")
    last_request: Optional[datetime] = Field(None, description="Last request timestamp")


# Global variables for system state
app_start_time = time.time()
request_count = 0
last_request_time = None
supervisor_agent = None
quality_assessor = None
memory_manager = None
api_operations_registry = []

# Security
security = HTTPBearer()

# Load API keys from configuration file.
# SECURITY: no hardcoded fallback tokens. If config load fails, the key manager
# (config/api_keys.py) generates a random key or reads ADMIN_API_KEY from the env.
from config.api_keys import api_key_manager
API_KEYS = api_key_manager.get_api_keys()
logger.info(f"Loaded {len(API_KEYS)} API key(s) from configuration")


def verify_api_key(credentials: HTTPAuthorizationCredentials = Security(security)) -> Dict[str, Any]:
    """Verify API key authentication."""
    token = credentials.credentials
    
    # Check with API key manager first
    try:
        from config.api_keys import api_key_manager
        if api_key_manager.validate_key(token):
            return api_key_manager.get_key_info(token)
    except Exception:
        pass  # Fall back to static keys
    
    # Fallback to static API_KEYS
    if token not in API_KEYS:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid API key",
            headers={"WWW-Authenticate": "Bearer"},
        )
    return API_KEYS[token]


# Create FastAPI app.
# SECURITY: interactive docs are disabled by default in production; enable them
# explicitly with EXPOSE_DOCS=1 for local development only.
_expose_docs = os.getenv("EXPOSE_DOCS", "0") == "1"
app = FastAPI(
    title="Agentic AI System API",
    description="Production-ready REST API for natural language query processing and API orchestration",
    version="1.0.0",
    docs_url="/docs" if _expose_docs else None,
    redoc_url="/redoc" if _expose_docs else None,
    openapi_url="/openapi.json" if _expose_docs else None,
)

# Add CORS middleware.
# SECURITY: wildcard "*" + allow_credentials is invalid and unsafe. Origins are
# read from ALLOWED_ORIGINS (comma-separated); default is no cross-origin access.
_allowed_origins = [
    o.strip() for o in os.getenv("ALLOWED_ORIGINS", "").split(",") if o.strip()
]
app.add_middleware(
    CORSMiddleware,
    allow_origins=_allowed_origins,
    allow_credentials=True,
    allow_methods=["GET", "POST"],
    allow_headers=["Authorization", "Content-Type"],
)


@app.on_event("startup")
async def startup_event():
    """Initialize the system on startup."""
    global supervisor_agent, quality_assessor, memory_manager, api_operations_registry
    
    logger.info("Starting Agentic AI System API Server...")
    
    try:
        # Initialize evaluation components
        quality_assessor = QualityAssessor()
        memory_manager = MemoryManager()
        
        # Initialize supervisor agent
        supervisor_agent = SupervisorAgent(config={
            'max_retries': 3,
            'retry_delay': 1.0
        })
        
        # Initialize worker agents
        query_parser = QueryParserAgent()
        api_matcher = APIMatcherAgent()
        api_executor = APIExecutorAgent(config={
            'cdms': {'enabled': False}  # Configure as needed
        })
        result_formatter = ResultFormatterAgent()
        
        # Initialize reflection evaluator
        reflection_evaluator = ReflectionEvaluator(config={
            'quality_assessor': quality_assessor,
            'memory_manager': memory_manager
        })
        
        # Register worker agents
        supervisor_agent.register_worker('query_parser', query_parser)
        supervisor_agent.register_worker('api_matcher', api_matcher)
        supervisor_agent.register_worker('api_executor', api_executor)
        supervisor_agent.register_worker('result_formatter', result_formatter)
        supervisor_agent.register_worker('evaluator', reflection_evaluator)
        
        # Load sample API operations
        sample_operations = [
            APIOperation(
                name="getUserProfile",
                method="GET",
                endpoint="/users/{userId}",
                description="Retrieve user profile information",
                parameters=[{"name": "userId", "type": "string", "required": True}],
                tags=["users", "profile"]
            ),
            APIOperation(
                name="searchDatasets",
                method="GET",
                endpoint="/datasets/search",
                description="Search for machine learning datasets",
                parameters=[{"name": "query", "type": "string"}, {"name": "limit", "type": "integer"}],
                tags=["datasets", "search", "ml"]
            ),
            APIOperation(
                name="getCDMSLabels",
                method="GET",
                endpoint="/cdms/labels",
                description="Retrieve CDMS labels for datasets",
                parameters=[{"name": "dataset_id", "type": "string"}],
                tags=["cdms", "labels", "metadata"]
            )
        ]
        
        # Add operations to matcher
        api_matcher.api_matcher.add_operations(sample_operations)
        api_operations_registry = sample_operations
        
        logger.info("Agentic AI System API Server started successfully")
        logger.info(f"Registered {len(api_operations_registry)} API operations")
        
    except Exception as e:
        logger.error(f"Failed to start API server: {e}")
        raise


@app.get("/", response_model=Dict[str, str])
async def root():
    """Root endpoint with basic information."""
    return {
        "message": "Agentic AI System API",
        "version": "1.0.0",
        "status": "operational",
        "docs": "/docs",
        "health": "/health"
    }


@app.get("/health", response_model=SystemStatus)
async def health_check():
    """Health check endpoint."""
    global request_count, last_request_time
    
    uptime = time.time() - app_start_time
    active_agents = len(supervisor_agent.worker_agents) if supervisor_agent else 0
    
    return SystemStatus(
        status="healthy",
        version="1.0.0",
        uptime=uptime,
        total_requests=request_count,
        active_agents=active_agents,
        last_request=last_request_time
    )


@app.post("/query", response_model=QueryResponse)
async def process_query(
    request: QueryRequest,
    background_tasks: BackgroundTasks,
    user: Dict[str, Any] = Depends(verify_api_key)
):
    """
    Process a natural language query through the agentic AI system.
    
    This endpoint:
    1. Parses the natural language query
    2. Matches it to available APIs
    3. Executes relevant API calls
    4. Returns structured results
    """
    global request_count, last_request_time
    
    request_id = str(uuid.uuid4())
    start_time = time.time()
    request_count += 1
    last_request_time = datetime.utcnow()
    
    logger.info(f"Processing query request {request_id}: {request.query}")
    
    try:
        if not supervisor_agent:
            raise HTTPException(
                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
                detail="System not initialized"
            )
        
        # Execute the agentic workflow
        result = supervisor_agent.execute({
            'query': request.query,
            'metadata': {
                'request_id': request_id,
                'user': user['name'],
                'max_api_calls': request.max_api_calls,
                'include_debug': request.include_debug,
                **request.metadata
            }
        })
        
        processing_time = time.time() - start_time
        
        # Format response
        if result.success:
            response_data = result.data.get('formatted_output', result.data)
            
            # Add debug information if requested
            if request.include_debug:
                response_data['debug'] = {
                    'execution_time': result.execution_time,
                    'agent_metadata': result.metadata,
                    'workflow_steps': ['parse_query', 'match_apis', 'execute_apis', 'format_results']
                }
            
            return QueryResponse(
                request_id=request_id,
                query=request.query,
                success=True,
                processing_time=processing_time,
                results=response_data,
                timestamp=datetime.utcnow()
            )
        else:
            return QueryResponse(
                request_id=request_id,
                query=request.query,
                success=False,
                processing_time=processing_time,
                results={},
                error_message=result.error_message,
                timestamp=datetime.utcnow()
            )
            
    except Exception as e:
        processing_time = time.time() - start_time
        # SECURITY: log full detail server-side, return a generic message to clients.
        logger.exception(f"Error processing query {request_id}: {e}")

        return QueryResponse(
            request_id=request_id,
            query=request.query,
            success=False,
            processing_time=processing_time,
            results={},
            error_message="Internal error while processing the query.",
            timestamp=datetime.utcnow()
        )


@app.get("/apis", response_model=List[Dict[str, Any]])
async def list_api_operations(user: Dict[str, Any] = Depends(verify_api_key)):
    """List all registered API operations."""
    operations = []
    for op in api_operations_registry:
        operations.append({
            'name': op.name,
            'method': op.method,
            'endpoint': op.endpoint,
            'description': op.description,
            'tags': op.tags,
            'parameters': op.parameters
        })
    return operations


@app.post("/apis", response_model=Dict[str, str])
async def register_api_operation(
    operation: APIOperationModel,
    user: Dict[str, Any] = Depends(verify_api_key)
):
    """Register a new API operation."""
    if "write" not in user.get("permissions", []):
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Insufficient permissions"
        )
    
    try:
        # Create API operation
        api_op = APIOperation(
            name=operation.name,
            method=operation.method,
            endpoint=operation.endpoint,
            description=operation.description,
            parameters=operation.parameters,
            tags=operation.tags
        )
        
        # Add to registry and matcher
        api_operations_registry.append(api_op)
        if supervisor_agent:
            api_matcher = supervisor_agent.worker_agents.get('api_matcher')
            if api_matcher:
                api_matcher.api_matcher.add_operation(api_op)
        
        logger.info(f"Registered new API operation: {operation.name}")
        
        return {
            "message": f"API operation '{operation.name}' registered successfully",
            "operation_id": operation.name
        }
        
    except Exception as e:
        logger.exception(f"Error registering API operation: {e}")
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Failed to register API operation."
        )


@app.get("/keywords/{query}")
async def extract_keywords(
    query: str,
    user: Dict[str, Any] = Depends(verify_api_key)
):
    """Extract keywords from a query (utility endpoint)."""
    try:
        extractor = KeywordExtractor()
        keywords = extractor.extract_keywords_hybrid(query)
        
        return {
            "query": query,
            "keywords": keywords,
            "count": len(keywords),
            "timestamp": datetime.utcnow()
        }
        
    except Exception as e:
        logger.exception(f"Error extracting keywords: {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Keyword extraction failed."
        )


@app.get("/stats", response_model=Dict[str, Any])
async def get_system_stats(user: Dict[str, Any] = Depends(verify_api_key)):
    """Get detailed system statistics."""
    if "admin" not in user.get("permissions", []):
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Admin permissions required"
        )
    
    stats = {
        "system": {
            "uptime": time.time() - app_start_time,
            "total_requests": request_count,
            "last_request": last_request_time,
            "registered_apis": len(api_operations_registry)
        },
        "agents": {},
        "configuration": {
            "max_retries": 3,
            "timeout": settings.api_timeout,
            "debug_mode": settings.debug
        }
    }
    
    # Add agent statistics
    if supervisor_agent:
        for agent_type, agent in supervisor_agent.worker_agents.items():
            agent_status = agent.get_status()
            stats["agents"][agent_type] = {
                "state": agent_status["state"],
                "execution_count": agent_status["execution_count"],
                "capabilities": agent.get_capabilities()
            }
    
    return stats


# Error handlers
@app.exception_handler(HTTPException)
async def http_exception_handler(request, exc):
    """Custom HTTP exception handler."""
    return JSONResponse(
        status_code=exc.status_code,
        content={
            "error": True,
            "message": exc.detail,
            "status_code": exc.status_code,
            "timestamp": datetime.utcnow().isoformat()
        }
    )


@app.exception_handler(Exception)
async def general_exception_handler(request, exc):
    """General exception handler."""
    logger.error(f"Unhandled exception: {exc}")
    return JSONResponse(
        status_code=500,
        content={
            "error": True,
            "message": "Internal server error",
            "status_code": 500,
            "timestamp": datetime.utcnow().isoformat()
        }
    )


if __name__ == "__main__":
    # Run the server. reload/host are opt-in via env so production defaults are safe.
    uvicorn.run(
        "api_server:app",
        host=os.getenv("API_HOST", "127.0.0.1"),
        port=int(os.getenv("API_PORT", "8000")),
        reload=os.getenv("API_RELOAD", "0") == "1",
        log_level=os.getenv("API_LOG_LEVEL", "info"),
    )