# main.py from fastapi import FastAPI, HTTPException, Query from pydantic import BaseModel from typing import Literal from dotenv import load_dotenv # Import our new factory function from app.llm_providers import get_llm_provider # --- 1. Application Setup --- load_dotenv() app = FastAPI( title="AI Model Hub Service", description="A extensible hub to route requests to various LLMs using a Factory Pattern.", version="0.0.0", ) # --- 2. Pydantic Models --- class ChatRequest(BaseModel): prompt: str # --- 3. API Endpoints --- @app.get("/") def read_root(): return {"status": "AI Model Hub is running!"} @app.post("/chat") async def chat_handler( request: ChatRequest, model: Literal["deepseek", "gemini"] = Query("deepseek", description="The AI model to use.") ): try: # Use the factory to get the correct provider instance provider = get_llm_provider(model) # Call the method on the instance. We don't need to know if it's # Gemini or DeepSeek, only that it fulfills the "contract". response_text = await provider.generate_response(request.prompt) return {"response": response_text, "model_used": model} except ValueError as e: # This catches errors from the factory (e.g., unsupported model) raise HTTPException(status_code=400, detail=str(e)) except Exception as e: # This catches errors from the provider's API call raise HTTPException(status_code=500, detail=f"An error occurred with the {model} API: {e}")