Update app.py

app.py

@@ -1,11 +1,4 @@
-"""
-app.py - Nexari Server (Web-tool + Detail-request fixes)
-Key changes:
- - When web search performed, pass structured WEB_DATA as an assistant message so the model MUST use it.
- - Detect "detailed/line-by-line" user requests and increase max_tokens & enforce numbered output format.
- - Minor safety & streaming robustness retained.
-"""
-
+# app.py - UPDATED (replace original)
 import re
 import json
 import asyncio
@@ -21,13 +14,15 @@ from tools_engine import analyze_intent, perform_web_search
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import gradio as gr
+import os
 
 logger = logging.getLogger("nexari")
 logging.basicConfig(level=logging.INFO)
 
-MODEL_ID = "Piyush-boss/Nexari-Qwen-3B-Full"
+MODEL_ID = os.environ.get("MODEL_ID", "Piyush-boss/Nexari-Qwen-3B-Full")
 tokenizer = None
 model = None
+device = "cpu"
 
 app = FastAPI()
 
@@ -45,28 +40,52 @@ def is_detailed_request(text: str) -> bool:
 
 @app.on_event("startup")
 async def startup_event():
-    global tokenizer, model
-    logger.info("Startup: loading tokenizer/model in background thread...")
+    global tokenizer, model, device
+    logger.info("Startup: initiating background model load...")
     try:
+        # choose device more safely
+        if torch.cuda.is_available():
+            device = "cuda"
+        else:
+            device = "cpu"
+
         def sync_load():
             tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+            # Use device_map when possible; use cpu if no GPU
             mdl = AutoModelForCausalLM.from_pretrained(
                 MODEL_ID,
-                dtype=None,
-                device_map="cpu",
-                low_cpu_mem_usage=True,
                 trust_remote_code=True,
+                low_cpu_mem_usage=(device == "cpu"),
+                device_map="auto" if device == "cuda" else None
             )
+            if device == "cpu":
+                mdl.to("cpu")
             return tok, mdl
+
         tokenizer, model = await asyncio.to_thread(sync_load)
-        logger.info("Model loaded successfully.")
+        logger.info("Model loaded successfully on %s.", device)
     except Exception as e:
-        logger.exception(f"Model loading failed: {e}")
+        logger.exception(f"Model loading failed at startup: {e}")
         tokenizer, model = None, None
 
-@app.on_event("shutdown")
-async def shutdown_event():
-    logger.info("Shutdown: cleanup if necessary.")
+def _build_prompt_from_messages(messages):
+    """
+    Build a single-string prompt for models in case tokenizer.apply_chat_template
+    is unavailable (fallback).
+    """
+    parts = []
+    for m in messages:
+        role = m.get("role","user")
+        content = m.get("content","")
+        if role == "system":
+            parts.append(f"[SYSTEM]\n{content}\n")
+        elif role == "user":
+            parts.append(f"[USER]\n{content}\n")
+        elif role == "assistant":
+            parts.append(f"[ASSISTANT]\n{content}\n")
+        else:
+            parts.append(content)
+    return "\n".join(parts)
 
 async def generate_response_stream(messages, max_tokens=600, temperature=0.85):
     try:
@@ -74,7 +93,6 @@ async def generate_response_stream(messages, max_tokens=600, temperature=0.85):
             messages = [{"role":"user","content":""}]
         last_user_msg = messages[-1].get("content","")
 
-        # initial thinking status
         yield f"data: {json.dumps({'status': 'Thinking...'})}\n\n"
         await asyncio.sleep(0)
 
@@ -85,19 +103,19 @@ async def generate_response_stream(messages, max_tokens=600, temperature=0.85):
         vibe_data = ""
         strategy_data = ""
 
-        # detect if user explicitly asked for long/detailed format
         want_detailed = is_detailed_request(last_user_msg)
         if want_detailed:
-            # bump tokens to allow long/line-by-line answer
             max_tokens = max(max_tokens, 1200)
-            temperature = min(temperature, 0.9)  # keep somewhat controlled
+            temperature = min(temperature, 0.9)
 
-        # Route based on intent
         if intent == "internet_search":
             yield f"data: {json.dumps({'status': 'Searching the web...'})}\n\n"
             await asyncio.sleep(0)
-            # perform_web_search now returns structured dict or empty string
-            tool_data_struct = perform_web_search(last_user_msg)
+            try:
+                tool_data_struct = perform_web_search(last_user_msg)
+            except Exception as e:
+                logger.exception("Web search failed: %s", e)
+                tool_data_struct = {"query": last_user_msg, "results": []}
             vibe_data = get_smart_context(last_user_msg)
             strategy_data = get_thinking_strategy(is_complex=True, detail=want_detailed)
 
@@ -116,7 +134,6 @@ async def generate_response_stream(messages, max_tokens=600, temperature=0.85):
             vibe_data = get_smart_context(last_user_msg)
             strategy_data = get_thinking_strategy(is_complex=False, detail=want_detailed)
 
-        # Base system instruction with explicit web-data usage rule
         base_system_instruction = (
             "### SYSTEM IDENTITY ###\n"
             "You are Nexari G1, an expressive and helpful AI created by Piyush.\n"
@@ -130,17 +147,14 @@ async def generate_response_stream(messages, max_tokens=600, temperature=0.85):
 
         final_system_prompt = f"{base_system_instruction}\n{vibe_data}\n{time_data}\n{strategy_data}"
 
-        # ensure system message present
+        # Ensure system message present
         if messages[0].get("role") != "system":
             messages.insert(0, {"role":"system","content": final_system_prompt})
         else:
             messages[0]["content"] = final_system_prompt
 
-        # If we have tool_data_struct (dict with items & sources), add as assistant message
         if tool_data_struct:
-            # create a clear WEB_DATA assistant message that model must consume
             web_block = "### WEB_DATA (from live search) ###\n"
-            # include numbered sources with short titles, snippets, and urls
             items = tool_data_struct.get("results", [])
             if items:
                 lines = []
@@ -153,12 +167,8 @@ async def generate_response_stream(messages, max_tokens=600, temperature=0.85):
                 web_block += "\n---\nINSTRUCTION: Use the WEB_DATA above to answer; cite relevant source numbers inline."
             else:
                 web_block += "No results found."
-
-            # Insert the web block as an assistant message so model treats it as retrieved evidence
-            # Insert after system message (index 1)
             messages.insert(1, {"role":"assistant","content": web_block})
 
-        # Model availability check
         if tokenizer is None or model is None:
             err = "Model is not loaded on server. Please check logs."
             payload = json.dumps({"choices":[{"delta":{"content": err}}]})
@@ -166,15 +176,22 @@ async def generate_response_stream(messages, max_tokens=600, temperature=0.85):
             yield "data: [DONE]\n\n"
             return
 
-        # prepare prompt
-        text_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        model_inputs = tokenizer([text_prompt], return_tensors="pt").to(model.device)
+        # Build prompt with tokenizer helper if available, else fallback
+        try:
+            if hasattr(tokenizer, "apply_chat_template"):
+                text_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+            else:
+                text_prompt = _build_prompt_from_messages(messages)
+        except Exception:
+            text_prompt = _build_prompt_from_messages(messages)
+
+        model_inputs = tokenizer(text_prompt, return_tensors="pt", truncation=True, max_length=4096).to(next(model.parameters()).device)
 
         status_msg = 'Reading results...' if tool_data_struct else 'Responding...'
         yield f"data: {json.dumps({'status': status_msg})}\n\n"
 
-        # Generation in thread
         def sync_generate():
+            # wrap generate in try to catch resource errors
             return model.generate(
                 **model_inputs,
                 max_new_tokens=max_tokens,
@@ -184,26 +201,29 @@ async def generate_response_stream(messages, max_tokens=600, temperature=0.85):
                 top_p=0.92,
                 repetition_penalty=1.08
             )
-        generated_ids = await asyncio.to_thread(sync_generate)
+        try:
+            generated_ids = await asyncio.to_thread(sync_generate)
+        except RuntimeError as e:
+            logger.exception("Generation failed (possible OOM): %s", e)
+            err_payload = json.dumps({"choices":[{"delta":{"content": "Model generation failed due to resource limits."}}]})
+            yield f"data: {err_payload}\n\n"
+            yield "data: [DONE]\n\n"
+            return
 
-        input_len = model_inputs.input_ids.shape[1]
+        input_len = model_inputs["input_ids"].shape[1]
         new_tokens = generated_ids[0][input_len:]
         raw_response = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
 
-        # Cleaning & safety
         cleaned = safe_replace_providers(raw_response)
         forbidden = ["I am a human","I have a physical body","I am alive"]
         for fc in forbidden:
             if fc.lower() in cleaned.lower():
                 cleaned = re.sub(re.escape(fc), "I am an AI — expressive and interactive.", cleaned, flags=re.IGNORECASE)
 
-        # If detailed requested, encourage numbered formatting if model didn't follow
         if want_detailed:
-            # simple heuristic: if no numbered lines present, add an instruction prefix
             if not re.search(r"^\s*\d+[\.\)]\s+", cleaned, re.M):
-                cleaned = "1) " + cleaned.replace("\n", "\n2) ")  # best-effort reformat
+                cleaned = "1) " + cleaned.replace("\n", "\n2) ")
 
-        # Format Answer tag
         cleaned = cleaned.replace("💡 **Answer:**", "\n\n---\n💡 **Answer:**")
 
         payload = json.dumps({"choices":[{"delta":{"content": cleaned}}]})
@@ -224,7 +244,8 @@ async def generate_response_stream(messages, max_tokens=600, temperature=0.85):
 
 @app.get("/api/status")
 def status():
-    return {"status":"online","mode":"Smart Override Enabled"}
+    ok = tokenizer is not None and model is not None
+    return {"status":"online" if ok else "degraded", "mode":"Smart Override Enabled", "model_loaded": ok}
 
 @app.post("/v1/chat/completions")
 async def chat_completions(request: Request):
@@ -236,9 +257,10 @@ async def chat_completions(request: Request):
         logger.exception(f"chat_completions endpoint error: {e}")
         return {"error": str(e)}
 
-# Mount gradio if create_ui returns Blocks (must not call .launch())
+# Mount gradio if create_ui returns Blocks
 try:
     demo = create_ui(lambda messages: "Use API")
+    # demo should be a Blocks or Component; mount safely
     app = gr.mount_gradio_app(app, demo, path="/")
     logger.info("Gradio mounted.")
 except Exception as e:
@@ -246,4 +268,4 @@ except Exception as e:
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 7860)))