Update app.py

app.py

@@ -1,38 +1,57 @@
 from fastapi import FastAPI
 from fastapi.responses import StreamingResponse, HTMLResponse
 from pydantic import BaseModel
-from transformers import pipeline
+from transformers import pipeline, TextStreamer
 import asyncio
+import httpx
+import time
 import queue
 import threading
-import time
 import random
-import httpx
-
-class ModelInput(BaseModel):
-    prompt: str
-    max_new_tokens: int = 64000
-
-app = FastAPI()
-
-# Your main generation model (DeepSeek)
+import re
+
+# =========================
+# CONFIG
+# =========================
+UPDATE_INTERVAL = 60  # seconds between KG updates
+MAX_KG_SIZE = 50      # limit stored KG nodes to avoid memory bloat
+
+# =========================
+# MODELS
+# =========================
+# Main generator
 generator = pipeline(
     "text-generation",
     model="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
     device="cpu"
 )
 
-# The summarization instruct model
-summarizer = pipeline(
+# Query + summarization model (SmolLM2 instruct)
+query_generator = pipeline(
     "text-generation",
     model="HuggingFaceTB/SmolLM2-360M-Instruct",
-    device="cpu",
-    max_length=512,  # keep summary short
-    do_sample=False
+    device="cpu"
 )
 
+summarizer = query_generator  # same model for now
+
+# =========================
+# KNOWLEDGE GRAPH
+# =========================
 knowledge_graph = {}
 
+# =========================
+# FASTAPI
+# =========================
+app = FastAPI()
+
+class ModelInput(BaseModel):
+    prompt: str
+    max_new_tokens: int = 64000
+
+# =========================
+# UTILS
+# =========================
 async def fetch_ddg_search(query: str):
     url = "https://api.duckduckgo.com/"
     params = {
@@ -44,101 +63,115 @@ async def fetch_ddg_search(query: str):
     }
     async with httpx.AsyncClient() as client:
         resp = await client.get(url, params=params, timeout=15)
-        data = resp.json()
-    return data
+        return resp.json()
 
 def clean_ddg_text(ddg_json):
-    # Take abstract text + top related topic texts concatenated
     abstract = ddg_json.get("AbstractText", "")
     related = ddg_json.get("RelatedTopics", [])
     related_texts = []
     for item in related:
-        if "Text" in item:
+        if isinstance(item, dict) and "Text" in item:
             related_texts.append(item["Text"])
-        elif "Name" in item and "Topics" in item:
+        elif isinstance(item, dict) and "Topics" in item:
             for sub in item["Topics"]:
                 if "Text" in sub:
                     related_texts.append(sub["Text"])
-    combined_text = abstract + " " + " ".join(related_texts)
-    # Simple clean up, trim length to avoid overloading
-    combined_text = combined_text.strip().replace("\n", " ")
-    if len(combined_text) > 1000:
-        combined_text = combined_text[:1000] + "..."
-    return combined_text
+    combined = (abstract + " " + " ".join(related_texts)).strip()
+    combined = re.sub(r"\s+", " ", combined)
+    if len(combined) > 1000:
+        combined = combined[:1000] + "..."
+    return combined
 
-def summarize_text(text: str):
-    # Run the instruct model to summarize/clean the text
-    prompt = f"Summarize this information concisely:\n{text}\nSummary:"
-    output = summarizer(prompt, max_length=256, do_sample=False)
-    return output[0]["generated_text"].strip()
-
-async def generate_random_query():
-    """Use your main model to invent a search query."""
+def generate_dynamic_query():
     prompt = (
-        "Generate one short, interesting, real-world question "
-        "about technology, startups, AI, or science that someone might Google. "
-        "Output only the query, nothing else."
+        "Generate a short, specific search query about technology, startups, AI, or science. "
+        "Be creative, realistic, and output only the query with no extra words."
     )
-    output = generator(prompt, max_new_tokens=20, do_sample=True, temperature=0.9, truncation=True)
-    query = output[0]["generated_text"].strip()
-    # Safety: strip weird trailing punctuation
-    return query.split("\n")[0].strip("?!. ")
+    output = query_generator(
+        prompt,
+        max_new_tokens=32,
+        truncation=True,
+        do_sample=True,
+        temperature=0.9
+    )
+    query = output[0]["generated_text"].strip().split("\n")[0]
+    return query
 
+def summarize_text(text: str):
+    prompt = f"Summarize this concisely:\n{text}\nSummary:"
+    output = summarizer(
+        prompt,
+        max_new_tokens=256,
+        truncation=True,
+        do_sample=False
+    )
+    return output[0]["generated_text"].strip()
+
+def inject_relevant_kg(prompt: str):
+    """Find relevant KG entries and inject into prompt."""
+    if not knowledge_graph:
+        return prompt
+    best_match = None
+    for key, node in knowledge_graph.items():
+        if any(word.lower() in prompt.lower() for word in key.split()):
+            best_match = node
+            break
+    if best_match:
+        return f"{prompt}\n\nRelevant knowledge from memory:\n{best_match['summary']}"
+    return prompt
+
+# =========================
+# BACKGROUND TASK
+# =========================
 async def update_knowledge_graph_periodically():
     while True:
         try:
-            query = await generate_random_query()
-            print(f"[KG Updater] Searching DuckDuckGo for query: {query}")
-
+            query = generate_dynamic_query()
+            print(f"[KG Updater] Searching DDG for query: {query}")
             ddg_data = await fetch_ddg_search(query)
             cleaned = clean_ddg_text(ddg_data)
 
             if not cleaned or len(cleaned) < 50:
-                print("[KG Updater] Too little info found, retrying with new query...")
-                continue  # Skip saving, go to next loop
-
-            # Summarize with instruct model
-            loop = asyncio.get_event_loop()
-            summary = await loop.run_in_executor(
-                None,
-                lambda: summarizer(
-                    f"Summarize this info concisely:\n{cleaned}",
-                    max_length=256,
-                    truncation=True,
-                    do_sample=False
-                )[0]["generated_text"].strip()
-            )
-
-            knowledge_graph[query] = {
-                "raw_text": cleaned,
-                "summary": summary,
-                "timestamp": time.time()
-            }
-            print(f"[KG Updater] Knowledge graph updated for query: {query}")
+                print("[KG Updater] Too little info found, retrying next cycle...")
+            else:
+                summary = summarize_text(cleaned)
+                knowledge_graph[query] = {
+                    "raw_text": cleaned,
+                    "summary": summary,
+                    "timestamp": time.time()
+                }
+                if len(knowledge_graph) > MAX_KG_SIZE:
+                    # remove oldest
+                    oldest_key = min(knowledge_graph, key=lambda k: knowledge_graph[k]['timestamp'])
+                    del knowledge_graph[oldest_key]
+                print(f"[KG Updater] Knowledge graph updated for query: {query}")
 
         except Exception as e:
             print(f"[KG Updater] Error: {e}")
 
-        await asyncio.sleep(60)  # Run forever, every 60s
-
+        await asyncio.sleep(UPDATE_INTERVAL)
 
 @app.on_event("startup")
 async def startup_event():
     asyncio.create_task(update_knowledge_graph_periodically())
 
-# Manual streaming endpoint (kept as-is)
+# =========================
+# STREAMING ENDPOINT
+# =========================
 @app.post("/generate/stream")
 async def generate_stream(input: ModelInput):
     q = queue.Queue()
+
     def run_generation():
         try:
-            streamer = pipeline("text-generation", model="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B", device="cpu").tokenizer
             streamer = TextStreamer(generator.tokenizer, skip_prompt=True)
             def enqueue_token(token):
                 q.put(token)
             streamer.put = enqueue_token
+
+            enriched_prompt = inject_relevant_kg(input.prompt)
             generator(
-                input.prompt,
+                enriched_prompt,
                 max_new_tokens=input.max_new_tokens,
                 do_sample=False,
                 streamer=streamer
@@ -147,6 +180,7 @@ async def generate_stream(input: ModelInput):
             q.put(f"[ERROR] {e}")
         finally:
             q.put(None)
+
     thread = threading.Thread(target=run_generation)
     thread.start()
 
@@ -157,23 +191,28 @@ async def generate_stream(input: ModelInput):
             if token is None:
                 break
             yield token
+
     return StreamingResponse(event_generator(), media_type="text/plain")
 
-# Endpoint to get KG
+# =========================
+# VIEW KG
+# =========================
 @app.get("/knowledge")
 async def get_knowledge():
     return knowledge_graph
 
-# Basic client page to test streaming
+# =========================
+# TEST CLIENT PAGE
+# =========================
 @app.get("/", response_class=HTMLResponse)
 async def root():
     return """
     <!DOCTYPE html>
     <html>
-    <head><title>Streaming Text Generation Client</title></head>
+    <head><title>Xylaria Cognitive Worker</title></head>
     <body>
-      <h2>Streaming Text Generation Demo</h2>
-      <textarea id="prompt" rows="4" cols="60">Write me a poem about tech startup struggles</textarea><br/>
+      <h2>Xylaria Cognitive Worker</h2>
+      <textarea id="prompt" rows="4" cols="60">Explain how AI startups secure funding</textarea><br/>
       <button onclick="startStreaming()">Generate</button>
       <pre id="output" style="white-space: pre-wrap; background:#eee; padding:10px; border-radius:5px; max-height:400px; overflow:auto;"></pre>
       <h3>Knowledge Graph</h3>
@@ -199,15 +238,13 @@ async def root():
               output.scrollTop = output.scrollHeight;
           }
       }
-
       async function fetchKG() {
           const kgPre = document.getElementById("kg");
           const res = await fetch("/knowledge");
           const data = await res.json();
           kgPre.textContent = JSON.stringify(data, null, 2);
       }
-
-      setInterval(fetchKG, 10000);  // update KG display every 10s
+      setInterval(fetchKG, 10000);
       window.onload = fetchKG;
       </script>
     </body>