Build a RAG system from scratch

python -m venv venv && source venv/bin/activate
pip install pypdf python-docx rank-bm25 numpy requests python-dotenv flask

python3 -m venv venv && source venv/bin/activate
pip install -r requirements.txt

python -m venv venv; venv\Scripts\Activate.ps1
pip install -r requirements.txt

from pathlib import Path
from typing import List, TypedDict

SUPPORTED_EXTS = {".pdf", ".docx", ".txt", ".md", ".markdown"}

class Page(TypedDict):
    source: str
    page_number: int
    text: str

def _extract_pdf(path: Path) -> List[Page]:
    from pypdf import PdfReader
    reader = PdfReader(str(path))
    pages = []
    for i, page in enumerate(reader.pages, start=1):
        text = (page.extract_text() or "").strip()
        if text:
            pages.append(Page(source=path.name, page_number=i, text=text))
    return pages

def extract_pages(path: Path) -> List[Page]:
    ext = path.suffix.lower()
    if ext == ".pdf":
        return _extract_pdf(path)
    if ext == ".docx":
        from docx import Document
        doc = Document(str(path))
        text = "\n".join(p.text for p in doc.paragraphs if p.text.strip())
        return [Page(source=path.name, page_number=1, text=text)] if text else []
    if ext in {".txt", ".md", ".markdown"}:
        text = path.read_text(encoding="utf-8", errors="replace").strip()
        return [Page(source=path.name, page_number=1, text=text)] if text else []
    return []

def _split_text(text, size, overlap):
    text = " ".join(text.split())            # normalize whitespace
    if len(text) <= size:
        return [text] if text else []
    chunks, start, n = [], 0, len(text)
    while start < n:
        end = min(start + size, n)
        if end < n:                          # prefer a clean break near the limit
            window = text[start:end]
            for sep in (". ", "! ", "? ", "\n", " "):
                pos = window.rfind(sep)
                if pos > size // 2:
                    end = start + pos + len(sep)
                    break
        chunk = text[start:end].strip()
        if chunk:
            chunks.append(chunk)
        if end >= n:
            break
        start = max(end - overlap, start + 1)   # step back to create overlap
    return chunks

def is_stale(config):
    """True if the cache is missing or the documents folder changed."""
    index_path = config.cache_dir / "index.json"
    if not index_path.exists():
        return True
    data = json.loads(index_path.read_text())
    return data.get("fingerprint") != _fingerprint(discover_files(config.docs_dir))

def build_index(config):
    files = discover_files(config.docs_dir)
    chunks = []
    for path in files:
        chunks.extend(chunk_pages(extract_pages(path)))
    config.cache_dir.mkdir(parents=True, exist_ok=True)
    (config.cache_dir / "index.json").write_text(
        json.dumps({"fingerprint": _fingerprint(files), "chunks": chunks}))
    return chunks, len(files)

import re
from rank_bm25 import BM25Okapi

def _tokenize(text):
    return re.findall(r"[a-z0-9]+", text.lower())

class Bm25Retriever:
    name = "bm25"
    def __init__(self, chunks):
        self.chunks = chunks
        self.bm25 = BM25Okapi([_tokenize(c["text"]) for c in chunks] or [[""]])

    def search(self, query, k):
        if not self.chunks:
            return []
        scores = self.bm25.get_scores(_tokenize(query))
        ranked = sorted(range(len(self.chunks)),
                        key=lambda i: scores[i], reverse=True)
        return [self.chunks[i] for i in ranked[:k]]

SYSTEM_PROMPT = """You are a careful assistant answering questions over a set
of the user's own documents. Follow these rules exactly:

1. Answer from the DOCUMENT CONTEXT below FIRST. For every claim that comes
   from the documents, cite the source like [filename:page].
2. If the answer is not contained in the document context, say so plainly:
   "This is not covered in your documents." You may then add general
   knowledge, but you MUST prefix it with
   "(general knowledge — not from your documents)".
3. Never invent document contents, quotes, or citations.
4. Be concise. Prefer the documents' own wording.
"""

def build_context(chunks):
    return "\n\n---\n\n".join(
        f"[{c['source']}:{c['page_number']}]\n{c['text']}" for c in chunks)

def build_user_prompt(question, chunks):
    context = build_context(chunks) if chunks else "(no relevant documents)"
    return f"DOCUMENT CONTEXT:\n{context}\n\nQUESTION:\n{question}"

from typing import Protocol

class LLMProvider(Protocol):
    name: str
    def is_available(self) -> bool: ...
    def chat(self, system: str, user: str) -> str: ...

def get_provider(name, config):
    if name == "claude":
        from .claude_code import ClaudeCodeProvider; return ClaudeCodeProvider(config)
    if name == "ollama":
        from .ollama import OllamaProvider; return OllamaProvider(config)
    if name == "gemini":
        from .gemini import GeminiProvider; return GeminiProvider(config)
    if name == "openai":
        from .openai import OpenAIProvider; return OpenAIProvider(config)
    raise ValueError(f"Unknown provider '{name}'")

import shutil, subprocess

class ClaudeCodeProvider:
    name = "claude"
    def __init__(self, config):
        self.bin = config.claude_bin
    def is_available(self):
        return shutil.which(self.bin) is not None
    def chat(self, system, user):
        result = subprocess.run(
            [self.bin, "-p", f"{system}\n\n{user}"],
            capture_output=True, text=True, timeout=180)
        return result.stdout.strip()

def answer(question, retriever, config):
    hits = retriever.search(question, config.top_k)
    provider = get_provider(config.provider, config)
    reply = provider.chat(SYSTEM_PROMPT, build_user_prompt(question, hits))
    print(reply)
    sources = []
    for h in hits:
        tag = f"{h['source']}:{h['page_number']}"
        if tag not in sources:
            sources.append(tag)
    print("Sources:", ", ".join(sources) or "(none)")

python -m localrag index
python -m localrag ask "How do I reset the device?"

class EmbeddingRetriever:
    name = "embeddings"
    def __init__(self, chunks, config):
        import numpy as np
        from .providers import embed_texts
        self.chunks, self.config = chunks, config
        vectors = np.asarray(
            embed_texts(config.embed_provider, config,
                        [c["text"] for c in chunks]), dtype="float32")
        self._vectors = self._normalize(vectors)

    def search(self, query, k):
        import numpy as np
        from .providers import embed_texts
        q = self._normalize(np.asarray(embed_texts(
            self.config.embed_provider, self.config, [query]),
            dtype="float32"))[0]
        sims = self._vectors @ q              # cosine (vectors normalized)
        ranked = np.argsort(sims)[::-1][:k]
        return [self.chunks[i] for i in ranked]

def build_retriever(chunks, config):
    if config.retriever == "embeddings":
        try:
            return EmbeddingRetriever(chunks, config)
        except Exception as exc:
            print(f"[localrag] Embeddings unavailable ({exc}). Falling back to BM25.")
    return Bm25Retriever(chunks)

RAG_RETRIEVER=embeddings RAG_EMBED_PROVIDER=ollama python -m localrag ask "power-cycle steps?"

@app.post("/api/upload")
def upload():
    for f in request.files.getlist("files"):
        name = secure_filename(f.filename)
        if Path(name).suffix.lower() in SUPPORTED_EXTS:
            f.save(base_config.docs_dir / name)
    chunks, n = refresh_index(base_config)      # rebuild on every drop
    return jsonify({"files": _list_files(base_config), "chunks": len(chunks)})

@app.post("/api/ask")
def ask():
    data = request.get_json()
    return jsonify(answer_question(_request_config(), data["question"]))

python -m localrag web      # http://127.0.0.1:5000

python -m localrag ask "How long is the warranty?"
python -m localrag ask "What is the capital of France?"

def test_bm25_finds_reset_instructions():
    chunks = chunk_pages(extract_pages(SAMPLE), size=400, overlap=80)
    hits = Bm25Retriever(chunks).search("how do I reset the device", k=3)
    assert "power button" in hits[0]["text"].lower()

pytest -q      # 4 passed