Newer
Older
cortex-hub / ai-hub / app / main.py
# 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}")