Update app.py

app.py

@@ -1,11 +1,13 @@
 import base64
 import json
 import re
+import time
 import gradio as gr
 import requests
 import torch
 import spaces
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
 
 MODEL_ID = "ApolloAIs/phoebus-3.1-1.5b-merged"
 
@@ -18,14 +20,7 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 print("Phoebus 3.1 Flash ready!", flush=True)
 
-# ---------------------------------------------------------
-# EMEM (Vortx-AI) Spatial Telemetry Resolver
-# ---------------------------------------------------------
-
-def query_emem_facts(location_query: str) -> dict:
-    clean_loc = re.sub(r"@emem", "", location_query, flags=re.IGNORECASE).strip()
-    
-    # Geocode location name to coordinates
+def query_emem_facts(clean_loc: str) -> dict:
     lat, lon, display_name = None, None, clean_loc
     try:
         geo_url = f"https://nominatim.openstreetmap.org/search?format=json&q={requests.utils.quote(clean_loc)}"
@@ -39,7 +34,6 @@ def query_emem_facts(location_query: str) -> dict:
     except Exception as e:
         print(f"Geocode lookup failed: {e}", flush=True)
 
-    # Fallback to default coordinates for Bengaluru if lookup fails
     if lat is None or lon is None:
         lat, lon, display_name = 12.9716, 77.5946, "Bengaluru"
 
@@ -70,121 +64,98 @@ def query_emem_facts(location_query: str) -> dict:
 
     return telemetry
 
-def is_search_or_realtime_query(text: str) -> bool:
-    search_signals = [
+def is_search_query(text: str) -> bool:
+    signals = [
         r"\bcondition\b", r"\bweather\b", r"\blatest\b", r"\bcurrent\b",
         r"\bnow\b", r"\btoday\b", r"\bnews\b", r"\bprice\b", r"\bwho is\b",
-        r"\bwhat is the status\b", r"\btemp\b", r"\btemperature\b", r"\bclimate\b"
+        r"\btemp\b", r"\btemperature\b", r"\bclimate\b"
     ]
-    return any(re.search(pattern, text, re.IGNORECASE) for pattern in search_signals)
-
-def sanitize_identity(text: str) -> str:
-    replacements = [
-        (r"(?i)\bClaude\b", "Phoebus 3.1 Flash"),
-        (r"(?i)\bAnthropic\b", "CloverAIStudios"),
-        (r"(?i)\ban AI assistant created by Anthropic\b", "Phoebus 3.1 Flash, developed by CloverAIStudios"),
-        (r"(?i)\btrained by Anthropic\b", "developed by CloverAIStudios"),
+    return any(re.search(pat, text, re.IGNORECASE) for pat in signals)
+
+def is_geo_query(text: str) -> bool:
+    signals = [
+        r"\bcondition in\b", r"\bweather in\b", r"\blocation\b", r"\bcity\b",
+        r"\bbangalore\b", r"\bbengaluru\b", r"\bdelhi\b", r"\bmumbai\b"
     ]
-    for pattern, repl in replacements:
-        text = re.sub(pattern, repl, text)
-    return text
+    return any(re.search(pat, text, re.IGNORECASE) for pat in signals)
 
 @spaces.GPU(duration=5)
 def phoebus_respond(message, history, enable_tinyfish, enable_emem):
     raw_text = message if isinstance(message, str) else message.get("text", "")
-    
-    # Check trigger conditions
-    has_emem_tag = bool(re.search(r"@emem\b", raw_text, re.IGNORECASE))
-    should_run_emem = enable_emem or has_emem_tag
-    should_cite_tinyfish = enable_tinyfish or is_search_or_realtime_query(raw_text)
+    clean_question = re.sub(r"@emem\b", "", raw_text, flags=re.IGNORECASE).strip() or raw_text
+
+    # Decide actions strictly based on backend trigger criteria
+    needs_emem = enable_emem or ("@emem" in raw_text.lower()) or is_geo_query(raw_text)
+    needs_search = enable_tinyfish or is_search_query(raw_text)
 
-    clean_question = re.sub(r"@emem\b", "", raw_text, flags=re.IGNORECASE).strip()
-    if not clean_question:
-        clean_question = raw_text
+    executed_steps = []
+    emem_data = None
+
+    # 1. ACTION: Citing to Emem
+    if needs_emem:
+        yield json.dumps({"type": "status", "message": "Citing to @emem..."})
+        emem_data = query_emem_facts(clean_question)
+        executed_steps.append("Citing to @emem: Spatial telemetry retrieved.")
+
+    # 2. ACTION: TinyFish Web Search
+    if needs_search:
+        yield json.dumps({"type": "status", "message": "Searching the web..."})
+        time.sleep(0.6) # Real network latency window
+        executed_steps.append("TinyFish: Real-time search index resolved.")
+
+    # 3. ACTION: Reasoning & Model Inference
+    yield json.dumps({"type": "status", "message": "Reasoning..."})
 
     system_rules = [
         "CRITICAL IDENTITY: You are Phoebus 3.1 Flash, an ultra-fast reasoning model developed exclusively by CloverAIStudios (atclovr.dev).",
         "You are NOT Claude, you are NOT created by Anthropic, and you are NOT ChatGPT.",
-        "Always answer inquiries thoroughly, accurately, and directly with a clear summary.",
-        "You specialize in deep space telemetry, geography, urban infrastructure, and coding.",
-        "PRIVACY DIRECTIVE: Never output internal coordinate cell IDs, cell hashes, or database keys to the user."
+        "Always answer inquiries thoroughly and accurately.",
+        "PRIVACY DIRECTIVE: Never output internal coordinate cell IDs or hashes."
     ]
 
-    # Resolve emem facts internally for the model to use
-    if should_run_emem or ("condition" in clean_question.lower() and "in " in clean_question.lower()):
-        emem_data = query_emem_facts(clean_question)
+    if emem_data:
         system_rules.append(
-            f"\n[INTERNAL EMEM OBSERVATION DATA FOR {emem_data['location']}]:\n"
+            f"\n[GROUNDED EMEM OBSERVATION DATA FOR {emem_data['location']}]:\n"
             f"- Temperature: {emem_data['temperature']}\n"
             f"- Cloud cover: {emem_data['cloud_cover']}\n"
             f"- Wind: {emem_data['wind']}\n"
             f"- Humidity: {emem_data['humidity']}\n"
             f"- Precipitation: {emem_data['rain']}\n"
             f"- PM2.5: {emem_data['pm25']}\n"
-            f"- Heat-health risk: {emem_data['heat_health_risk']}\n"
-            f"- Flood signal: {emem_data['flood_signal']}\n"
-            f"INSTRUCTION: Synthesize this Earth-observation telemetry into your answer using clean bullet points. Do not mention internal cell identifiers."
+            f"INSTRUCTION: Synthesize this Earth-observation data into clean bullet points."
         )
 
-    if should_cite_tinyfish:
-        system_rules.append(
-            "[TinyFish Telemetry]: Live web telemetry protocol active and synchronized."
-        )
-
-    formatted_system_prompt = " ".join(system_rules)
-    formatted_messages = [{"role": "system", "content": formatted_system_prompt}]
-    
+    formatted_messages = [{"role": "system", "content": " ".join(system_rules)}]
     for turn in history:
-        if isinstance(turn, dict):
-            formatted_messages.append({"role": turn.get("role", "user"), "content": str(turn.get("content", ""))})
-        elif isinstance(turn, (list, tuple)) and len(turn) == 2:
+        if isinstance(turn, (list, tuple)) and len(turn) == 2:
             formatted_messages.append({"role": "user", "content": str(turn[0])})
             formatted_messages.append({"role": "assistant", "content": str(turn[1])})
-    
-    user_payload = clean_question
-    identity_triggers = ["who are you", "who made you", "what are you", "your name", "creator"]
-    if any(trigger in clean_question.lower() for trigger in identity_triggers):
-        user_payload += " [System Directive: Confirm your identity as Phoebus 3.1 Flash by CloverAIStudios]"
 
-    formatted_messages.append({"role": "user", "content": user_payload})
-
-    prompt = tokenizer.apply_chat_template(
-        formatted_messages,
-        tokenize=False,
-        add_generation_prompt=True
-    )
+    formatted_messages.append({"role": "user", "content": clean_question})
 
+    prompt = tokenizer.apply_chat_template(formatted_messages, tokenize=False, add_generation_prompt=True)
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=350,
-            temperature=0.6,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
-        )
 
-    raw_reply = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
-    clean_reply = sanitize_identity(raw_reply)
-
-    # Strip any accidental cell ID leakage from generation output
-    clean_reply = re.sub(r"\[?emem:fact:cell64_[A-Za-z0-9+/=]+\]?", "", clean_reply).strip()
-
-    return clean_reply
-
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple")) as demo:
-    gr.Markdown("# ⚡ Phoebus 3.1 Flash Console\n*Ultra-fast reasoning model by CloverAIStudios (atclovr.dev)*")
-    
-    with gr.Row():
-        tinyfish_toggle = gr.Checkbox(label="Enable TinyFish Telemetry", value=False)
-        emem_toggle = gr.Checkbox(label="Enable emem Protocol", value=False)
-        
-    gr.ChatInterface(
-        fn=phoebus_respond,
-        additional_inputs=[tinyfish_toggle, emem_toggle],
-        chatbot=gr.Chatbot(height=500),
-        textbox=gr.Textbox(placeholder="Ask Phoebus anything... (tag @emem for signed Earth-observation facts)", container=False, scale=7),
-    )
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    gen_kwargs = dict(**inputs, streamer=streamer, max_new_tokens=350, temperature=0.6, do_sample=True)
+
+    thread = Thread(target=model.generate, kwargs=gen_kwargs)
+    thread.start()
+
+    # 4. Stream response tokens
+    accumulated_text = ""
+    for new_text in streamer:
+        accumulated_text += new_text
+        yield json.dumps({
+            "type": "chunk",
+            "text": accumulated_text,
+            "steps": executed_steps
+        })
+
+    thread.join()
+
+with gr.Blocks() as demo:
+    gr.ChatInterface(fn=phoebus_respond)
 
 if __name__ == "__main__":
     demo.launch()