Update app.py

app.py

@@ -18,9 +18,13 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-OLLAMA_BASE = "http://localhost:11434"
-MODEL = os.environ.get("DEFAULT_MODEL", "qwen2.5-coder:7b-instruct-q4_K_M")
-API_TOKEN = os.environ.get("API_TOKEN", "")
+OLLAMA_BASE  = "http://localhost:11434"
+MODEL        = os.environ.get("DEFAULT_MODEL", "qwen2.5-coder:7b-instruct-q4_K_M")
+API_TOKEN    = os.environ.get("API_TOKEN", "")
+# Free CPU tier: keep context small or requests will timeout after 5 min
+MAX_CTX      = 4096
+MAX_OUT      = 1024
+TIMEOUT      = 240  # 4 min hard limit — under HF's 5 min kill
 
 
 def verify_token(creds: HTTPAuthorizationCredentials = Depends(security)):
@@ -33,7 +37,7 @@ def verify_token(creds: HTTPAuthorizationCredentials = Depends(security)):
 
 @app.get("/")
 async def root():
-    return {"status": "ok", "model": MODEL}
+    return {"status": "ok", "model": MODEL, "max_ctx": MAX_CTX}
 
 
 @app.get("/health")
@@ -41,7 +45,7 @@ async def health():
     try:
         r = requests.get(f"{OLLAMA_BASE}/api/tags", timeout=5)
         models = [m["name"] for m in r.json().get("models", [])]
-        return {"status": "ok", "model": MODEL, "available_models": models}
+        return {"status": "ok", "model": MODEL, "available_models": models, "max_ctx": MAX_CTX}
     except Exception as e:
         return {"status": "starting", "error": str(e)}
 
@@ -67,7 +71,8 @@ async def chat_completions(request: Request, token: str = Depends(verify_token))
         "messages": body.get("messages", []),
         "stream": stream,
         "options": {
-            "num_ctx": body.get("max_tokens", 32768),
+            "num_ctx": MAX_CTX,
+            "num_predict": min(body.get("max_tokens", MAX_OUT), MAX_OUT),
             "temperature": body.get("temperature", 0.7),
         }
     }
@@ -77,7 +82,7 @@ async def chat_completions(request: Request, token: str = Depends(verify_token))
             try:
                 with requests.post(
                     f"{OLLAMA_BASE}/v1/chat/completions",
-                    json=payload, stream=True, timeout=300
+                    json=payload, stream=True, timeout=TIMEOUT
                 ) as r:
                     for chunk in r.iter_content(chunk_size=None):
                         if chunk:
@@ -86,8 +91,11 @@ async def chat_completions(request: Request, token: str = Depends(verify_token))
                 yield f"data: {{\"error\": \"{e}\"}}\n\n".encode()
         return StreamingResponse(generate(), media_type="text/event-stream")
 
-    r = requests.post(f"{OLLAMA_BASE}/v1/chat/completions", json=payload, timeout=300)
-    return r.json()
+    try:
+        r = requests.post(f"{OLLAMA_BASE}/v1/chat/completions", json=payload, timeout=TIMEOUT)
+        return r.json()
+    except requests.Timeout:
+        raise HTTPException(504, "Inference timeout — try a shorter prompt")
 
 
 @app.post("/v1/messages")
@@ -101,7 +109,8 @@ async def messages(request: Request, token: str = Depends(verify_token)):
         "messages": body.get("messages", []),
         "stream": stream,
         "options": {
-            "num_ctx": body.get("max_tokens", 32768),
+            "num_ctx": MAX_CTX,
+            "num_predict": min(body.get("max_tokens", MAX_OUT), MAX_OUT),
             "temperature": body.get("temperature", 0.7),
         }
     }
@@ -117,7 +126,7 @@ async def messages(request: Request, token: str = Depends(verify_token)):
             try:
                 with requests.post(
                     f"{OLLAMA_BASE}/v1/chat/completions",
-                    json=payload, stream=True, timeout=300
+                    json=payload, stream=True, timeout=TIMEOUT
                 ) as r:
                     buf = ""
                     for chunk in r.iter_content(chunk_size=None):
@@ -153,21 +162,24 @@ async def messages(request: Request, token: str = Depends(verify_token)):
 
         return StreamingResponse(generate_anthropic(), media_type="text/event-stream")
 
-    r = requests.post(f"{OLLAMA_BASE}/v1/chat/completions", json=payload, timeout=300)
-    data = r.json()
-    content = (data.get("choices") or [{}])[0].get("message", {}).get("content", "")
-    return {
-        "id": data.get("id", f"msg_{int(time.time())}"),
-        "type": "message",
-        "role": "assistant",
-        "content": [{"type": "text", "text": content}],
-        "model": model,
-        "stop_reason": "end_turn",
-        "usage": {
-            "input_tokens": data.get("usage", {}).get("prompt_tokens", 0),
-            "output_tokens": data.get("usage", {}).get("completion_tokens", 0)
+    try:
+        r = requests.post(f"{OLLAMA_BASE}/v1/chat/completions", json=payload, timeout=TIMEOUT)
+        data = r.json()
+        content = (data.get("choices") or [{}])[0].get("message", {}).get("content", "")
+        return {
+            "id": data.get("id", f"msg_{int(time.time())}"),
+            "type": "message",
+            "role": "assistant",
+            "content": [{"type": "text", "text": content}],
+            "model": model,
+            "stop_reason": "end_turn",
+            "usage": {
+                "input_tokens": data.get("usage", {}).get("prompt_tokens", 0),
+                "output_tokens": data.get("usage", {}).get("completion_tokens", 0)
+            }
         }
-    }
+    except requests.Timeout:
+        raise HTTPException(504, "Inference timeout — try a shorter prompt")
 
 
 if __name__ == "__main__":