Update app.py

app.py

@@ -1,13 +1,13 @@
 import os
-import httpx
+import asyncio
 from fastapi import FastAPI, HTTPException, Query
 from fastapi.responses import StreamingResponse
+from openai import AsyncOpenAI
 from collections import defaultdict
-from typing import AsyncGenerator
 
 app = FastAPI()
 
-# Model list (unchanged)
+# Define available models
 AVAILABLE_MODELS = {
     "openai/gpt-4.1": "OpenAI GPT-4.1",
     "openai/gpt-4.1-mini": "OpenAI GPT-4.1-mini",
@@ -56,70 +56,79 @@ AVAILABLE_MODELS = {
     "mistral-ai/mistral-small-2503": "Mistral Small 3.1"
 }
 
-# In-memory history
+# Chat memory (in-memory)
 chat_histories = defaultdict(list)
+MAX_HISTORY = 10  # limit memory to avoid crashes
 
-# Async generator for AI response
-async def generate_ai_response(chat_id: str, model: str) -> AsyncGenerator[str, None]:
+# Generate response stream
+async def generate_ai_response(chat_id: str, model: str):
     token = os.getenv("GITHUB_TOKEN")
     if not token:
         raise HTTPException(status_code=500, detail="GitHub token not configured")
 
+    endpoint = "https://models.github.ai/inference"
+
     if model not in AVAILABLE_MODELS:
-        raise HTTPException(status_code=400, detail=f"Invalid model. Choose from: {', '.join(AVAILABLE_MODELS.keys())}")
+        raise HTTPException(
+            status_code=400,
+            detail=f"Model not available. Choose from: {', '.join(AVAILABLE_MODELS.keys())}"
+        )
+
+    client = AsyncOpenAI(base_url=endpoint, api_key=token)
 
-    headers = {
-        "Authorization": f"Bearer {token}",
-        "Content-Type": "application/json"
-    }
+    try:
+        stream = await asyncio.wait_for(
+            client.chat.completions.create(
+                messages=chat_histories[chat_id],
+                model=model,
+                temperature=1.0,
+                top_p=1.0,
+                stream=True
+            ),
+            timeout=60  # Prevent hangs
+        )
 
-    payload = {
-        "model": model,
-        "messages": chat_histories[chat_id],
-        "stream": True,
-        "temperature": 1.0,
-        "top_p": 1.0
-    }
+        async for chunk in stream:
+            if chunk.choices and chunk.choices[0].delta.content:
+                content = chunk.choices[0].delta.content
+                yield content
+                chat_histories[chat_id].append({"role": "assistant", "content": content})
+                chat_histories[chat_id] = chat_histories[chat_id][-MAX_HISTORY:]
 
-    async with httpx.AsyncClient(timeout=60.0) as client:
-        try:
-            async with client.stream("POST", "https://models.github.ai/inference", headers=headers, json=payload) as response:
-                async for line in response.aiter_lines():
-                    if line.startswith("data:"):
-                        data = line[len("data:"):].strip()
-                        if data == "[DONE]":
-                            break
-                        if data:
-                            yield f"{data}\n"
-                            # Optionally: append to chat history
-                            chat_histories[chat_id].append({"role": "assistant", "content": data})
-        except Exception as e:
-            yield f"Error: {str(e)}"
+    except asyncio.TimeoutError:
+        yield "Error: Response timed out."
+        raise HTTPException(status_code=504, detail="Model timed out.")
+    except Exception as err:
+        yield f"Error: {str(err)}"
+        raise HTTPException(status_code=500, detail="AI generation failed")
 
-# Generate response endpoint
+# Chat endpoint
 @app.post("/generate")
 async def generate_response(
-    chat_id: str = Query(..., description="Chat session ID"),
-    prompt: str = Query(..., description="User input message"),
+    chat_id: str = Query(..., description="Unique chat ID"),
+    prompt: str = Query(..., description="User message"),
     model: str = Query("openai/gpt-4.1-mini", description="Model to use")
 ):
     if not prompt:
         raise HTTPException(status_code=400, detail="Prompt cannot be empty")
-    
+
     chat_histories[chat_id].append({"role": "user", "content": prompt})
+    chat_histories[chat_id] = chat_histories[chat_id][-MAX_HISTORY:]
 
     return StreamingResponse(
         generate_ai_response(chat_id, model),
         media_type="text/event-stream"
     )
 
-# Reset chat history endpoint
+# Optional: reset chat history
 @app.post("/reset")
-async def reset_chat(chat_id: str = Query(...)):
+async def reset_chat(chat_id: str = Query(..., description="ID of chat to reset")):
     if chat_id in chat_histories:
         chat_histories[chat_id].clear()
         return {"message": f"Chat {chat_id} history reset."}
-    raise HTTPException(status_code=404, detail="Chat ID not found")
+    else:
+        raise HTTPException(status_code=404, detail="Chat ID not found")
 
+# For ASGI servers like Uvicorn
 def get_app():
     return app