Project: chat with documents

docs/*.md
   │  load + heading-aware chunking
   ▼
chunks (id, source, heading, text)
   │  embed                    │ tokenize
   ▼                           ▼
vector index              BM25 keyword index
   └───────────┬───────────────┘
               │ hybrid search + RRF + threshold
               ▼
        top-k chunks ──► grounded prompt ──► LLM (or offline extractive answer)
                                                │
                                                ▼
                                   answer + [n] citations + sources list

from pathlib import Path

DOCS = {
    "refunds.md": """# Refunds
## Annual plans
Annual plans can be refunded within 30 days of purchase. Refunds go back to the original card.
## Monthly plans
Monthly plans are not refunded. You can cancel a monthly plan at any time from Settings.
""",
    "backups.md": """# Backups
## Schedule
Backups run every night at 2 AM in your local time zone.
## Retention
Deleted files are kept for 90 days. After 90 days they are permanently removed.
## Errors
Error E-504 means the backup server timed out. Wait five minutes and retry.
""",
    "security.md": """# Security
## Encryption
All backups are encrypted with AES-256 at rest and TLS in transit.
## Admin accounts
Admins must enable two-factor authentication within 7 days of account creation.
""",
}

folder = Path("docs")
folder.mkdir(exist_ok=True)
for name, text in DOCS.items():
    (folder / name).write_text(text, encoding="utf-8")
print(sorted(p.name for p in folder.glob("*.md")))

"""docqa.py - chat with a folder of Markdown documents."""
from __future__ import annotations

import hashlib
import math
import os
import re
import sys
from collections import Counter
from pathlib import Path

import numpy as np

FALLBACK = "I couldn't find that in the documentation."
STOPWORDS = {"a", "an", "the", "is", "are", "do", "does", "i", "my", "to", "of", "in",
             "for", "and", "what", "how", "can", "it", "on", "be", "you", "your", "after"}
MIN_SIMILARITY = 0.2  # vector hits below this are treated as unrelated


# ---------- loading and chunking (lesson: Loading and chunking) ----------
def load_chunks(folder: Path) -> list[dict]:
    chunks = []
    for path in sorted(folder.glob("*.md")):
        path_headings: list[str] = []
        for block in re.split(r"\n(?=#)", path.read_text(encoding="utf-8")):
            lines = block.strip().splitlines()
            if not lines or not lines[0].startswith("#"):
                continue
            level = len(lines[0]) - len(lines[0].lstrip("#"))
            path_headings[:] = path_headings[: level - 1] + [lines[0].lstrip("#").strip()]
            body = " ".join(l.strip() for l in lines[1:] if l.strip())
            if body:
                heading = " > ".join(path_headings)
                chunks.append({"id": f"{path.name}#{len(chunks)}", "source": path.name,
                               "heading": heading, "text": f"{heading}: {body}"})
    return chunks


# ---------- embeddings (lesson: Embeddings) - swap in a real model here ----------
def words(text: str) -> list[str]:
    out = []
    for w in re.findall(r"[a-z0-9-]+", text.lower()):
        if w in STOPWORDS:
            continue
        for suffix in ("ing", "ed", "s"):
            if w.endswith(suffix) and len(w) > len(suffix) + 3:
                w = w[: -len(suffix)]
                break
        out.append(w)
    return out


def embed(text: str, dim: int = 512) -> np.ndarray:
    vec = np.zeros(dim, dtype=np.float32)
    for w in words(text):
        vec[int(hashlib.md5(w.encode()).hexdigest(), 16) % dim] += 1
    n = np.linalg.norm(vec)
    return vec / n if n else vec


# ---------- retrieval (lessons: Vector stores, Retrieval patterns) ----------
class Index:
    def __init__(self, chunks: list[dict]):
        self.chunks = chunks
        self.vectors = np.stack([embed(c["text"]) for c in chunks])
        self.tokens = [words(c["text"]) for c in chunks]
        self.df = Counter(t for toks in self.tokens for t in set(toks))
        self.avg_len = sum(map(len, self.tokens)) / len(self.tokens)

    def vector_rank(self, query: str) -> list[int]:
        scores = self.vectors @ embed(query)
        return [int(i) for i in np.argsort(-scores) if scores[i] >= MIN_SIMILARITY]

    def bm25_rank(self, query: str, k1: float = 1.5, b: float = 0.75) -> list[int]:
        n, scores = len(self.chunks), []
        for i, toks in enumerate(self.tokens):
            tf, s = Counter(toks), 0.0
            for t in words(query):
                if t in tf:
                    idf = math.log(1 + (n - self.df[t] + 0.5) / (self.df[t] + 0.5))
                    s += idf * tf[t] * (k1 + 1) / (tf[t] + k1 * (1 - b + b * len(toks) / self.avg_len))
            scores.append(s)
        return [i for i in sorted(range(n), key=lambda i: -scores[i]) if scores[i] > 0]

    def search(self, query: str, k: int = 3) -> list[dict]:
        fused: dict[int, float] = {}
        for ranking in (self.vector_rank(query), self.bm25_rank(query)):
            for rank, i in enumerate(ranking, start=1):
                fused[i] = fused.get(i, 0.0) + 1 / (60 + rank)
        best = sorted(fused, key=lambda i: -fused[i])[:k]
        return [{**self.chunks[i], "score": round(fused[i], 4)} for i in best]


# ---------- generation (lesson: Generation with sources) ----------
def build_messages(question: str, hits: list[dict]) -> list[dict]:
    sources = "\n".join(f"[{i}] ({h['source']}) {h['text']}" for i, h in enumerate(hits, 1))
    system = ("Answer using ONLY the numbered sources. End each sentence with a citation like [1]. "
              f'If the answer is not in the sources, reply exactly: "{FALLBACK}"')
    return [{"role": "system", "content": system},
            {"role": "user", "content": f"<sources>\n{sources}\n</sources>\n\nQuestion: {question}"}]


def extractive_answer(question: str, hits: list[dict]) -> str:
    """Offline mode: return the source sentence that best matches the question."""
    q = set(words(question))
    best, best_overlap = None, 0
    for n, h in enumerate(hits, start=1):
        body = h["text"].split(": ", 1)[-1]
        for sentence in re.split(r"(?<=\.)\s+", body):
            overlap = len(q & set(words(sentence)))
            if overlap > best_overlap:
                best, best_overlap = f"{sentence} [{n}]", overlap
    return best if best_overlap >= 2 else FALLBACK


def llm_answer(question: str, hits: list[dict]) -> str:
    from openai import OpenAI
    response = OpenAI().chat.completions.create(
        model=os.getenv("MODEL_NAME", "gpt-4.1-mini"),
        messages=build_messages(question, hits), temperature=0)
    return response.choices[0].message.content


def ask(index: Index, question: str) -> str:
    hits = index.search(question)
    if not hits:
        return FALLBACK
    use_llm = os.getenv("DOCQA_USE_LLM") == "1"
    answer = llm_answer(question, hits) if use_llm else extractive_answer(question, hits)
    if answer == FALLBACK:
        return answer
    sources = "\n".join(f"  [{i}] {h['source']} - {h['heading']}" for i, h in enumerate(hits, 1))
    return f"{answer}\nSources:\n{sources}"


if __name__ == "__main__":
    index = Index(load_chunks(Path(sys.argv[1] if len(sys.argv) > 1 else "docs")))
    print(f"indexed {len(index.chunks)} chunks")
    for question in sys.argv[2:] or ["How long are deleted files kept?"]:
        print(f"\nQ: {question}\nA: {ask(index, question)}")

python make_docs.py
python docqa.py docs "How long are deleted files kept?" "What does error E-504 mean?" "Can I get a refund on a monthly plan?" "Who is the CEO?"

export OPENAI_API_KEY=sk-...
export DOCQA_USE_LLM=1
python docqa.py docs "Can I get a refund on a monthly plan?"