Update app.py

app.py

@@ -1,52 +1,52 @@
-import os
-import time
-import json
-import numpy as np
-import logging
+import os, time, json, numpy as np, logging
 from typing import List
 from huggingface_hub import HfApi, hf_hub_download, list_repo_files
+from sentence_transformers import SentenceTransformer
 from openai import OpenAI
 
-# ---------------------------
 # Logging setup
-# ---------------------------
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("ohamlab_agent")
 
 # ---------------------------
-# Config
+# Environment / Config
 # ---------------------------
-HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("OPENAI_API_KEY") or os.environ.get("HUGGINGFACE_TOKEN")
+HF_TOKEN = (
+    os.environ.get("HF_TOKEN")
+    or os.environ.get("OPENAI_API_KEY")
+    or os.environ.get("HUGGINGFACE_TOKEN")
+)
 if not HF_TOKEN:
-    logger.critical("Missing HF_TOKEN / OPENAI_API_KEY / HUGGINGFACE_TOKEN environment variable.")
-    raise RuntimeError("ERROR: set env var HF_TOKEN or OPENAI_API_KEY with your Hugging Face / Router token.")
+    raise RuntimeError("Missing HF_TOKEN / OPENAI_API_KEY / HUGGINGFACE_TOKEN.")
 
-MODEL_ID = "openai/gpt-oss-20b"  # chat model via HF router
-EMBED_MODEL = "text-embedding-3-small"
+CHAT_MODEL_ID = "openai/gpt-oss-20b"   # via Hugging Face router
+EMBED_MODEL_ID = "sentence-transformers/all-MiniLM-L6-v2"
 
 HF_REPO = "rahul7star/OhamLab-LLM"
-HF_REPO_DIR = "./hf_capsules"  # local cache
+HF_REPO_DIR = "./hf_capsules"
 os.makedirs(HF_REPO_DIR, exist_ok=True)
 
 # ---------------------------
-# Client (OpenAI router via HF)
+# Clients
 # ---------------------------
 try:
     client = OpenAI(base_url="https://router.huggingface.co/v1", api_key=HF_TOKEN)
-    logger.info("✅ OpenAI client initialized via HF router.")
+    logger.info("✅ OpenAI client via Hugging Face router initialized.")
 except Exception as e:
-    logger.exception("❌ Failed initializing OpenAI client: %s", e)
+    logger.exception("Failed initializing chat client.")
     raise
 
+embedder = SentenceTransformer(EMBED_MODEL_ID)
+logger.info(f"✅ Loaded local embedding model: {EMBED_MODEL_ID}")
+
 # ---------------------------
-# Knowledge Loader
+# Load Markdown Knowledge
 # ---------------------------
 def load_markdown_files(repo_id: str, local_dir: str) -> List[str]:
-    """Downloads all .md files from Hugging Face repo."""
     api = HfApi(token=HF_TOKEN)
     files = list_repo_files(repo_id, repo_type="model", token=HF_TOKEN)
     md_files = [f for f in files if f.endswith(".md")]
-    logger.info(f"📘 Found {len(md_files)} markdown files in {repo_id}")
+    logger.info(f"📘 Found {len(md_files)} markdown files.")
 
     chunks = []
     for f in md_files:
@@ -54,7 +54,6 @@ def load_markdown_files(repo_id: str, local_dir: str) -> List[str]:
             path = hf_hub_download(repo_id=repo_id, filename=f, local_dir=local_dir, token=HF_TOKEN)
             with open(path, "r", encoding="utf-8") as fh:
                 content = fh.read()
-            # Split into 400–600 char segments
             buf = ""
             for line in content.splitlines():
                 buf += line.strip() + " "
@@ -64,83 +63,55 @@ def load_markdown_files(repo_id: str, local_dir: str) -> List[str]:
             if buf:
                 chunks.append(buf.strip())
         except Exception as e:
-            logger.warning(f"⚠️ Failed loading {f}: {e}")
-    logger.info(f"✅ Loaded {len(chunks)} knowledge chunks.")
+            logger.warning(f"⚠️ Failed to read {f}: {e}")
+    logger.info(f"✅ Loaded {len(chunks)} text chunks.")
     return chunks
 
-def get_embeddings(texts: List[str]) -> np.ndarray:
-    """Batch embed text list."""
-    if not texts:
-        return np.zeros((1, 1536))
-    res = client.embeddings.create(model=EMBED_MODEL, input=texts)
-    return np.array([r.embedding for r in res.data])
-
-# ---------------------------
-# Load knowledge and embeddings
-# ---------------------------
-logger.info("🔍 Loading OhamLab knowledge base...")
 KNOWLEDGE_CHUNKS = load_markdown_files(HF_REPO, HF_REPO_DIR)
-logger.info("📊 Generating embeddings...")
-KNOWLEDGE_EMBS = get_embeddings(KNOWLEDGE_CHUNKS)
+logger.info("📊 Creating embeddings...")
+KNOWLEDGE_EMBS = embedder.encode(KNOWLEDGE_CHUNKS, normalize_embeddings=True)
 logger.info(f"🧠 Knowledge base ready ({len(KNOWLEDGE_CHUNKS)} chunks).")
 
 # ---------------------------
-# Retrieval helper
+# Retrieval
 # ---------------------------
 def get_relevant_context(query: str, top_k: int = 3) -> str:
-    if not KNOWLEDGE_CHUNKS or not query:
-        return ""
-    q_emb = get_embeddings([query])[0]
-    sims = np.dot(KNOWLEDGE_EMBS, q_emb) / (
-        np.linalg.norm(KNOWLEDGE_EMBS, axis=1) * np.linalg.norm(q_emb)
-    )
+    q_emb = embedder.encode([query], normalize_embeddings=True)[0]
+    sims = np.dot(KNOWLEDGE_EMBS, q_emb)
     top_idx = np.argsort(sims)[-top_k:][::-1]
     return "\n\n".join(KNOWLEDGE_CHUNKS[i] for i in top_idx)
 
 # ---------------------------
-# Chat Logic
+# Chat
 # ---------------------------
 SYSTEM_PROMPT = (
     "You are OhamLab AI — factual, concise, and context-aware.\n"
-    "Use OhamLab Markdown knowledge if relevant.\n"
-    "Never invent information; be clear and professional."
+    "If applicable, use knowledge from OhamLab Markdown corpus."
 )
 
 def chat(query: str, history: List[dict]) -> str:
     context = get_relevant_context(query)
-    user_input = f"{query}\n\n[Context]\n{context[:1500]}" if context else query
-
+    user_input = f"{query}\n\n[Context]\n{context[:1200]}" if context else query
     msgs = history + [{"role": "user", "content": user_input}]
     try:
         resp = client.chat.completions.create(
-            model=MODEL_ID,
+            model=CHAT_MODEL_ID,
             messages=msgs,
-            temperature=0.5,
-            max_tokens=800,
+            temperature=0.6,
+            max_tokens=700,
         )
         return resp.choices[0].message.content.strip()
     except Exception as e:
         logger.error(f"Chat error: {e}")
-        return "I’m having trouble processing that. Please try again."
+        return "There was a problem generating the response."
 
-# ---------------------------
-# Example usage
-# ---------------------------
 if __name__ == "__main__":
-    logger.info("🚀 Starting OhamLab AI — Knowledge Mode")
-    history = [{"role": "system", "content": SYSTEM_PROMPT}]
+    logger.info("🚀 OhamLab AI — Knowledge Chat Ready")
+    hist = [{"role": "system", "content": SYSTEM_PROMPT}]
     while True:
-        try:
-            q = input("\n🧠 Ask OhamLab → ").strip()
-            if not q:
-                continue
-            if q.lower() in ("exit", "quit"):
-                break
-            ans = chat(q, history)
-            print("\n💬", ans)
-            history.append({"role": "user", "content": q})
-            history.append({"role": "assistant", "content": ans})
-        except KeyboardInterrupt:
+        q = input("\n💬 Ask → ").strip()
+        if q.lower() in ["exit", "quit"]:
             break
-        except Exception as e:
-            logger.exception(f"Main loop error: {e}")
+        ans = chat(q, hist)
+        print("\n🤖", ans)
+        hist.extend([{"role": "user", "content": q}, {"role": "assistant", "content": ans}])