Update app.py

app.py

@@ -4,7 +4,7 @@ import subprocess
 import re
 
 # =============================================================
-# 1. Bootstrap: Install prebuilt CPU/ZeroGPU llama-cpp-python
+# 1. Bootstrap prebuilt llama-cpp-python
 # =============================================================
 def bootstrap_llama_cpp():
     print("Bootstrapping llama-cpp-python...")
@@ -42,9 +42,9 @@ except (ImportError, RuntimeError, OSError):
     from llama_cpp import Llama
 
 import spaces
-from typing import List, Literal, Optional
+from typing import List, Literal, Optional, Any
 from pydantic import BaseModel, Field
-from fastapi import HTTPException, Header, Request
+from fastapi import HTTPException, Header
 import gradio as gr
 
 # =============================================================
@@ -69,14 +69,11 @@ SYSTEM_PROMPT = (
     "earnings, and low traffic delays."
 )
 
-# =============================================================
-# 3. Core ZeroGPU Inference function
-# =============================================================
 @spaces.GPU
-def generate_response(prompt: str, history_messages: list = None) -> str:
+def run_inference(prompt: str, history: List[dict] = None) -> str:
     messages = [{"role": "system", "content": SYSTEM_PROMPT}]
-    if history_messages:
-        for m in history_messages:
+    if history:
+        for m in history:
             messages.append({"role": m["role"], "content": m["content"]})
     messages.append({"role": "user", "content": prompt})
     response = llm.create_chat_completion(
@@ -85,7 +82,7 @@ def generate_response(prompt: str, history_messages: list = None) -> str:
     return response["choices"][0]["message"]["content"]
 
 # =============================================================
-# 4. Gradio Chat UI with @spaces.GPU
+# 3. Gradio Interface (Satisfies ZeroGPU startup check)
 # =============================================================
 @spaces.GPU
 def gradio_chat(user_message, history):
@@ -95,7 +92,7 @@ def gradio_chat(user_message, history):
             hist_formatted.append({"role": "user", "content": user_msg})
             if bot_msg:
                 hist_formatted.append({"role": "assistant", "content": bot_msg})
-    return generate_response(user_message, hist_formatted)
+    return run_inference(user_message, hist_formatted)
 
 demo = gr.ChatInterface(
     fn=gradio_chat,
@@ -104,7 +101,7 @@ demo = gr.ChatInterface(
 )
 
 # =============================================================
-# 5. Attach POST /chat to demo.app using Gradio's native router
+# 4. Custom API Schemas (Matching CommitCon/src/phase5_api.py)
 # =============================================================
 class HistoryMessage(BaseModel):
     role: Literal["user", "assistant"]
@@ -113,15 +110,15 @@ class HistoryMessage(BaseModel):
 class ChatRequest(BaseModel):
     driver_id: Optional[str] = "DRV_DEFAULT"
     driver_prompt: str = Field(min_length=1, max_length=2000)
+    recommendation_data: Optional[dict[str, Any]] = None
     current_zone_id: Optional[str] = None
     history: List[HistoryMessage] = Field(default_factory=list)
 
-@demo.app.post("/chat")
-def chat_endpoint(payload: ChatRequest, request: Request, x_api_key: Optional[str] = Header(None)):
+def chat_endpoint(payload: ChatRequest, x_api_key: Optional[str] = Header(None, alias="X-API-Key")):
     if API_KEY and x_api_key != API_KEY:
         raise HTTPException(status_code=401, detail="Unauthorized: Invalid API Key")
     formatted_history = [{"role": m.role, "content": m.content} for m in payload.history]
-    answer = generate_response(payload.driver_prompt, formatted_history)
+    answer = run_inference(payload.driver_prompt, formatted_history)
     return {
         "status": "success",
         "answer": answer,
@@ -129,8 +126,11 @@ def chat_endpoint(payload: ChatRequest, request: Request, x_api_key: Optional[st
         "zone_context": payload.current_zone_id
     }
 
+# Register the POST /chat route directly into Gradio's underlying FastAPI router
+demo.app.add_api_route("/chat", chat_endpoint, methods=["POST"])
+
 # =============================================================
-# 6. Launch with demo.launch() (DO NOT USE UVICORN DIRECTLY ON ZEROGPU)
+# 5. Native ZeroGPU Launch
 # =============================================================
 if __name__ == "__main__":
     demo.launch()