Update app.py

app.py

@@ -1,10 +1,11 @@
 """
-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.
 """
 
 import os
+import re
 import json
 import time
 import textwrap
@@ -25,11 +26,12 @@ HF_TOKEN = (
 if not HF_TOKEN:
     raise RuntimeError("❌ Missing HF_TOKEN / OPENAI_API_KEY / HUGGINGFACE_TOKEN environment variable.")
 
-MODEL_ID = "openai/gpt-oss-20b"                # Chat model (via HF router)
-EMBED_MODEL = "BAAI/bge-small-en-v1.5"         # Router-compatible embedding model
-HF_REPO = "rahul7star/OhamLab-LLM"             # Knowledge repo
-CACHE_PATH = "/tmp/ohamlab_emb_cache.json"     # Cache file
+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)
 
@@ -37,7 +39,7 @@ api = HfApi(token=HF_TOKEN)
 # 2. Load and Chunk Markdown Files
 # ---------------------------
 def load_ohamlab_knowledge():
-    """Load all .md files from Hugging Face repo and split into ~500-char chunks."""
+    """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 = []
@@ -56,25 +58,13 @@ def load_ohamlab_knowledge():
                 chunks.append({"file": f, "text": buf.strip()})
         except Exception as e:
             print(f"⚠️ Failed to load {f}: {e}")
-    print(f"📚 Loaded {len(chunks)} text chunks from {len(md_files)} markdown files.")
     return chunks
 
 # ---------------------------
-# 3. Embeddings (with Cache & Retry)
+# 3. Generate or Load Embeddings (with Cache)
 # ---------------------------
-def create_embeddings_with_retry(texts, retries=3, delay=2):
-    """Generate embeddings with retries on failure."""
-    for attempt in range(retries):
-        try:
-            res = client.embeddings.create(model=EMBED_MODEL, input=texts)
-            return [d.embedding for d in res.data]
-        except Exception as e:
-            print(f"⚠️ Embedding attempt {attempt+1} failed: {e}")
-            time.sleep(delay)
-    raise RuntimeError("❌ Failed to generate embeddings after retries.")
-
 def get_embeddings_with_cache():
-    """Generate or load cached embeddings for OhamLab knowledge."""
+    """Generate or load cached embeddings for OhamLab context."""
     if os.path.exists(CACHE_PATH):
         try:
             with open(CACHE_PATH, "r") as f:
@@ -89,12 +79,16 @@ def get_embeddings_with_cache():
     chunks = load_ohamlab_knowledge()
     texts = [c["text"] for c in 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)]
@@ -109,9 +103,9 @@ OHAMLAB_TEXTS, OHAMLAB_EMBS = get_embeddings_with_cache()
 # 4. Semantic Retrieval
 # ---------------------------
 def retrieve_knowledge(query, top_k=3):
-    """Retrieve top-k most relevant text snippets from markdown knowledge bank."""
+    """Retrieve top-k most relevant text snippets."""
     try:
-        q_emb = create_embeddings_with_retry([query])[0]
+        q_emb = client.embeddings.create(model=EMBED_MODEL, input=[query]).data[0].embedding
         sims = np.dot(OHAMLAB_EMBS, q_emb) / (
             np.linalg.norm(OHAMLAB_EMBS, axis=1) * np.linalg.norm(q_emb)
         )
@@ -122,17 +116,17 @@ def retrieve_knowledge(query, top_k=3):
         return ""
 
 # ---------------------------
-# 5. System Prompt
+# 5. System Prompt with Context Injection
 # ---------------------------
 def build_system_prompt(context: str, mode: str = "chat") -> str:
     return textwrap.dedent(f"""
-    You are OhamLab — AI Intelligence Software.
+    You are OhamLab — AI Intelligence Software
 
     Guidelines:
-    - Only answer using information retrieved from the OhamLab knowledge bank (.md files).
-    - Do NOT hallucinate. If the answer is not found, respond: "I could not find an answer in my knowledge base."
-    - Always answer with clarity, precision, and factual grounding.
+    - 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 (Retrieved Snippets) ---
@@ -141,7 +135,7 @@ def build_system_prompt(context: str, mode: str = "chat") -> str:
     """).strip()
 
 # ---------------------------
-# 6. Response Generation
+# 6. Model Call
 # ---------------------------
 def generate_response(user_input, history, mode="chat"):
     context = retrieve_knowledge(user_input)
@@ -153,28 +147,45 @@ def generate_response(user_input, history, mode="chat"):
         resp = client.chat.completions.create(
             model=MODEL_ID,
             messages=messages,
-            temperature=0.0,  # deterministic for knowledge-only answers
+            temperature=0.7,
             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 generating your response."
+        return "⚠️ OahmLab encountered a temporary issue generating your response."
 
 # ---------------------------
-# 7. Chat Logic
+# 7. Gradio Chat UI
 # ---------------------------
+import traceback
+import gradio as gr
+
+# ---------------------------
+# Chat Logic
+# ---------------------------
+
+
+
 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, ""
+
     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:
@@ -183,33 +194,99 @@ def chat_with_model(user_message, chat_history):
         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 []
 
+
 # ---------------------------
-# 8. Gradio UI
+# Gradio Chat UI
 # ---------------------------
+
 def build_ui():
     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; }
-            #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 .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%; }
-            .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; }
-        """
+            .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,
@@ -218,24 +295,29 @@ def build_ui():
             avatar_images=[None, None],
         )
 
+        # Input box (full width)
         with gr.Row():
             msg = gr.Textbox(
-                placeholder="Ask OhamLab anything ...",
+                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")
-            clear = gr.Button("Clear", variant="secondary")
+            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)
+
     return demo
 
+
 # ---------------------------
 # Entrypoint
 # ---------------------------
@@ -243,3 +325,4 @@ if __name__ == "__main__":
     print("🚀 Starting OhamLab Assistant...")
     demo = build_ui()
     demo.launch(server_name="0.0.0.0", server_port=7860)
+