Spaces:
Sleeping
Sleeping
import os | |
from pathlib import Path | |
import re | |
# Disable Chroma telemetry (optional) | |
os.environ["CHROMA_TELEMETRY_ENABLED"] = "false" | |
# Check if DB exists, else build | |
DB_DIR = Path(__file__).parent / "db" | |
if not DB_DIR.exists() or not any(DB_DIR.iterdir()): | |
print("π¦ No DB found. Building vectorstore...") | |
import scripts.load_documents | |
import scripts.chunk_and_embed | |
import scripts.setup_vectorstore | |
else: | |
print("β DB found. Skipping build.") | |
import gradio as gr | |
from scripts.router_chain import build_router_chain | |
OPENAI_KEY = os.getenv("OPENAI_API_KEY", None) | |
MODEL_NAME = os.getenv("OPENAI_MODEL", "gpt-4o-mini") | |
if not OPENAI_KEY: | |
print("WARNING: OPENAI_API_KEY not set. The app may fail at runtime.") | |
# Build the router once (keeps vectorstore & models in memory) | |
router = build_router_chain(model_name=MODEL_NAME) | |
def chat_fn(message, history): | |
if not message: | |
return history, "" | |
# call router | |
result = router.invoke({"input": message}) | |
# RetrievalQA returns dict with 'result' key (and maybe 'source_documents') | |
answer = result.get("result") if isinstance(result, dict) else str(result) | |
# append sources if present | |
sources = None | |
if isinstance(result, dict) and "source_documents" in result and result["source_documents"]: | |
try: | |
sources = list({str(d.metadata.get("source", "unknown")) for d in result["source_documents"]}) | |
except Exception: | |
sources = None | |
if sources: | |
answer = f"{answer}\n\nπ Sources: {', '.join(sources)}" | |
def format_answer(answer): | |
# Wrap LaTeX formulas in a span so MathJax can render them | |
answer = re.sub(r"\$\$(.+?)\$\$", r'<span class="math">$$\1$$</span>', answer) | |
return f"<div>{answer}</div>" | |
answer = format_answer(answer) | |
history.append((message, answer)) | |
return history, "" | |
CSS = """ | |
* { direction: rtl; text-align: right; font-family: 'Vazir', sans-serif; } | |
.gr-chat-message { white-space: pre-wrap; } | |
.math { font-size: 1.2em; } | |
""" | |
with gr.Blocks(css=CSS) as demo: | |
# demo.load(lambda: None, [], [], _js=""" | |
# const script = document.createElement('script'); | |
# script.src = "https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"; | |
# document.head.appendChild(script); | |
# """) | |
gr.Markdown("## π SCR Course Assistant β Chat with course files") | |
# chatbot = gr.Chatbot(elem_id="chatbot", type="messages") | |
chatbot = gr.Chatbot(elem_id="chatbot", type="tuples") | |
txt = gr.Textbox(show_label=False, placeholder="Ask about the course...") | |
txt.submit(chat_fn, [txt, chatbot], [chatbot, txt]) | |
txt.submit(lambda: None, None, txt) # clear input | |
if __name__ == "__main__": | |
demo.launch(server_port=int(os.getenv("PORT", 7860))) | |