Update app.py

app.py

@@ -1,25 +1,29 @@
 import os
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Query
 from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
 from openai import AsyncOpenAI
+from typing import Optional
 
 app = FastAPI()
 
-async def generate_ai_response(prompt: str , model: str):
-    # Configuration for unofficial GitHub AI endpoint
+class GenerateRequest(BaseModel):
+    prompt: str
+
+async def generate_ai_response(prompt: str, model: str):
+    # Configuration for AI endpoint
     token = os.getenv("GITHUB_TOKEN")
+    endpoint = os.getenv("AI_SERVER_URL", "https://models.github.ai/inference")  # Default fallback
+    
     if not token:
         raise HTTPException(status_code=500, detail="GitHub token not configured")
-    
-    endpoint = "https://models.github.ai/inference"
-    
 
     client = AsyncOpenAI(base_url=endpoint, api_key=token)
 
-    try:
+    try
         stream = await client.chat.completions.create(
             messages=[
-                {"role": "system", "content": "You are a helpful assistant named Orion and made by Abdullah Ali"},
+                {"role": "system", "content": "You are a helpful assistant named Orion, created by Abdullah Ali"},
                 {"role": "user", "content": prompt}
             ],
             model=model,
@@ -34,17 +38,34 @@ async def generate_ai_response(prompt: str , model: str):
 
     except Exception as err:
         yield f"Error: {str(err)}"
-        raise HTTPException(status_code=500, detail="AI generation failed")
+        raise HTTPException(status_code=500, detail=f"AI generation failed: {str(err)}")
 
-@app.post("/generate")
-async def generate_response(prompt: str , model: str):
-    if not prompt:
+@app.post("/generate", summary="Generate AI response", response_description="Streaming AI response")
+async def generate_response(
+    model: str = Query("default-model", description="The AI model to use"),
+    prompt: Optional[str] = Query(None, description="The input text prompt for the AI"),
+    request: Optional[GenerateRequest] = None
+):
+    """
+    Generate a streaming AI response based on the provided prompt and model.
+    
+    - **model**: The AI model to use (specified as a query parameter, defaults to default-model)
+    - **prompt**: The input text prompt for the AI (can be in query parameter or request body)
+    """
+    # Determine prompt source: query parameter or request body
+    final_prompt = prompt if prompt is not None else (request.prompt if request is not None else None)
+    
+    if not final_prompt or not final_prompt.strip():
         raise HTTPException(status_code=400, detail="Prompt cannot be empty")
     
+    if not model or not model.strip():
+        raise HTTPException(status_code=400, detail="Model cannot be empty")
+    
     return StreamingResponse(
-        generate_ai_response(prompt , model),
+        generate_ai_response(final_prompt, model),
         media_type="text/event-stream"
     )
 
 def get_app():
-    return app
+    return app
+