# One-time: Postgres with the pgvector extension
#   CREATE EXTENSION vector;
#   CREATE TABLE chunks (
#     id bigserial PRIMARY KEY,
#     doc_id text, source text, chunk text,
#     embedding vector(1024)
#   );

import psycopg, ollama
db = psycopg.connect("postgresql://localhost/kb")

def embed(text):
    return ollama.embed(model="bge-m3", input=text)["embeddings"][0]

# ---- INGEST (runs on a schedule) --------------------------------
for doc_id, source, chunk in load_and_split_your_documents():
    db.execute("INSERT INTO chunks (doc_id, source, chunk, embedding)"
               " VALUES (%s, %s, %s, %s)",
               (doc_id, source, chunk, embed(chunk)))
db.commit()

# ---- QUERY (runs per question) ----------------------------------
def answer(question):
    rows = db.execute(
        "SELECT chunk, source FROM chunks"
        " ORDER BY embedding <=> %s::vector"      # <=> is cosine distance
        " LIMIT 5", (embed(question),)).fetchall()

    context = "\n\n".join(f"[{i+1}] (source: {s})\n{c}"
                           for i, (c, s) in enumerate(rows))

    return ollama.chat(model="qwen3:8b", messages=[
        {"role": "system", "content":
         "Answer using ONLY the numbered sources below. Cite them as [1], [2]. "
         "If the sources do not contain the answer, say so and stop. Do not guess."},
        {"role": "user", "content": f"{context}\n\nQuestion: {question}"},
    ])["message"]["content"]

# Structural chunking, with the metadata that makes filtering possible later
import re

def split_by_heading(text, doc_id, source, max_chars=1500):
    sections = re.split(r"\n(?=#{1,3}\s)", text)     # split before markdown headings
    for sec in sections:
        heading = sec.split("\n", 1)[0].lstrip("# ").strip()
        # a long section still needs splitting -- but now every piece knows
        # which section it came from, which is what makes citations useful
        for i in range(0, len(sec), max_chars):
            yield {
                "doc_id":  doc_id,
                "source":  source,
                "heading": heading,
                "chunk":   sec[i:i + max_chars],
                # capture these AT INGEST. Adding them later means a full re-index.
                "updated_at":  file_mtime(source),
                "acl_groups":  groups_that_can_read(source),
            }

-- Hybrid search in one SQL statement: keyword and vector, fused by rank.
WITH semantic AS (
  SELECT id, ROW_NUMBER() OVER (ORDER BY embedding <=> %(qvec)s::vector) AS rank
  FROM chunks WHERE acl_groups && %(user_groups)s     -- permission filter, applied FIRST
  ORDER BY embedding <=> %(qvec)s::vector LIMIT 50
),
keyword AS (
  SELECT id, ROW_NUMBER() OVER (
           ORDER BY ts_rank_cd(tsv, websearch_to_tsquery(%(q)s)) DESC) AS rank
  FROM chunks
  WHERE tsv @@ websearch_to_tsquery(%(q)s) AND acl_groups && %(user_groups)s
  LIMIT 50
)
SELECT c.id, c.chunk, c.source,
       COALESCE(1.0 / (60 + s.rank), 0) +
       COALESCE(1.0 / (60 + k.rank), 0) AS score      -- Reciprocal Rank Fusion
FROM chunks c
LEFT JOIN semantic s ON s.id = c.id
LEFT JOIN keyword  k ON k.id = c.id
WHERE s.id IS NOT NULL OR k.id IS NOT NULL
ORDER BY score DESC
LIMIT 50;      -- 50 candidates go to the reranker, which returns the 5 you use

# Reranking: read the question and each candidate TOGETHER, then keep the best few
from FlagEmbedding import FlagReranker
reranker = FlagReranker("BAAI/bge-reranker-v2-m3", use_fp16=True)

pairs  = [[question, c["chunk"]] for c in candidates]        # 50 candidates
scores = reranker.compute_score(pairs)

top = [c for _, c in sorted(zip(scores, candidates),
                            key=lambda x: -x[0])][:5]        # keep 5

# Worth measuring rather than assuming: on most corpora this is the single
# largest accuracy gain per unit of effort in the whole pipeline. It costs
# roughly 50-200ms, which is why it runs on 50 candidates and not 50,000.

# Retrieval evaluation. No framework needed -- this is the whole thing.
def evaluate(eval_set, retrieve, k=10):
    hits, rr = 0, 0.0
    for q, gold_doc_id in eval_set:
        ids = [c["doc_id"] for c in retrieve(q, limit=k)]
        if gold_doc_id in ids:
            hits += 1
            rr   += 1.0 / (ids.index(gold_doc_id) + 1)
    n = len(eval_set)
    return {"recall@k": hits / n, "mrr": rr / n}

print("vector only  ", evaluate(eval_set, retrieve_vector))
print("hybrid       ", evaluate(eval_set, retrieve_hybrid))
print("hybrid+rerank", evaluate(eval_set, retrieve_hybrid_reranked))
# vector only   {'recall@k': 0.72, 'mrr': 0.51}
# hybrid        {'recall@k': 0.88, 'mrr': 0.63}
# hybrid+rerank {'recall@k': 0.88, 'mrr': 0.81}