""" 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 import traceback import numpy as np 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") or os.environ.get("HUGGINGFACE_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 = "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 and Chunk Markdown Files # --------------------------- def load_ohamlab_knowledge(): """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() buf = "" for line in content.splitlines(): buf += line.strip() + " " if len(buf) >= 500: chunks.append({"file": f, "text": buf.strip()}) buf = "" if buf: chunks.append({"file": f, "text": buf.strip()}) except Exception as e: print(f"⚠️ Failed to load {f}: {e}") return chunks # --------------------------- # 3. Generate or Load Embeddings (with Cache) # --------------------------- def get_embeddings_with_cache(): """Generate or load cached embeddings for OhamLab context.""" if os.path.exists(CACHE_PATH): try: 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 from {CACHE_PATH} ({len(embs)} chunks)") return texts, embs except Exception: print("⚠️ Cache corrupted, regenerating embeddings...") 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] 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") 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. Semantic Retrieval # --------------------------- def retrieve_knowledge(query, top_k=3): """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) / ( np.linalg.norm(OHAMLAB_EMBS, axis=1) * np.linalg.norm(q_emb) ) top_idx = np.argsort(sims)[-top_k:][::-1] return "\n\n".join(OHAMLAB_TEXTS[i] for i in top_idx) except Exception as e: print(f"⚠️ Retrieval error: {e}") return "" # --------------------------- # 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 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 (Retrieved Snippets) --- {context[:1800]} --- End Context --- """).strip() # --------------------------- # 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=1200, ) return resp.choices[0].message.content.strip() except Exception as e: print(f"⚠️ Model call failed: {e}") return "⚠️ OahmLab encountered a temporary issue generating your response." # --------------------------- # 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: bot_reply = generate_response(user_message, history) 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"), 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], ) # Input box (full width) with gr.Row(): 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) return demo # --------------------------- # Entrypoint # --------------------------- if __name__ == "__main__": print("🚀 Starting OhamLab Assistant...") demo = build_ui() demo.launch(server_name="0.0.0.0", server_port=7860)