Update app.py

app.py

@@ -1,19 +1,24 @@
-# app.py  -- robust SSE status handling and chunk sanitization
-# + system identity prompt (Nexari-G1, creator: Piyush)
+# app.py  -- upgraded to use duckduckgo-based search and a lightweight local image generator
 import os
 import json
 import logging
 import asyncio
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from typing import Any, Dict
 
-# Local model modules (expect these to exist in your project)
+# Local model modules
 import router_model
 import coder_model
 import chat_model
 
+# New utilities
+import intent_model
+import web_search           # duckduckgo-based wrapper
+import image_gen            # lightweight CPU image generator
+from time_utils import parse_time_iso
+
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("nexari.app")
 
@@ -38,15 +43,17 @@ async def startup_event():
     coder_model.BASE_DIR = os.path.join(MODEL_DIR, "coder")
     chat_model.BASE_DIR = os.path.join(MODEL_DIR, "chat")
 
+    # load models + intent model concurrently
     tasks = [
         asyncio.create_task(router_model.load_model_async()),
         asyncio.create_task(coder_model.load_model_async()),
         asyncio.create_task(chat_model.load_model_async()),
+        asyncio.create_task(intent_model.load_model_async()),
     ]
     results = await asyncio.gather(*tasks, return_exceptions=True)
     for i, r in enumerate(results):
         if isinstance(r, Exception):
-            logger.error("Model loader %d failed: %s", i, r)
+            logger.error("Loader %d failed: %s", i, r)
     logger.info("Startup complete.")
 
 class Message(BaseModel):
@@ -57,57 +64,26 @@ class ChatRequest(BaseModel):
     messages: list[Message]
     stream: bool = True
     temperature: float = 0.7
-
-def get_intent(last_user_message: str):
-    # 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 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}],
-            temperature=0.1, max_tokens=50
-        )
-        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"
-        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("Router failure: %s", e)
-        return "chat", "neutral"
+    # allow explicit user-requested tools (no API keys required)
+    use_web_search: bool = False
+    use_image_gen: bool = False
+    # optional: time parsing example parameter
+    time_hint: str = None
+    # optional image params (size/generate style) — for future extension
+    image_params: dict = None
+
+def get_intent_and_sentiment(last_user_message: str):
+    if not getattr(intent_model, "model", None):
+        raise RuntimeError("Intent model not loaded. Start-up failed or model missing.")
+    return intent_model.classify(last_user_message)
 
 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:
@@ -118,15 +94,12 @@ def sanitize_chunk(chunk: Any) -> Dict[str, Any]:
                     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]"}
 
-# Static system identity prefix to include in system prompts:
 SYSTEM_IDENTITY_PREFIX = (
     "You are Nexari-G1, an AI assistant created by Piyush (developer name: Piyush). "
     "always understand the user behaviour and request. "
@@ -137,23 +110,32 @@ SYSTEM_IDENTITY_PREFIX = (
 
 @app.post("/v1/chat/completions")
 async def chat_endpoint(request: ChatRequest):
-    # Validate incoming
     messages = [m.dict() for m in request.messages] if request.messages else []
     if not messages:
-        return {"error": "No messages provided."}
+        raise HTTPException(status_code=400, detail="No messages provided.")
     last = messages[-1]['content']
 
-    intent, sentiment = get_intent(last)
+    if request.time_hint:
+        try:
+            parsed = parse_time_iso(request.time_hint)
+            logger.info("Parsed user time_hint -> %s", parsed.isoformat())
+        except Exception as e:
+            logger.warning("time_hint parse failed: %s", e)
+
+    try:
+        intent, sentiment = get_intent_and_sentiment(last)
+    except Exception as e:
+        logger.exception("Intent detection failed: %s", e)
+        raise HTTPException(status_code=500, detail=f"Intent detection failed: {e}")
 
     selected_model = None
-    # base system message will always include identity prefix
     sys_msg = SYSTEM_IDENTITY_PREFIX + "You are a helpful assistant."
-    status_indicator = "Thinking..."  # default if not changed below
+    status_indicator = "Thinking..."
 
     if intent == "coding":
         if not getattr(coder_model, "model", None):
             logger.error("Coder model not available.")
-            return {"error":"Coder model not available."}
+            raise HTTPException(status_code=500, detail="Coder model not available.")
         selected_model = coder_model.model
         sys_msg = SYSTEM_IDENTITY_PREFIX + "You are an expert Coding Assistant. Write clean, efficient code with comments where helpful."
         status_indicator = "Coding..."
@@ -161,7 +143,7 @@ async def chat_endpoint(request: ChatRequest):
     elif intent == "reasoning":
         if not getattr(chat_model, "model", None):
             logger.error("Chat model not available for reasoning.")
-            return {"error":"Model not available."}
+            raise HTTPException(status_code=500, detail="Chat model not available.")
         selected_model = chat_model.model
         sys_msg = SYSTEM_IDENTITY_PREFIX + "You are a reasoning-focused assistant. Walk through your thinking clearly and show steps if relevant."
         status_indicator = "Reasoning..."
@@ -169,64 +151,76 @@ async def chat_endpoint(request: ChatRequest):
     else:
         if not getattr(chat_model, "model", None):
             logger.error("Chat model missing.")
-            return {"error":"Chat model not available."}
+            raise HTTPException(status_code=500, detail="Chat model not available.")
         selected_model = chat_model.model
         logger.info("Intent: CHAT (%s)", sentiment)
         if sentiment == "sad":
             sys_msg = SYSTEM_IDENTITY_PREFIX + "You are empathic and calm. Provide supportive, concise responses."
             status_indicator = "Empathizing..."
         else:
-            # default chat system message with identity included
             sys_msg = SYSTEM_IDENTITY_PREFIX + "You are a helpful conversational assistant."
 
-    # ensure system prompt is present (first message)
     if messages[0].get("role") != "system":
         messages.insert(0, {"role":"system","content": sys_msg})
     else:
-        # replace existing system content to ensure identity is present and consistent
         messages[0]["content"] = sys_msg
 
-    # Streaming generator
+    tool_context = {}
+    # Web search via duckduckgo (no API keys)
+    if request.use_web_search:
+        try:
+            logger.info("User requested web search for: %s", last)
+            snippets = web_search.search(last, max_results=3)
+            tool_context['web_search'] = snippets
+            messages.append({"role":"system","content": f"Web search results (top 3):\n{json.dumps(snippets)[:4000]}"})
+        except Exception as e:
+            logger.exception("Web search failed: %s", e)
+            messages.append({"role":"system","content": f"[Web search failed: {e}]"})
+    # Image generation using local CPU-friendly generator
+    if request.use_image_gen:
+        try:
+            logger.info("User requested image generation for: %s", last)
+            # call a synchronous local generator which returns metadata (path)
+            img_meta = image_gen.generate_image(prompt=last, params=request.image_params or {})
+            tool_context['image_result'] = img_meta
+            messages.append({"role":"system","content": f"Image generated: {json.dumps(img_meta)}"})
+        except Exception as e:
+            logger.exception("Image generation failed: %s", e)
+            messages.append({"role":"system","content": f"[Image generation failed: {e}]"})
+
+
     def iter_response():
         try:
-            # 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
             )
 
-            # Iterate model generator and sanitize every chunk so it cannot inject a status
             for chunk in stream:
                 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("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")
+    return StreamingResponse(iter_response(), media_type="text/event-stream")