Update app.py

app.py

@@ -1,3 +1,4 @@
+# app.py  -- robust SSE status handling and chunk sanitization
 import os
 import json
 import logging
@@ -5,31 +6,31 @@ import asyncio
 from fastapi import FastAPI
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
+from typing import Any, Dict
 
-# separated modules
+# Local model modules (expect these to exist in your project)
 import router_model
 import coder_model
 import chat_model
 
 logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger("nexari")
+logger = logging.getLogger("nexari.app")
 
 app = FastAPI()
 
-# Primary MODEL_DIR
 MODEL_DIR = "./models"
 
 def ensure_model_dir_or_fail():
     try:
         os.makedirs(MODEL_DIR, exist_ok=True)
-        logger.info(f"Model dir ensured: {MODEL_DIR}")
+        logger.info("Model directory ensured: %s", MODEL_DIR)
     except Exception as e:
-        logger.critical(f"CRITICAL: Unable to create model directory '{MODEL_DIR}': {e}")
+        logger.critical("Unable to create model dir: %s", e)
         raise
 
 @app.on_event("startup")
 async def startup_event():
-    logger.info("⏳ STARTUP: ensuring model dir and loading models...")
+    logger.info("Startup: ensure model dir and set base dirs...")
     ensure_model_dir_or_fail()
 
     router_model.BASE_DIR = os.path.join(MODEL_DIR, "router")
@@ -42,10 +43,10 @@ async def startup_event():
         asyncio.create_task(chat_model.load_model_async()),
     ]
     results = await asyncio.gather(*tasks, return_exceptions=True)
-    for idx, res in enumerate(results):
-        if isinstance(res, Exception):
-            logger.error(f"Model loader {idx} failed during startup: {res}")
-    logger.info("✅ Startup complete.")
+    for i, r in enumerate(results):
+        if isinstance(r, Exception):
+            logger.error("Model loader %d failed: %s", i, r)
+    logger.info("Startup complete.")
 
 class Message(BaseModel):
     role: str
@@ -57,108 +58,159 @@ class ChatRequest(BaseModel):
     temperature: float = 0.7
 
 def get_intent(last_user_message: str):
-    # Fallback if router not loaded
-    if not router_model.model:
-        logger.warning("Router model not loaded — using rule-based fallback.")
-        text = last_user_message.lower()
-        if any(tok in text for tok in ["code", "function", "bug", "error", "fix", "html", "css", "python", "js"]):
+    # If router model missing, use a simple rule
+    if not getattr(router_model, "model", None):
+        text = (last_user_message or "").lower()
+        if any(tok in text for tok in ["code", "bug", "fix", "error", "function", "python", "js", "html", "css"]):
             return "coding", "neutral"
+        if any(tok in text for tok in ["why", "how", "prove", "reason", "think"]):
+            return "reasoning", "neutral"
         return "chat", "neutral"
 
-    sys_prompt = "Analyze intent. Return JSON: {'intent': 'coding' or 'chat' or 'reasoning', 'sentiment': 'neutral' or 'sad'}"
+    sys_prompt = "Analyze intent. Return JSON like {'intent':'coding'|'chat'|'reasoning', 'sentiment':'neutral'|'sad'}"
     try:
         res = router_model.model.create_chat_completion(
-            messages=[{"role": "system", "content": sys_prompt}, {"role": "user", "content": last_user_message}],
+            messages=[{"role":"system","content":sys_prompt},{"role":"user","content": last_user_message}],
             temperature=0.1, max_tokens=50
         )
-        content = res['choices'][0]['message']['content'].lower()
-        
+        content = ""
+        try:
+            content = res['choices'][0]['message']['content'].lower()
+        except Exception:
+            try:
+                content = res['choices'][0]['text'].lower()
+            except Exception:
+                content = ""
         if "coding" in content:
             return "coding", "neutral"
-        # Reasoning intent detection
         if "reasoning" in content or "think" in content or "solve" in content:
             return "reasoning", "neutral"
         if "sad" in content:
             return "chat", "sad"
         return "chat", "neutral"
     except Exception as e:
-        logger.exception(f"Router intent detection failed: {e}")
+        logger.exception("Router failure: %s", e)
         return "chat", "neutral"
 
+def sanitize_chunk(chunk: Any) -> Dict[str, Any]:
+    """
+    Ensure chunk is a JSON-serializable mapping for SSE.
+    Remove any 'status' fields so we never send an unintended status overwrite.
+    """
+    # If chunk is already a dict-like
+    if isinstance(chunk, dict):
+        # shallow copy to avoid mutating model internals
+        out = {}
+        for k, v in chunk.items():
+            if k == "status":
+                # drop status fields from model-chunks; log for diagnostics
+                logger.debug("Dropping status field from model chunk: %s", v)
+                continue
+            # try to keep strings and numbers; for complex objects convert to str
+            if isinstance(v, (str, int, float, bool, type(None))):
+                out[k] = v
+            else:
+                try:
+                    json.dumps(v)
+                    out[k] = v
+                except Exception:
+                    out[k] = str(v)
+        return out
+    else:
+        # Not a dict: coerce into a safe dict with text key
+        try:
+            # if it's bytes or similar, convert
+            txt = str(chunk)
+            return {"text": txt}
+        except Exception:
+            return {"text": "[UNSERIALIZABLE_CHUNK]"}
+
 @app.post("/v1/chat/completions")
 async def chat_endpoint(request: ChatRequest):
-    messages = [m.dict() for m in request.messages]
+    # Validate incoming
+    messages = [m.dict() for m in request.messages] if request.messages else []
     if not messages:
         return {"error": "No messages provided."}
-    last_msg = messages[-1]['content']
+    last = messages[-1]['content']
 
-    # 1. Detect Intent & Sentiment
-    intent, sentiment = get_intent(last_msg)
+    intent, sentiment = get_intent(last)
 
     selected_model = None
-    sys_msg = "You are a helpful AI."
-    
-    # 2. Set Status Indicator Text
-    status_indicator = "Thinking..." # Default
+    sys_msg = "You are a helpful assistant."
+    status_indicator = "Thinking..."  # default if not changed below
 
     if intent == "coding":
-        if not coder_model.model:
-            logger.error("Client requested coding intent but coder model is not loaded.")
-            return {"error": "Coder model not available."}
+        if not getattr(coder_model, "model", None):
+            logger.error("Coder model not available.")
+            return {"error":"Coder model not available."}
         selected_model = coder_model.model
-        sys_msg = "You are an expert Coding Assistant. Write clean, efficient code."
+        sys_msg = "You are a coding expert. Provide clean code."
         status_indicator = "Coding..."
-        logger.info("Using: CODER")
-
+        logger.info("Intent: CODING")
     elif intent == "reasoning":
-        # Currently using Chat model for reasoning, but indicating "Reasoning..."
-        # If you have a specific reasoning model (e.g. DeepSeek-R1), add logic here.
-        if not chat_model.model:
-             return {"error": "Model not available."}
+        if not getattr(chat_model, "model", None):
+            logger.error("Chat model not available for reasoning.")
+            return {"error":"Model not available."}
         selected_model = chat_model.model
-        status_indicator = "Reasoning..." 
-        logger.info("Using: REASONING (via Chat Model)")
-        
+        status_indicator = "Reasoning..."
+        logger.info("Intent: REASONING")
     else:
-        if not chat_model.model:
-            logger.error("Chat model not loaded.")
-            return {"error": "Chat model not available."}
+        if not getattr(chat_model, "model", None):
+            logger.error("Chat model missing.")
+            return {"error":"Chat model not available."}
         selected_model = chat_model.model
-        logger.info(f"Using: CHAT (Sentiment: {sentiment})")
+        logger.info("Intent: CHAT (%s)", sentiment)
         if sentiment == "sad":
-            sys_msg = "You are an empathetic friend."
+            sys_msg = "You are empathic and calm."
             status_indicator = "Empathizing..."
 
-    # Ensure system prompt
-    if messages[0]['role'] != "system":
-        messages.insert(0, {"role": "system", "content": sys_msg})
+    # ensure system prompt is present
+    if messages[0].get("role") != "system":
+        messages.insert(0, {"role":"system","content": sys_msg})
     else:
-        messages[0]['content'] = sys_msg
+        messages[0]["content"] = sys_msg
 
-    # 3. Stream Response with Status Packet
+    # Streaming generator
     def iter_response():
         try:
-            # Send status packet first (frontend expects this to update its indicator)
-            yield f"data: {json.dumps({'status': status_indicator})}\n\n"
-
-            # Start model streaming generator
+            # 1) Send a single authoritative SSE status event (event: status)
+            # Use event field so client can handle it separately from data token stream.
+            status_payload = json.dumps({"status": status_indicator})
+            event_payload = f"event: status\n"
+            event_payload += f"data: {status_payload}\n\n"
+            logger.info("Sending authoritative status event: %s", status_indicator)
+            yield event_payload
+
+            # 2) small flush hint to reduce buffering and help client parse promptly
+            yield ":\n\n"
+
+            # 3) Start streaming model output
             stream = selected_model.create_chat_completion(
                 messages=messages,
                 temperature=request.temperature,
                 stream=True
             )
 
-            # **CRITICAL: send the status packet again immediately after starting the stream**
-            # This guarantees the frontend will receive the server's authoritative status even
-            # if it creates the indicator DOM slightly later (race condition on the client).
-            yield f"data: {json.dumps({'status': status_indicator})}\n\n"
-
+            # Iterate model generator and sanitize every chunk so it cannot inject a status
             for chunk in stream:
-                yield f"data: {json.dumps(chunk)}\n\n"
+                safe = sanitize_chunk(chunk)
+                try:
+                    yield f"data: {json.dumps(safe)}\n\n"
+                except Exception:
+                    # fallback to a safe string representation
+                    yield f"data: {json.dumps({'text': str(safe)})}\n\n"
+
+            # 4) final done marker
             yield "data: [DONE]\n\n"
+            logger.info("Stream finished for request (status was: %s)", status_indicator)
+
         except Exception as e:
-            logger.exception(f"Error while streaming: {e}")
-            yield f"data: {json.dumps({'error': str(e)})}\n\n"
+            logger.exception("Streaming error: %s", e)
+            # send explicit error object
+            try:
+                yield f"data: {json.dumps({'error': str(e)})}\n\n"
+            except Exception:
+                yield "data: {\"error\":\"streaming failure\"}\n\n"
             yield "data: [DONE]\n\n"
 
     return StreamingResponse(iter_response(), media_type="text/event-stream")