import os
import re
import json
import pickle
from urllib.parse import quote
import numpy as np
import gradio as gr
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer
from openai import OpenAI
# ============================================================
# Configuration
# ============================================================
BUILD_DIR = "brainchat_build"
CHUNKS_PATH = os.path.join(BUILD_DIR, "chunks.pkl")
TOKENS_PATH = os.path.join(BUILD_DIR, "tokenized_chunks.pkl")
EMBED_PATH = os.path.join(BUILD_DIR, "embeddings.npy")
CONFIG_PATH = os.path.join(BUILD_DIR, "config.json")
# Put ONE of these logo files in your Space repo root (same folder as app.py)
LOGO_CANDIDATES = [
"Brain chat-09.png",
"brainchat_logo.png.png",
"Brain Chat Imagen.svg",
"ebcbb9f5-022f-473a-bf51-7e7974f794b4.png",
]
MODEL_NAME_TEXT = os.getenv("OPENAI_MODEL", "gpt-4o-mini")
# ============================================================
# Globals (lazy loaded)
# ============================================================
BM25 = None
CHUNKS = None
EMBEDDINGS = None
EMBED_MODEL = None
CLIENT = None
# ============================================================
# Utilities
# ============================================================
def tokenize(text: str):
return re.findall(r"\w+", text.lower(), flags=re.UNICODE)
def ensure_loaded():
global BM25, CHUNKS, EMBEDDINGS, EMBED_MODEL, CLIENT
if CHUNKS is None:
missing = [p for p in [CHUNKS_PATH, TOKENS_PATH, EMBED_PATH, CONFIG_PATH] if not os.path.exists(p)]
if missing:
raise FileNotFoundError(
"Missing build files. Make sure you ran the build step and committed brainchat_build/.\n"
+ "\n".join(missing)
)
with open(CHUNKS_PATH, "rb") as f:
CHUNKS = pickle.load(f)
with open(TOKENS_PATH, "rb") as f:
tokenized_chunks = pickle.load(f)
EMBEDDINGS = np.load(EMBED_PATH)
with open(CONFIG_PATH, "r", encoding="utf-8") as f:
cfg = json.load(f)
BM25 = BM25Okapi(tokenized_chunks)
EMBED_MODEL = SentenceTransformer(cfg["embedding_model"])
if CLIENT is None:
api_key = os.getenv("OPENAI_API_KEY")
if not api_key:
raise ValueError("OPENAI_API_KEY is missing. Add it in your Space Secrets.")
CLIENT = OpenAI(api_key=api_key)
def search_hybrid(query: str, shortlist_k: int = 30, final_k: int = 5):
ensure_loaded()
q_tokens = tokenize(query)
bm25_scores = BM25.get_scores(q_tokens)
shortlist_idx = np.argsort(bm25_scores)[::-1][:shortlist_k]
qvec = EMBED_MODEL.encode([query], normalize_embeddings=True).astype("float32")[0]
shortlist_emb = EMBEDDINGS[shortlist_idx]
dense_scores = shortlist_emb @ qvec
rerank = np.argsort(dense_scores)[::-1][:final_k]
final_idx = shortlist_idx[rerank]
return [CHUNKS[int(i)] for i in final_idx]
def build_context(records):
blocks = []
for i, r in enumerate(records, start=1):
blocks.append(
f"""[Source {i}]
Book: {r.get('book','')}
Section: {r.get('section_title','')}
Pages: {r.get('page_start','')}-{r.get('page_end','')}
Text:
{r.get('text','')}"""
)
return "\n\n".join(blocks)
def make_sources(records):
seen = set()
lines = []
for r in records:
key = (r.get("book"), r.get("section_title"), r.get("page_start"), r.get("page_end"))
if key in seen:
continue
seen.add(key)
lines.append(
f"• {r.get('book','')} | {r.get('section_title','')} | pp. {r.get('page_start','')}-{r.get('page_end','')}"
)
return "\n".join(lines)
def choose_quiz_count(user_text: str, selector: str) -> int:
if selector in {"3", "5", "7"}:
return int(selector)
t = user_text.lower()
if any(k in t for k in ["mock test", "final exam", "exam practice", "full test"]):
return 7
if any(k in t for k in ["detailed", "revision", "comprehensive", "study"]):
return 5
return 3
def language_instruction(language_mode: str) -> str:
if language_mode == "English":
return "Answer only in English."
if language_mode == "Spanish":
return "Answer only in Spanish."
if language_mode == "Bilingual":
return "Answer first in English, then provide a Spanish version under the heading 'Español:'."
return "If the user writes in Spanish, answer in Spanish; otherwise answer in English."
def build_tutor_prompt(mode: str, language_mode: str, question: str, context: str) -> str:
mode_map = {
"Explain": (
"Explain clearly like a friendly tutor using simple language. "
"Use short headings if helpful."
),
"Detailed": (
"Give a detailed explanation. Include key terms and clinical relevance only if supported by the context."
),
"Short Notes": "Write concise revision notes using bullet points.",
"Flashcards": "Create 6 flashcards in Q/A format.",
"Case-Based": (
"Create a short clinical scenario (2–4 lines) and then explain the underlying concept using the context."
),
}
return f"""
You are BrainChat, an interactive neurology and neuroanatomy tutor.
Rules:
- Use ONLY the provided context from the books.
- If the answer is not supported by the context, say exactly:
Not found in the course material.
- Do not invent facts outside the context.
- {language_instruction(language_mode)}
Teaching style:
{mode_map.get(mode, mode_map['Explain'])}
Context:
{context}
Student question:
{question}
""".strip()
def build_quiz_generation_prompt(language_mode: str, topic: str, context: str, n_questions: int) -> str:
return f"""
You are BrainChat, an interactive tutor.
Rules:
- Use ONLY the provided context.
- Create exactly {n_questions} quiz questions.
- Questions should be short, clear, and course-aligned.
- Provide a short answer key per question.
- Return VALID JSON only.
- {language_instruction(language_mode)}
Required JSON format:
{{
"title": "short quiz title",
"questions": [
{{"q": "question 1", "answer_key": "expected short answer"}},
{{"q": "question 2", "answer_key": "expected short answer"}}
]
}}
Context:
{context}
Topic:
{topic}
""".strip()
def build_quiz_evaluation_prompt(language_mode: str, quiz_data: dict, user_answers: str) -> str:
quiz_json = json.dumps(quiz_data, ensure_ascii=False)
return f"""
You are BrainChat, an interactive tutor.
Task:
Evaluate the student's answers fairly against the answer keys.
Accept semantically correct answers even if wording differs.
Return VALID JSON only.
Required JSON format:
{{
"score_obtained": 0,
"score_total": 0,
"summary": "short overall feedback",
"results": [
{{
"question": "question text",
"answer_key": "expected short answer",
"student_answer": "student answer",
"result": "Correct / Partially Correct / Incorrect",
"feedback": "short explanation"
}}
],
"improvement_tip": "one short study suggestion"
}}
Quiz:
{quiz_json}
Student answers:
{user_answers}
Language:
{language_instruction(language_mode)}
""".strip()
def chat_text(prompt: str) -> str:
ensure_loaded()
resp = CLIENT.chat.completions.create(
model=MODEL_NAME_TEXT,
temperature=0.2,
messages=[
{"role": "system", "content": "You are a helpful educational assistant."},
{"role": "user", "content": prompt},
],
)
return resp.choices[0].message.content.strip()
def chat_json(prompt: str) -> dict:
ensure_loaded()
resp = CLIENT.chat.completions.create(
model=MODEL_NAME_TEXT,
temperature=0.2,
response_format={"type": "json_object"},
messages=[
{"role": "system", "content": "Return only valid JSON."},
{"role": "user", "content": prompt},
],
)
return json.loads(resp.choices[0].message.content)
# ============================================================
# Logo + Header HTML
# ============================================================
def find_logo_file():
for name in LOGO_CANDIDATES:
if os.path.exists(name):
return name
return None
def logo_img_tag(size_px: int = 88) -> str:
logo_file = find_logo_file()
if logo_file:
url = f"/gradio_api/file={quote(logo_file)}"
return f''
return '