MaryamKarimi080's picture
Update app.py
6064d5e verified
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)))