🤖 Build a company-docs RAG assistant with Qwen

import os

EMBED_MODEL = "qwen3-embedding:0.6b"
CHAT_MODEL = os.environ.get("QWEN_CHAT_MODEL", "qwen3.5:9b")   # qwen3.8:27b on a 24 GB GPU
DB_PATH = "rag.db"

# Search results scoring below this are not used as sources. Score distributions differ by embedding model and documents,
# so set it in step 7 between the lowest top score of answerable questions and the highest of unanswerable ones.
MIN_SCORE = float(os.environ.get("RAG_MIN_SCORE", "0.35"))
DOCS_DIR = "docs"

# Instruction added to search queries only (the format recommended on the Qwen3 Embedding model card)
QUERY_TASK = "Given a web search query, retrieve relevant passages that answer the query"

import re
from dataclasses import dataclass
from pathlib import Path


@dataclass
class Chunk:
    source: str      # file name
    heading: str     # nearest heading — shown as the source location in answers
    text: str


def split_markdown(path: Path, max_chars: int = 500, overlap: int = 80) -> list[Chunk]:
    """Split by headings (#) first, then re-split only long sections by character count.
    Character counts behave predictably across languages, unlike word counts."""
    chunks: list[Chunk] = []
    heading, buf = path.stem, []

    def flush():
        body = "\n".join(buf).strip()
        start = 0
        while body and start < len(body):
            chunks.append(Chunk(path.name, heading, body[start:start + max_chars]))
            start += max_chars - overlap
        buf.clear()

    for line in path.read_text(encoding="utf-8").splitlines():
        m = re.match(r"^(#{1,3})\s+(.*)", line)
        if m:
            flush()
            heading = m.group(2).strip()
        else:
            buf.append(line)
    flush()
    return chunks


if __name__ == "__main__":
    for c in split_markdown(Path("docs/leave-policy.md")):
        print(f"[{c.source} > {c.heading}] {c.text[:60]}")

import hashlib
import json
import math
import sqlite3
from pathlib import Path
from ollama import embed
from chunk import split_markdown
from config import DB_PATH, DOCS_DIR, EMBED_MODEL

db = sqlite3.connect(DB_PATH, check_same_thread=False)
db.executescript("""
CREATE TABLE IF NOT EXISTS files (name TEXT PRIMARY KEY, sha256 TEXT);
CREATE TABLE IF NOT EXISTS chunks (id INTEGER PRIMARY KEY, source TEXT, heading TEXT, text TEXT, vec TEXT);
""")


def normalize(vec: list[float]) -> list[float]:
    norm = math.sqrt(sum(x * x for x in vec)) or 1.0
    return [x / norm for x in vec]


def index_file(path: Path) -> int:
    chunks = split_markdown(path)
    vecs = embed(model=EMBED_MODEL, input=[c.text for c in chunks]).embeddings   # no instruction on documents
    db.execute("DELETE FROM chunks WHERE source = ?", (path.name,))
    db.executemany(
        "INSERT INTO chunks (source, heading, text, vec) VALUES (?, ?, ?, ?)",
        [(c.source, c.heading, c.text, json.dumps(normalize(v))) for c, v in zip(chunks, vecs)],
    )
    return len(chunks)


def sync() -> None:
    """Re-index only changed documents — files with the same hash are not embedded again"""
    current = {p.name: p for p in Path(DOCS_DIR).glob("*.md")}
    known = dict(db.execute("SELECT name, sha256 FROM files"))
    for name, path in current.items():
        digest = hashlib.sha256(path.read_bytes()).hexdigest()
        if known.get(name) == digest:
            continue
        n = index_file(path)
        db.execute("INSERT OR REPLACE INTO files VALUES (?, ?)", (name, digest))
        print(f"indexed {name}: {n} chunks")
    for name in set(known) - set(current):          # clean up deleted documents
        db.execute("DELETE FROM chunks WHERE source = ?", (name,))
        db.execute("DELETE FROM files WHERE name = ?", (name,))
        print(f"removed {name}")
    db.commit()


if __name__ == "__main__":
    sync()
    print("total chunks:", db.execute("SELECT COUNT(*) FROM chunks").fetchone()[0])

import json
import sys
from ollama import embed
from config import EMBED_MODEL, QUERY_TASK
from store import db, normalize


def search(question: str, k: int = 3) -> list[dict]:
    q = normalize(embed(model=EMBED_MODEL, input=[f"Instruct: {QUERY_TASK}\nQuery:{question}"]).embeddings[0])
    scored = []
    for source, heading, text, vec in db.execute("SELECT source, heading, text, vec FROM chunks"):
        score = sum(a * b for a, b in zip(q, json.loads(vec)))     # normalized vectors, so dot product = cosine similarity
        scored.append({"source": source, "heading": heading, "text": text, "score": round(score, 3)})
    return sorted(scored, key=lambda r: -r["score"])[:k]


if __name__ == "__main__":
    for r in search(sys.argv[1] if len(sys.argv) > 1 else "How many vacation days do I get?"):
        print(f"{r['score']:.3f}  [{r['source']} > {r['heading']}] {r['text'][:50]}")

import sys
from ollama import chat
from config import CHAT_MODEL, MIN_SCORE
from search import search

SYSTEM = """You are a company policy assistant. Always answer in the language of the question.
- Answer only from the numbered [n] sources below, and add the source number to every sentence, like [1].
- If the sources do not contain the answer, reply only "I could not find this in the company documents." and do not guess."""


def answer(question: str, min_score: float = MIN_SCORE) -> dict:
    hits = [h for h in search(question) if h["score"] >= min_score]
    if not hits:   # no relevant document → do not call the model at all, so it cannot make up an answer
        return {"answer": "I could not find this in the company documents.", "sources": []}
    context = "\n\n".join(f"[{i}] ({h['source']} > {h['heading']})\n{h['text']}" for i, h in enumerate(hits, 1))
    resp = chat(
        model=CHAT_MODEL,
        messages=[{"role": "system", "content": SYSTEM},
                  {"role": "user", "content": f"Sources:\n{context}\n\nQuestion: {question}"}],
        think=False,
        options={"temperature": 0.2, "num_ctx": 8192},
    )
    return {"answer": resp.message.content, "sources": hits}


if __name__ == "__main__":
    result = answer(sys.argv[1] if len(sys.argv) > 1 else "I'm in my second year. How many vacation days do I have?")
    print(result["answer"])
    for i, s in enumerate(result["sources"], 1):
        print(f"  [{i}] {s['source']} > {s['heading']} ({s['score']})")

import re
from ollama import chat
from pydantic import BaseModel
from config import CHAT_MODEL


class Verdict(BaseModel):
    supported: bool
    unsupported_claims: list[str]


def check_citations(answer: str, n_sources: int) -> list[str]:
    """Rule check: find sentences without a source number or citing a number that does not exist"""
    problems = []
    # Move citations that follow the punctuation ("15 days. [1]") in front of it ("15 days [1].") before splitting sentences
    text = re.sub(r"([.!?])\s*((?:\[\d+\]\s*)+)", lambda m: " " + m.group(2).strip() + m.group(1), answer.strip())
    for sentence in re.split(r"(?<=[.!?])\s+", text):
        if not sentence or "could not find" in sentence:
            continue
        cited = [int(n) for n in re.findall(r"\[(\d+)\]", sentence)]
        if not cited:
            problems.append(f"no source: {sentence}")
        elif any(n < 1 or n > n_sources for n in cited):
            problems.append(f"unknown source number: {sentence}")
    return problems


def check_faithfulness(answer: str, sources: list[dict]) -> Verdict:
    """Model check: judge again whether each claim in the answer is backed by the sources"""
    context = "\n\n".join(f"[{i}] {s['text']}" for i, s in enumerate(sources, 1))
    resp = chat(
        model=CHAT_MODEL,
        messages=[{"role": "user", "content": f"Sources:\n{context}\n\nAnswer:\n{answer}\n\nFind any claims in the answer not supported by the sources."}],
        format=Verdict.model_json_schema(),
        think=False,
        options={"temperature": 0},
    )
    return Verdict.model_validate_json(resp.message.content)


if __name__ == "__main__":
    from answer import answer as ask
    result = ask("I'm in my second year. How many vacation days do I have?")
    print("rule check:", check_citations(result["answer"], len(result["sources"])) or "passed")
    print("model check:", check_faithfulness(result["answer"], result["sources"]))

{"question": "How many vacation days do I get after one year?", "source": "leave-policy.md", "must_include": ["15 days"]}
{"question": "What do I do with my laptop when I leave the company?", "source": "equipment.md", "must_include": ["return"]}
{"question": "Where do I request a VPN account?", "source": "security.md", "must_include": ["portal"]}
{"question": "Where can I see the cafeteria menu?", "source": "-", "must_include": ["could not find"]}

import json
from answer import answer
from config import MIN_SCORE
from search import search

cases = [json.loads(line) for line in open("data/qa.jsonl", encoding="utf-8")]

retrieval_hits = answer_hits = 0
answerable_scores, unanswerable_scores = [], []
for case in cases:
    results = search(case["question"], k=3)
    top = results[0]["score"] if results else 0.0
    if case["source"] == "-":                       # the documents do not answer this question
        unanswerable_scores.append(top)
        retrieved = top < MIN_SCORE                  # correct only if nothing clears the threshold
    else:
        answerable_scores.append(top)
        retrieved = case["source"] in [r["source"] for r in results]
    retrieval_hits += retrieved
    reply = answer(case["question"])["answer"]
    ok = all(k.lower() in reply.lower() for k in case["must_include"])
    answer_hits += ok
    print(f"{'✓' if retrieved and ok else '✗'} {case['question']}  top score {top:.3f} · retrieval {'O' if retrieved else 'X'} · answer {'O' if ok else 'X'}")

n = len(cases)
print(f"retrieval {retrieval_hits}/{n} · answers {answer_hits}/{n}")
if answerable_scores and unanswerable_scores:
    print(f"lowest top score of answerable questions {min(answerable_scores):.3f} / highest of unanswerable {max(unanswerable_scores):.3f}"
          " → set MIN_SCORE between the two")

from contextlib import asynccontextmanager
from fastapi import FastAPI
from pydantic import BaseModel, Field
from answer import answer
from store import sync
from verify import check_citations


@asynccontextmanager
async def lifespan(app: FastAPI):
    sync()                      # on startup, re-index only changed documents
    yield


app = FastAPI(title="Company Docs Assistant", lifespan=lifespan)


class Question(BaseModel):
    text: str = Field(min_length=1, max_length=500)


@app.post("/ask")
def ask(q: Question) -> dict:
    result = answer(q.text)
    problems = check_citations(result["answer"], len(result["sources"]))
    return {
        "answer": result["answer"],
        "sources": [{"source": s["source"], "heading": s["heading"], "score": s["score"]} for s in result["sources"]],
        "needs_review": bool(problems),    # flag answers with unsupported sentences as "needs review" in the UI
        "problems": problems,
    }