ComputeAgent/main.py

# This allows importing modules from the top-level project directory
import os
import sys
sys.path.append("/home/hivenet")
"""
ComputeAgent FastAPI Main Application

This is the main entry point for the ComputeAgent FastAPI application.
It creates the FastAPI app, includes routers, and manages the application lifecycle.

Features:
- FastAPI application setup
- Router inclusion for modular organization
- Application lifecycle management (startup/shutdown)
- CORS middleware configuration
- Global error handlers
- Background task management for memory operations
- Interactive API documentation

Usage:
    python main.py

Or with uvicorn directly:
    uvicorn main:app --host 0.0.0.0 --port 8000 --reload

Author: ComputeAgent Team
License: Private
"""

import asyncio
import logging
from contextlib import asynccontextmanager

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, RedirectResponse
from fastapi.middleware.cors import CORSMiddleware
import uvicorn

# Import the compute agent router and initialization function
from ComputeAgent.routers.compute_agent_HITL import compute_agent_router, initialize_agent

# Initialize logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("ComputeAgent Main")


@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    Lifespan context manager for FastAPI application.
    Handles startup and shutdown events.
    
    Args:
        app: FastAPI application instance
    """
    # Startup
    logger.info("=" * 80)
    logger.info("🚀 Starting ComputeAgent API Application...")
    logger.info("=" * 80)
    
    try:
        # Initialize the ComputeAgent
        await initialize_agent()
        logger.info("✅ ComputeAgent API ready to serve requests")
        
    except Exception as e:
        logger.error(f"❌ Failed to initialize application: {e}")
        raise
    
    logger.info("=" * 80)
    logger.info("📚 API Documentation available at:")
    logger.info("   - Swagger UI: http://localhost:8000/docs")
    logger.info("   - ReDoc: http://localhost:8000/redoc")
    logger.info("=" * 80)
    
    yield
    
    # Shutdown
    logger.info("=" * 80)
    logger.info("👋 Shutting down ComputeAgent API Application...")
    
    logger.info("✅ ComputeAgent API shutdown complete")
    logger.info("=" * 80)


# Create FastAPI application
app = FastAPI(
    title="ComputeAgent API",
    description="""
    AI-powered agent system for model deployment and compute workflows.
    
    ## Features
    
    * **Model Deployment**: Deploy AI models from HuggingFace with capacity estimation
    * **React Agent**: Execute compute tasks with MCP tool integration
    * **Memory Management**: Persistent conversations across sessions
    * **Streaming Support**: Real-time updates via Server-Sent Events
    * **Human-in-the-Loop**: Approval workflow for capacity decisions
    
    ## Endpoints
    
    ### ComputeAgent
    - **POST /api/compute/query** - Process queries (non-streaming)
    - **POST /api/compute/query/stream** - Process queries (streaming)
    - **POST /api/compute/memory/clear** - Clear conversation memory
    - **POST /api/compute/memory/inspect** - Inspect memory status
    - **GET /api/compute/health** - Health check
    - **GET /api/compute/examples** - Example queries
    - **GET /api/compute/info** - Router information
    
    ## Getting Started
    
    1. Check API health: `GET /api/compute/health`
    2. Get example queries: `GET /api/compute/examples`
    3. Process a query: `POST /api/compute/query`
    
    For streaming responses, use: `POST /api/compute/query/stream`
    """,
    version="1.0.0",
    lifespan=lifespan,
    docs_url="/docs",
    redoc_url="/redoc"
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Configure appropriately for production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Include routers
app.include_router(compute_agent_router)

# Note: Root endpoint ("/") is handled by Gradio when mounted in app.py
# Removed redirect to /docs to allow Gradio to be at root path

@app.get("/health", tags=["root"])
async def global_health_check():
    """
    Global health check endpoint.
    
    Returns:
        Application health status
    """
    return {
        "status": "healthy",
        "application": "ComputeAgent API",
        "version": "1.0.0",
        "docs": "/docs",
        "compute_agent_health": "/api/compute/health"
    }


# Global error handlers
@app.exception_handler(404)
async def not_found_handler(request: Request, exc: Exception):
    """
    Custom 404 handler for not found endpoints.
    
    Args:
        request: The incoming request
        exc: The exception raised
        
    Returns:
        JSON response with error details
    """
    return JSONResponse(
        status_code=404,
        content={
            "success": False,
            "error": "Endpoint not found",
            "path": str(request.url.path),
            "message": "The requested endpoint does not exist. Visit /docs for available endpoints."
        }
    )


@app.exception_handler(500)
async def internal_error_handler(request: Request, exc: Exception):
    """
    Custom 500 handler for internal server errors.
    
    Args:
        request: The incoming request
        exc: The exception raised
        
    Returns:
        JSON response with error details
    """
    logger.error(f"Internal server error on {request.url.path}: {exc}")
    return JSONResponse(
        status_code=500,
        content={
            "success": False,
            "error": "Internal server error",
            "detail": str(exc),
            "message": "An unexpected error occurred. Please try again or contact support."
        }
    )


@app.exception_handler(Exception)
async def general_exception_handler(request: Request, exc: Exception):
    """
    General exception handler for uncaught exceptions.
    
    Args:
        request: The incoming request
        exc: The exception raised
        
    Returns:
        JSON response with error details
    """
    logger.error(f"Unhandled exception on {request.url.path}: {exc}", exc_info=True)
    return JSONResponse(
        status_code=500,
        content={
            "success": False,
            "error": "Unexpected error",
            "detail": str(exc),
            "message": "An unexpected error occurred. Please check logs for details."
        }
    )


# Middleware for logging
@app.middleware("http")
async def log_requests(request: Request, call_next):
    """
    Middleware to log all incoming requests.
    
    Args:
        request: The incoming request
        call_next: The next middleware or route handler
        
    Returns:
        Response from the route handler
    """
    logger.info(f"📨 {request.method} {request.url.path}")
    
    try:
        response = await call_next(request)
        logger.info(f"✅ {request.method} {request.url.path} - Status: {response.status_code}")
        return response
    except Exception as e:
        logger.error(f"❌ {request.method} {request.url.path} - Error: {e}")
        raise