rag.py

"""
RAG support: chunking extracted text, indexing it into Chroma, and
retrieving relevant chunks for a given query.
"""

import chromadb
import re

def chunk_text(text: str, chunk_size: int = 1000, overlap: int = 100) -> list[str]:
    """
    Split raw text into overlapping chunks.
    """
    start = 0
    chunks = []
    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        start = end - overlap
    return chunks


def build_collection(chunks: list[str]):
    """
    Embed and index the given chunks into a fresh Chroma collection.
    """
    if not chunks:
        return None
    client = chromadb.Client()
    try:
        client.delete_collection("study_material")
    except Exception:
        pass
    collection = client.create_collection("study_material")
    collection.add(documents=chunks, ids=[f"chunk_{i}" for i in range(len(chunks))])
    return collection

def retrieve_relevant_chunks(collection, query: str, k: int = 4) -> list[str]:
    """
    Given a user's question, return the top-k most relevant chunks.
    """
    if collection is None: return []
    num_docs = collection.count()
    if num_docs == 0:
        return []
    actual_k = min(k, num_docs)
    results = collection.query(query_texts=[query], n_results=actual_k)
    return results["documents"][0]

def find_direct_reference(study_context: str, query: str) -> str:
    """ch
    If the query references a specific numbered item (e.g. "question 5",
    "problem 12", "chapter 3"), search the raw text directly for that label
    and return a window of text starting there. This catches literal
    lookups that semantic embedding search is bad at.

    Returns an empty string if no reference pattern is found in the query,
    or if that label doesn't appear anywhere in study_context.
    """
    match = re.search(r'\b(question|problem|exercise|chapter)\s*#?\s*(\d+)\b', query, re.IGNORECASE)
    if not match:
        return ""

    label, number = match.group(1), match.group(2)
    # Look for that exact label + number in the source text
    pattern = re.compile(rf'\b{label}\s*#?\s*{number}\b', re.IGNORECASE)
    found = pattern.search(study_context)
    if not found:
        return ""

    # Grab a chunk of text starting at the match — enough to likely contain
    # the full question, without running all the way to the end of the doc.
    start = found.start()
    return study_context[start:start + 2000]

def clean_math_output(text: str) -> str:
    # \[ ... \] -> $$ ... $$ , \( ... \) -> $ ... $
    text = re.sub(r'\\\[(.*?)\\\]', r'$$\1$$', text, flags=re.DOTALL)
    text = re.sub(r'\\\((.*?)\\\)', r'$\1$', text, flags=re.DOTALL)

    # Catches cases where the model skips delimiters entirely and just uses
    # plain ( ) around LaTeX-flavored content.
    text = re.sub(r'\(([^()]*\\[a-zA-Z]+[^()]*)\)', r'$\1$', text)

    # Strip cosmetic sizing commands
    text = re.sub(r'\\[Bb]igl|\\[Bb]igr|\\[Bb]ig\b', '', text)

    # Fix stray semicolons before \text{}
    text = re.sub(r';(\\text)', r'\\,\1', text)

    return text