Update app.py

app.py

@@ -1,111 +1,94 @@
-import os, re, time, json, textwrap, traceback, numpy as np
+import os
+import time
+import json
+import numpy as np
+import logging
 from typing import List
-import gradio as gr
+from huggingface_hub import HfApi, hf_hub_download, list_repo_files
 from openai import OpenAI
-from huggingface_hub import list_repo_files, hf_hub_download
 
-# ==========================================================
-# CONFIG
-# ==========================================================
-REPO_ID = "rahul7star/OhamLab-LLM"
-API_KEY_ENV = "OPENAI_API_KEY"
-HF_TOKEN = os.getenv(API_KEY_ENV)
+# ---------------------------
+# Logging setup
+# ---------------------------
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("ohamlab_agent")
+
+# ---------------------------
+# Config
+# ---------------------------
+HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("OPENAI_API_KEY") or os.environ.get("HUGGINGFACE_TOKEN")
 if not HF_TOKEN:
-    raise RuntimeError(f"Missing {API_KEY_ENV}")
+    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.")
 
 MODEL_ID = "openai/gpt-oss-20b"  # chat model via HF router
 EMBED_MODEL = "text-embedding-3-small"
-CACHE_FILE = "/tmp/ohamlab_emb_cache.json"
-
-client = OpenAI(api_key=HF_TOKEN)
-LOG_FILE = "ohamlab_chat.log"
-
-def log(msg: str):
-    ts = time.strftime("%Y-%m-%d %H:%M:%S")
-    line = f"[{ts}] {msg}"
-    print(line)
-    try:
-        with open(LOG_FILE, "a", encoding="utf-8") as f:
-            f.write(line + "\n")
-    except Exception:
-        pass
 
-# ==========================================================
-# KNOWLEDGE SCANNER + EMBEDDINGS
-# ==========================================================
-def load_all_md_from_repo(repo_id: str) -> List[str]:
-    """Scan all .md files in the repo and return concatenated chunks."""
-    try:
-        files = list_repo_files(repo_id=repo_id, repo_type="model", token=HF_TOKEN)
-        md_files = [f for f in files if f.lower().endswith(".md")]
-        if not md_files:
-            log("⚠️ No markdown files found.")
-            return []
-
-        chunks = []
-        for fname in md_files:
-            try:
-                local_path = hf_hub_download(repo_id, filename=fname, token=HF_TOKEN)
-                with open(local_path, "r", encoding="utf-8") as f:
-                    text = f.read()
-                text = re.sub(r"<[^>]+>", "", text)  # strip HTML
-                # Split into ~500 char chunks
-                buf = ""
-                for line in text.splitlines():
-                    buf += line.strip() + " "
-                    if len(buf) > 500:
-                        chunks.append(buf.strip())
-                        buf = ""
-                if buf:
+HF_REPO = "rahul7star/OhamLab-LLM"
+HF_REPO_DIR = "./hf_capsules"  # local cache
+os.makedirs(HF_REPO_DIR, exist_ok=True)
+
+# ---------------------------
+# Client (OpenAI router via HF)
+# ---------------------------
+try:
+    client = OpenAI(base_url="https://router.huggingface.co/v1", api_key=HF_TOKEN)
+    logger.info("✅ OpenAI client initialized via HF router.")
+except Exception as e:
+    logger.exception("❌ Failed initializing OpenAI client: %s", e)
+    raise
+
+# ---------------------------
+# Knowledge Loader
+# ---------------------------
+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}")
+
+    chunks = []
+    for f in md_files:
+        try:
+            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() + " "
+                if len(buf) > 500:
                     chunks.append(buf.strip())
-                log(f"Loaded {fname} ({len(text)} chars → {len(chunks)} chunks)")
-            except Exception as e:
-                log(f"⚠️ Failed to load {fname}: {e}")
-        return chunks
-    except Exception as e:
-        log(f"⚠️ Repo scan failed: {e}")
-        return []
+                    buf = ""
+            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.")
+    return chunks
 
 def get_embeddings(texts: List[str]) -> np.ndarray:
+    """Batch embed text list."""
     if not texts:
         return np.zeros((1, 1536))
-    try:
-        res = client.embeddings.create(model=EMBED_MODEL, input=texts)
-        return np.array([r.embedding for r in res.data])
-    except Exception as e:
-        log(f"Embedding error: {e}")
-        return np.zeros((len(texts), 1536))
-
-def load_knowledge_cache() -> tuple[list[str], np.ndarray]:
-    """Load embeddings from cache or regenerate from repo."""
-    if os.path.exists(CACHE_FILE):
-        try:
-            with open(CACHE_FILE, "r", encoding="utf-8") as f:
-                data = json.load(f)
-            chunks = data["chunks"]
-            embs = np.array(data["embs"])
-            log(f"Loaded cached embeddings: {len(chunks)} chunks.")
-            return chunks, embs
-        except Exception:
-            pass
-
-    log("Scanning Markdown files in repo for knowledge base...")
-    chunks = load_all_md_from_repo(REPO_ID)
-    embs = get_embeddings(chunks)
-    try:
-        json.dump({"chunks": chunks, "embs": embs.tolist()}, open(CACHE_FILE, "w"))
-    except Exception:
-        pass
-    log(f"Knowledge base ready: {len(chunks)} chunks.")
-    return chunks, embs
-
-KNOWLEDGE_CHUNKS, KNOWLEDGE_EMBS = load_knowledge_cache()
-
-# ==========================================================
-# RETRIEVAL
-# ==========================================================
-def get_relevant_context(query: str, top_k=3) -> str:
-    if not KNOWLEDGE_CHUNKS or not query.strip():
+    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(f"🧠 Knowledge base ready ({len(KNOWLEDGE_CHUNKS)} chunks).")
+
+# ---------------------------
+# Retrieval helper
+# ---------------------------
+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) / (
@@ -114,85 +97,50 @@ def get_relevant_context(query: str, top_k=3) -> str:
     top_idx = np.argsort(sims)[-top_k:][::-1]
     return "\n\n".join(KNOWLEDGE_CHUNKS[i] for i in top_idx)
 
-# ==========================================================
-# CHAT ENGINE
-# ==========================================================
+# ---------------------------
+# Chat Logic
+# ---------------------------
 SYSTEM_PROMPT = (
-    "You are OhamLab AI — a concise, factual chat assistant.\n"
-    "When relevant, use the OhamLab knowledge base provided in context.\n"
-    "Never show code unless explicitly requested. Keep tone professional and calm."
+    "You are OhamLab AI — factual, concise, and context-aware.\n"
+    "Use OhamLab Markdown knowledge if relevant.\n"
+    "Never invent information; be clear and professional."
 )
 
-history = [{"role": "system", "content": SYSTEM_PROMPT}]
-chat_ui_history = []
-MAX_HISTORY_CHARS = 3000
-
-def trim_msgs(msgs, max_chars=MAX_HISTORY_CHARS):
-    sys = msgs[0]
-    tail, total = [], len(sys["content"])
-    for m in reversed(msgs[1:]):
-        seg = len(m["content"])
-        if total + seg > max_chars: break
-        tail.append(m)
-        total += seg
-    return [sys] + list(reversed(tail))
-
-def chat(user_msg, chat_hist):
-    global history, chat_ui_history
-    user_msg = (user_msg or "").strip()
-    if not user_msg:
-        return chat_ui_history, ""
-
-    context = get_relevant_context(user_msg)
-    if context:
-        user_msg += f"\n\n[Context]\n{context[:1500]}"
-
-    history.append({"role": "user", "content": user_msg})
-    trimmed = trim_msgs(history)
+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
 
+    msgs = history + [{"role": "user", "content": user_input}]
     try:
         resp = client.chat.completions.create(
             model=MODEL_ID,
-            messages=trimmed,
-            max_tokens=600,
-            temperature=0.6,
+            messages=msgs,
+            temperature=0.5,
+            max_tokens=800,
         )
-        reply = resp.choices[0].message.content.strip()
+        return resp.choices[0].message.content.strip()
     except Exception as e:
-        log(f"Chat error: {e}")
-        reply = "I'm experiencing a temporary issue. Please try again shortly."
-
-    history.append({"role": "assistant", "content": reply})
-    chat_ui_history.append((user_msg, reply))
-    history[:] = trim_msgs(history)
-
-    log(f"USER: {user_msg[:60]} | BOT: {reply[:60]}")
-    return chat_ui_history, ""
-
-def reset():
-    global history, chat_ui_history
-    history = [{"role": "system", "content": SYSTEM_PROMPT}]
-    chat_ui_history = []
-    log("Chat reset.")
-    return []
-
-# ==========================================================
-# UI
-# ==========================================================
-def build_ui():
-    with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("### 💬 OhamLab AI — Knowledge Chat (Markdown-Aware)")
-        chatbot = gr.Chatbot(height=350)
-        msg = gr.Textbox(placeholder="Ask anything about OhamLab...", lines=1)
-        send = gr.Button("Send", variant="primary")
-        clear = gr.Button("Clear")
-
-        send.click(chat, inputs=[msg, chatbot], outputs=[chatbot, msg])
-        msg.submit(chat, inputs=[msg, chatbot], outputs=[chatbot, msg])
-        clear.click(reset, None, chatbot)
-    demo.launch(server_name="0.0.0.0", server_port=7860)
-    return demo
+        logger.error(f"Chat error: {e}")
+        return "I’m having trouble processing that. Please try again."
 
+# ---------------------------
+# Example usage
+# ---------------------------
 if __name__ == "__main__":
-    log("Starting OhamLab AI Chat Agent (Markdown Knowledge)...")
-    build_ui()
+    logger.info("🚀 Starting OhamLab AI — Knowledge Mode")
+    history = [{"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:
+            break
+        except Exception as e:
+            logger.exception(f"Main loop error: {e}")