Update app.py

app.py

@@ -1,18 +1,28 @@
 import os
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Query
 from fastapi.responses import StreamingResponse
 from openai import AsyncOpenAI
 
 app = FastAPI()
 
-async def generate_ai_response(prompt: str):
+# Define available models (you can expand this list)
+AVAILABLE_MODELS = {
+    "openai/gpt-4.1-mini": "OpenAI GPT-4 Mini",
+    "deepseek/DeepSeek-V3-0324": "DeepSeek V3 0324",
+    # Add more models as needed
+}
+
+async def generate_ai_response(prompt: str, model: str):
     # Configuration for unofficial GitHub AI endpoint
     token = os.getenv("GITHUB_TOKEN")
     if not token:
         raise HTTPException(status_code=500, detail="GitHub token not configured")
     
     endpoint = "https://models.github.ai/inference"
-    model = "openai/gpt-4.1-mini"  # Unofficial model name
+    
+    # Validate the model
+    if model not in AVAILABLE_MODELS:
+        raise HTTPException(status_code=400, detail=f"Model not available. Choose from: {', '.join(AVAILABLE_MODELS.keys())}")
 
     client = AsyncOpenAI(base_url=endpoint, api_key=token)
 
@@ -37,12 +47,15 @@ async def generate_ai_response(prompt: str):
         raise HTTPException(status_code=500, detail="AI generation failed")
 
 @app.post("/generate")
-async def generate_response(prompt: str):
+async def generate_response(
+    prompt: str = Query(..., description="The prompt for the AI"),
+    model: str = Query("openai/gpt-4.1-mini", description="The model to use for generation")
+):
     if not prompt:
         raise HTTPException(status_code=400, detail="Prompt cannot be empty")
     
     return StreamingResponse(
-        generate_ai_response(prompt),
+        generate_ai_response(prompt, model),
         media_type="text/event-stream"
     )