from fastapi import FastAPI, HTTPException, Header from pydantic import BaseModel from dataclasses import dataclass from typing import Optional from enum import Enum from dotenv import load_dotenv from memory import TokenBufferMemory from agent_factory import AgentFactory from base_agent import ManagerAgent from tools import initialize_registry from llm_provider import GroqProvider, GeminiProvider, QuotaExhaustedError, ProviderDownError, ProviderError import time import uvicorn import logging import os import sys load_dotenv() logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", handlers=[ logging.FileHandler("agent_debug.log"), logging.StreamHandler(sys.stdout) ] ) logger = logging.getLogger("SystemBackend") class ChatRequest(BaseModel): query: str provider: Optional[str] = None api_key: Optional[str] = None class ChatResponse(BaseModel): success: bool response: Optional[str] = None provider_used: Optional[str] = None agent_used: Optional[str] = None error_type: Optional[str] = None required_provider: Optional[str] = None message: Optional[str] = None app = FastAPI() registry = initialize_registry() factory = AgentFactory(registry) BASE_DIR = os.path.dirname(os.path.abspath(__file__)) YAML_PATH = os.path.join(BASE_DIR, "agents.yaml") try: workers_map = factory.load_from_yaml(YAML_PATH) logger.info(f"✅ Successfully loaded agents: {list(workers_map.keys())}") except FileNotFoundError: logger.critical("❌ 'agents.yaml' not found! Please create this configuration file.") sys.exit(1) agent_memory = TokenBufferMemory(max_tokens=4096) manager_agent = ManagerAgent( name="Manager", sub_agents=workers_map, memory=agent_memory ) class ProviderStatus(Enum): ACTIVE = "Active" DOWN = "Down" QUOTA_EXHAUSTED = "QuotaExhausted" ERROR = "Error" @dataclass class ProviderState: name: str status: ProviderStatus reset_time: float class ProviderManager: def __init__(self): self.providers = { "groq": ProviderState(name="groq", status=ProviderStatus.ACTIVE, reset_time=0.0), "gemini": ProviderState(name="gemini", status=ProviderStatus.ACTIVE, reset_time=0.0), } def update_status(self, provider_name: str, status: ProviderStatus): self.providers[provider_name].status = status if status == ProviderStatus.QUOTA_EXHAUSTED: self.providers[provider_name].reset_time = time.time() + (24 * 60 * 60) # 24 hours elif status == ProviderStatus.DOWN: self.providers[provider_name].reset_time = time.time() + 60 # 60 seconds def get_provider(self): for name, state in self.providers.items(): if state.status == ProviderStatus.ACTIVE: return name elif state.status in [ProviderStatus.DOWN, ProviderStatus.QUOTA_EXHAUSTED]: if time.time() > state.reset_time: state.status = ProviderStatus.ACTIVE return name return None def is_provider_active(self, provider_name: str) -> bool: if provider_name not in self.providers: return False state = self.providers[provider_name] if state.status in [ProviderStatus.DOWN, ProviderStatus.QUOTA_EXHAUSTED]: if time.time() > state.reset_time: state.status = ProviderStatus.ACTIVE return True return False return state.status == ProviderStatus.ACTIVE provider_manager = ProviderManager() class RequestBody(BaseModel): query: str provider: Optional[str] = None api_key: Optional[str] = None def has_server_key(name: str) -> bool: if name == "groq" and os.getenv("GROQ_API_KEY"): return True if name == "gemini" and os.getenv("GEMINI_API_KEY"): return True return False def get_provider_instance(name: str, key: str): if name == "groq": return GroqProvider(api_key=key) elif name == "gemini": return GeminiProvider(api_key=key) raise ValueError(f"Unknown provider: {name}") @app.post("/chat", response_model=ChatResponse) async def chat_endpoint(request: ChatRequest): # CASE 1: MANUAL OVERRIDE (User specifically asked for a provider) if request.provider: target = request.provider.lower() # A. Check if valid/active if not provider_manager.is_provider_active(target): return ChatResponse( success=False, error_type="provider_down", required_provider=target, message=f"Requested provider '{target}' is currently unavailable (Down/Quota)." ) # B. Resolve Key final_key = request.api_key if request.api_key else None if not final_key and has_server_key(target): final_key = os.getenv(f"{target.upper()}_API_KEY") logger.info(f"🔍 DEBUG CHECK: Target={target}, Key_Type={type(final_key)}, Has_Key={bool(final_key)}") if not final_key: return ChatResponse( success=False, error_type="needs_key", required_provider=target, message=f"API Key missing for {target}." ) # C. Execute (NO LOOP - Fail fast if user preference fails) try: logger.info(f"🔄 Executing User Preference: {target}") llm = get_provider_instance(target, final_key) result = manager_agent.process_query(request.query, llm) return ChatResponse( success=True, response=result.content, provider_used=target, agent_used=manager_agent.name ) except (QuotaExhaustedError, ProviderDownError) as e: provider_manager.update_status(target, ProviderStatus.QUOTA_EXHAUSTED) # or check exception type return ChatResponse(success=False, error_type="provider_down", message=str(e)) except Exception as e: logger.error(f"Server Error: {e}") return ChatResponse(success=False, error_type="server_error", message=str(e)) # CASE 2: AUTO-PILOT (Loop through available providers) else: while True: current = provider_manager.get_provider() if not current: return ChatResponse( success=False, error_type="all_down", message="All providers are currently down or exhausted." ) # Check Key for the Auto-Selected candidate final_key = None if has_server_key(current): final_key = os.getenv(f"{current.upper()}_API_KEY") if not final_key: # If Auto-Router picks a provider we have no key for, we must ask the user. return ChatResponse( success=False, error_type="needs_key", required_provider=current, message=f"Auto-switching to {current}, but API Key is missing." ) try: logger.info(f"🔄 Auto-Routing via: {current}") llm = get_provider_instance(current, final_key) result = manager_agent.process_query(request.query, llm) return ChatResponse( success=True, response=result.content, provider_used=current, agent_used=manager_agent.name ) except QuotaExhaustedError: provider_manager.update_status(current, ProviderStatus.QUOTA_EXHAUSTED) continue # Try next in loop except ProviderDownError: provider_manager.update_status(current, ProviderStatus.DOWN) continue # Try next in loop except Exception as e: logger.error(f"Critical Error: {e}") return ChatResponse(success=False, error_type="server_error", message=str(e)) if __name__ == "__main__": port = int(os.getenv("PORT", 7860)) uvicorn.run(app, host="0.0.0.0", port=port)