main_v2.py

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import os, random
from openai import AzureOpenAI

app = FastAPI(title="GraTech AI Hub")

# Multi-model load balancing
MODELS = {
    "gpt4": {"endpoint": os.getenv("AZURE_OPENAI_ENDPOINT"), "key": os.getenv("AZURE_OPENAI_API_KEY"), "deployment": "gpt-4o"},
    "deepseek": {"endpoint": os.getenv("AZURE_FOUNDRY_ENDPOINT"), "key": os.getenv("AZURE_FOUNDRY_API_KEY"), "deployment": "DeepSeek-V3.1"},
    "llama": {"endpoint": os.getenv("AZURE_FOUNDRY_ENDPOINT"), "key": os.getenv("AZURE_FOUNDRY_API_KEY"), "deployment": "llama-3-dot1"}
}

class ChatRequest(BaseModel):
    message: str
    model: str = "auto"  # auto, gpt4, deepseek, llama
    
@app.post("/api/chat")
async def chat(req: ChatRequest):
    model_key = req.model if req.model != "auto" else random.choice(list(MODELS.keys()))
    config = MODELS[model_key]
    
    client = AzureOpenAI(azure_endpoint=config["endpoint"], api_key=config["key"], api_version="2024-06-01")
    
    response = client.chat.completions.create(
        model=config["deployment"],
        messages=[{"role": "user", "content": req.message}],
        temperature=0.3
    )
    
    return {"reply": response.choices[0].message.content, "model_used": model_key}

@app.get("/api/health")
async def health():
    return {"status": "healthy", "models": list(MODELS.keys())}