Update app.py

app.py

@@ -1,5 +1,5 @@
 """
-OhamLab — AI Intelligence
+OhamLab — AI Intelligence 
 Loads knowledge from rahul7star/OhamLab-LLM markdown corpus, caches embeddings,
 and provides retrieval-augmented chat through Hugging Face router.
 """
@@ -15,9 +15,9 @@ import gradio as gr
 from openai import OpenAI
 from huggingface_hub import HfApi, hf_hub_download, list_repo_files
 
-# =========================================================
+# ---------------------------
 # 1. Configuration
-# =========================================================
+# ---------------------------
 HF_TOKEN = (
     os.environ.get("HF_TOKEN")
     or os.environ.get("OPENAI_API_KEY")
@@ -26,32 +26,28 @@ HF_TOKEN = (
 if not HF_TOKEN:
     raise RuntimeError("❌ Missing HF_TOKEN / OPENAI_API_KEY / HUGGINGFACE_TOKEN environment variable.")
 
-MODEL_ID = "openai/gpt-4o-mini"  # Chat model
-EMBED_MODEL = "text-embedding-3-small"  # Embedding model
-HF_REPO = "rahul7star/OhamLab-LLM"
-CACHE_PATH = "/tmp/ohamlab_emb_cache.json"
+MODEL_ID = "openai/gpt-oss-20b"               # Chat model (via HF router)
+EMBED_MODEL = "text-embedding-3-small"        # Embedding model
+HF_REPO = "rahul7star/OhamLab-LLM"            # Knowledge repo
+CACHE_PATH = "/tmp/ohamlab_emb_cache.json"    # Cache file
 
+# Client
 client = OpenAI(base_url="https://router.huggingface.co/v1", api_key=HF_TOKEN)
 api = HfApi(token=HF_TOKEN)
 
-
-# =========================================================
-# 2. Load Markdown Corpus
-# =========================================================
+# ---------------------------
+# 2. Load and Chunk Markdown Files
+# ---------------------------
 def load_ohamlab_knowledge():
-    """Load all markdown files and split into ~500-character chunks."""
-    print(f"📂 Loading markdown files from {HF_REPO}...")
+    """Loads all .md files from Hugging Face repo and splits into ~500-char chunks."""
     files = list_repo_files(HF_REPO, repo_type="model", token=HF_TOKEN)
     md_files = [f for f in files if f.endswith(".md")]
     chunks = []
-
     for f in md_files:
         try:
             path = hf_hub_download(HF_REPO, filename=f, token=HF_TOKEN)
             with open(path, "r", encoding="utf-8") as fh:
-                content = fh.read().strip()
-
-            # clean + split into ~500 chars
+                content = fh.read()
             buf = ""
             for line in content.splitlines():
                 buf += line.strip() + " "
@@ -60,71 +56,54 @@ def load_ohamlab_knowledge():
                     buf = ""
             if buf:
                 chunks.append({"file": f, "text": buf.strip()})
-            print(f"✅ Loaded {f} ({len(content)} chars)")
         except Exception as e:
             print(f"⚠️ Failed to load {f}: {e}")
-
-    print(f"📘 Total chunks: {len(chunks)}")
     return chunks
 
-
-# =========================================================
-# 3. Embedding Generation with Retry + Cache
-# =========================================================
-def create_embeddings_with_retry(texts, retries=3, delay=2):
-    """Create embeddings with retry logic."""
-    for attempt in range(1, retries + 1):
-        try:
-            response = client.embeddings.create(model=EMBED_MODEL, input=texts)
-            return [d.embedding for d in response.data]
-        except Exception as e:
-            print(f"⚠️ Embedding attempt {attempt} failed: {e}")
-            if attempt == retries:
-                raise RuntimeError("❌ Failed to generate embeddings after retries.")
-            time.sleep(delay)
-
-
+# ---------------------------
+# 3. Generate or Load Embeddings (with Cache)
+# ---------------------------
 def get_embeddings_with_cache():
-    """Generate or load cached embeddings."""
+    """Generate or load cached embeddings for OhamLab context."""
     if os.path.exists(CACHE_PATH):
         try:
-            with open(CACHE_PATH, "r", encoding="utf-8") as f:
+            with open(CACHE_PATH, "r") as f:
                 cache = json.load(f)
             texts = [c["text"] for c in cache]
             embs = np.array([c["embedding"] for c in cache])
-            print(f"✅ Loaded cached embeddings ({len(embs)} chunks)")
+            print(f"✅ Loaded cached embeddings from {CACHE_PATH} ({len(embs)} chunks)")
             return texts, embs
-        except Exception as e:
-            print(f"⚠️ Cache load failed: {e}")
+        except Exception:
+            print("⚠️ Cache corrupted, regenerating embeddings...")
 
-    # Load and embed new
     chunks = load_ohamlab_knowledge()
     texts = [c["text"] for c in chunks]
-    print(f"📘 Generating embeddings for {len(texts)} chunks...")
-
+    print(f"📘 Generating embeddings for {len(texts)} OhamLab chunks...")
     all_embs = []
     for i in range(0, len(texts), 50):
         batch = texts[i:i + 50]
-        embs = create_embeddings_with_retry(batch)
-        all_embs.extend(embs)
+        try:
+            res = client.embeddings.create(model=EMBED_MODEL, input=batch)
+            embs = [d.embedding for d in res.data]
+            all_embs.extend(embs)
+        except Exception as e:
+            print(f"⚠️ Embedding batch failed ({i}): {e}")
+            all_embs.extend([[0.0] * 1536] * len(batch))  # fallback
         time.sleep(0.5)
 
     data = [{"text": t, "embedding": e} for t, e in zip(texts, all_embs)]
-    with open(CACHE_PATH, "w", encoding="utf-8") as f:
+    with open(CACHE_PATH, "w") as f:
         json.dump(data, f)
     print(f"💾 Cached embeddings to {CACHE_PATH}")
-
     return texts, np.array(all_embs)
 
-
 OHAMLAB_TEXTS, OHAMLAB_EMBS = get_embeddings_with_cache()
 
-
-# =========================================================
-# 4. Retrieval
-# =========================================================
+# ---------------------------
+# 4. Semantic Retrieval
+# ---------------------------
 def retrieve_knowledge(query, top_k=3):
-    """Return top-k most relevant text snippets."""
+    """Retrieve top-k most relevant text snippets."""
     try:
         q_emb = client.embeddings.create(model=EMBED_MODEL, input=[query]).data[0].embedding
         sims = np.dot(OHAMLAB_EMBS, q_emb) / (
@@ -136,94 +115,202 @@ def retrieve_knowledge(query, top_k=3):
         print(f"⚠️ Retrieval error: {e}")
         return ""
 
-
-# =========================================================
-# 5. Prompt Construction
-# =========================================================
-def build_system_prompt(context, mode="chat"):
+# ---------------------------
+# 5. System Prompt with Context Injection
+# ---------------------------
+def build_system_prompt(context: str, mode: str = "chat") -> str:
     return textwrap.dedent(f"""
-    You are OhamLab — an AI Research Assistant for OhamLab.
+    You are OhamLab — AI Intelligence Software
 
-    Rules:
-    - Use precise, factual, and confident tone.
-    - If relevant, use retrieved OhamLab context.
-    - If unknown, politely say "I could not find an answer in my knowledge base."
+    Guidelines:
+    - Always answer with clarity, scientific accuracy, and concise insight.
+    - Incorporate OhamLab research knowledge when relevant.
+    - Avoid code unless explicitly requested.
+    - Be confident but label speculation clearly.
     - Mode: {mode.upper()}
 
-    --- OhamLab Context ---
+    --- OhamLab Context (Retrieved Snippets) ---
     {context[:1800]}
     --- End Context ---
     """).strip()
 
-
-# =========================================================
-# 6. Chat Generation
-# =========================================================
+# ---------------------------
+# 6. Model Call
+# ---------------------------
 def generate_response(user_input, history, mode="chat"):
     context = retrieve_knowledge(user_input)
     sys_prompt = build_system_prompt(context, mode)
-
     messages = [{"role": "system", "content": sys_prompt}] + history + [
         {"role": "user", "content": user_input}
     ]
-
     try:
         resp = client.chat.completions.create(
             model=MODEL_ID,
             messages=messages,
             temperature=0.7,
-            max_tokens=800,
+            max_tokens=1200,
         )
         return resp.choices[0].message.content.strip()
     except Exception as e:
         print(f"⚠️ Model call failed: {e}")
-        return "⚠️ OhamLab encountered a temporary issue."
+        return "⚠️ OahmLab encountered a temporary issue generating your response."
+
+# ---------------------------
+# 7. Gradio Chat UI
+# ---------------------------
+import traceback
+import gradio as gr
+
+# ---------------------------
+# Chat Logic
+# ---------------------------
+
 
 
-# =========================================================
-# 7. Gradio UI
-# =========================================================
 def chat_with_model(user_message, chat_history):
+    """
+    Maintains full conversational context and returns updated chat history.
+    The assistant speaks as 'OhamLab'.
+    """
     if not user_message:
         return chat_history, ""
 
-    chat_history = chat_history or []
-    history = [{"role": m["role"], "content": m["content"]} for m in chat_history]
+    if chat_history is None:
+        chat_history = []
+
+    # Convert Gradio message list (dict-based) to usable context
+    history = [
+        {"role": m["role"], "content": m["content"]}
+        for m in chat_history
+        if isinstance(m, dict) and "role" in m
+    ]
+
+    # Append current user message
     history.append({"role": "user", "content": user_message})
 
     try:
         bot_reply = generate_response(user_message, history)
-    except Exception:
-        bot_reply = f"⚠️ Internal error:\n\n{traceback.format_exc()}"
+    except Exception as e:
+        tb = traceback.format_exc()
+        bot_reply = f"⚠️ OhamLab encountered an error:\n\n{e}\n\n{tb}"
 
+    # Add OhamLab's response as assistant role
     history.append({"role": "assistant", "content": bot_reply})
+
     return history, ""
 
 
 def reset_chat():
+    """Resets the chat session."""
     return []
 
 
+# ---------------------------
+# Gradio Chat UI
+# ---------------------------
+
 def build_ui():
-    with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo")) as demo:
+    with gr.Blocks(
+        theme=gr.themes.Soft(primary_hue="indigo"),
+        css="""
+            /* --- Hide share/delete icons --- */
+            #ohamlab .wrap.svelte-1lcyrj3 > div > div > button {
+                display: none !important;
+            }
+            [data-testid="share-btn"],
+            [data-testid="delete-btn"],
+            .message-controls,
+            .message-actions {
+                display: none !important;
+                visibility: hidden !important;
+            }
+
+            /* --- User (Right) Message Bubble --- */
+            #ohamlab .message.user {
+                background-color: #4f46e5 !important;
+                color: white !important;
+                border-radius: 14px !important;
+                align-self: flex-end !important;
+                text-align: right !important;
+                margin-left: 25%;
+            }
+
+            /* --- OhamLab (Left) Message Bubble --- */
+            #ohamlab .message.assistant {
+                background-color: #f8f9fa !important;
+                color: #111 !important;
+                border-radius: 14px !important;
+                align-self: flex-start !important;
+                text-align: left !important;
+                margin-right: 25%;
+            }
+
+            
+            #ohamlab .chatbot .wrap.svelte-1lcyrj3 > div > div > button {
+                   display: none !important; /* hide share/delete icons */
+               }
+         
+            /* --- Overall Container --- */
+            .gradio-container {
+                max-width: 900px !important;
+                margin: auto;
+                padding-top: .5rem;
+            }
+            textarea {
+                resize: none !important;
+                border-radius: 12px !important;
+                border: 1px solid #d1d5db !important;
+                box-shadow: 0 1px 3px rgba(0,0,0,0.08);
+            }
+            button.primary {
+                background-color: #4f46e5 !important;
+                color: white !important;
+                border-radius: 10px !important;
+                padding: 0.6rem 1.4rem !important;
+                font-weight: 600;
+                transition: all 0.2s ease-in-out;
+            }
+            button.primary:hover {
+                background-color: #4338ca !important;
+            }
+            button.secondary {
+                background-color: #f3f4f6 !important;
+                border-radius: 10px !important;
+                color: #374151 !important;
+                font-weight: 500;
+                transition: all 0.2s ease-in-out;
+            }
+            button.secondary:hover {
+                background-color: #e5e7eb !important;
+            }
+        """,
+    ) as demo:
+        
+        # Chatbot area
         chatbot = gr.Chatbot(
             label="💠 OhamLab Conversation",
             height=520,
             elem_id="ohamlab",
             type="messages",
+            avatar_images=[None, None],
         )
 
-        msg = gr.Textbox(
-            placeholder="Ask OhamLab anything...",
-            lines=3,
-            show_label=False,
-            container=False,
-        )
-
+        # Input box (full width)
         with gr.Row():
-            send = gr.Button("Send", variant="primary")
-            clear = gr.Button("Clear", variant="secondary")
-
+            msg = gr.Textbox(
+                placeholder="Ask OhamLab anything ..",
+                lines=3,
+                show_label=False,
+                scale=12,
+                container=False,
+            )
+
+        # Buttons (Send + Clear)
+        with gr.Row(equal_height=True, variant="compact"):
+            send = gr.Button("Send", variant="primary", elem_classes=["primary"])
+            clear = gr.Button("Clear", variant="secondary", elem_classes=["secondary"])
+
+        # Wiring
         send.click(chat_with_model, inputs=[msg, chatbot], outputs=[chatbot, msg])
         msg.submit(chat_with_model, inputs=[msg, chatbot], outputs=[chatbot, msg])
         clear.click(reset_chat, outputs=chatbot)
@@ -231,9 +318,9 @@ def build_ui():
     return demo
 
 
-# =========================================================
-# 8. Entrypoint
-# =========================================================
+# ---------------------------
+# Entrypoint
+# ---------------------------
 if __name__ == "__main__":
     print("🚀 Starting OhamLab Assistant...")
     demo = build_ui()