Update app.py

app.py

@@ -1,14 +1,16 @@
 from fastapi import FastAPI, HTTPException
+from fastapi.responses import StreamingResponse
 import requests
+import json
 
 app = FastAPI()
 
-# Ollama internal URL (running locally in the container)
+# Ollama internal URL
 OLLAMA_BASE_URL = "http://localhost:11434"
 
-# Proxy endpoint to Ollama's API
+# Generic proxy for other API endpoints
 @app.get("/api/{path:path}")
-async def ollama_proxy(path: str, query: str = None):
+async def ollama_proxy_get(path: str, query: str = None):
     url = f"{OLLAMA_BASE_URL}/api/{path}"
     params = {"query": query} if query else {}
     try:
@@ -18,16 +20,25 @@ async def ollama_proxy(path: str, query: str = None):
     except requests.exceptions.RequestException as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-@app.post("/api/{path:path}")
-async def ollama_proxy_post(path: str, body: dict):
-    url = f"{OLLAMA_BASE_URL}/api/{path}"
+# Handle /api/chat specifically
+@app.post("/api/chat")
+async def ollama_chat(body: dict):
+    url = f"{OLLAMA_BASE_URL}/api/chat"
     try:
-        response = requests.post(url, json=body)
+        # Forward the request to Ollama with streaming support
+        response = requests.post(url, json=body, stream=True)
         response.raise_for_status()
-        return response.json()
+
+        # Stream the response back to the client
+        def generate():
+            for chunk in response.iter_lines():
+                if chunk:
+                    yield chunk + b"\n"
+
+        return StreamingResponse(generate(), media_type="text/event-stream")
     except requests.exceptions.RequestException as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        raise HTTPException(status_code=500, detail=f"Ollama error: {str(e)}")
 
 @app.get("/")
 async def root():
-    return {"message": "Ollama running on Hugging Face Spaces! use the space url"}
+    return {"message": "Ollama API proxy running on Hugging Face Spaces!"}