2. RAG with pgvector

"""PostgreSQL + pgvector: the vector store of the assistant (schema "assistant")."""
import psycopg
from pgvector.psycopg import register_vector_async
from psycopg_pool import AsyncConnectionPool

from .config import settings

SCHEMA = f"""
CREATE EXTENSION IF NOT EXISTS vector;
CREATE SCHEMA IF NOT EXISTS assistant;
CREATE TABLE IF NOT EXISTS assistant.documents (
    id          TEXT PRIMARY KEY,                       -- e.g. "event:123"
    kind        TEXT NOT NULL,                          -- "event" (add "faq", "policy" ... later)
    title       TEXT NOT NULL,
    content     TEXT NOT NULL,                          -- the text that was embedded
    embedding   vector({settings.embed_dim}) NOT NULL,
    updated_at  TIMESTAMPTZ NOT NULL DEFAULT now()
);
-- approximate nearest neighbours (HNSW) on cosine distance, the metric of the embedding model
CREATE INDEX IF NOT EXISTS documents_embedding_hnsw
    ON assistant.documents USING hnsw (embedding vector_cosine_ops);
"""


async def migrate() -> None:
    """Create the extension and the table before the pool registers the vector type."""
    async with await psycopg.AsyncConnection.connect(settings.db_dsn, autocommit=True) as conn:
        await conn.execute(SCHEMA)


pool = AsyncConnectionPool(settings.db_dsn, configure=register_vector_async, min_size=1, max_size=10, open=False)

"""Retrieval-augmented generation: catalog events -> text -> embeddings -> pgvector -> top-k."""
import asyncio
from dataclasses import dataclass

import httpx
import numpy as np

from .config import settings
from .db import pool
from .llm import embed

# nomic-embed-text was trained with task prefixes: documents and queries are embedded differently
DOC_PREFIX, QUERY_PREFIX = "search_document: ", "search_query: "


@dataclass
class Hit:
    id: str
    title: str
    content: str
    score: float  # cosine similarity, 1 = same direction


def event_text(e: dict) -> str:
    """The text a fan would search for: what, where, when, how much."""
    v = e["venue"]
    return (f"{e['title']} ({e['category'].lower()}) at {v['name']}, {v['city']} "
            f"(capacity {v['capacity']}), on {e['startsAt']}. Tickets from R$ {e['priceFrom']}.")


async def fetch_events(limit: int) -> list[dict]:
    """Events 1..limit from the catalog API (8 requests in parallel)."""
    sem = asyncio.Semaphore(8)
    async with httpx.AsyncClient(base_url=settings.catalog_url, timeout=10) as http:
        async def one(event_id: int) -> dict | None:
            async with sem:
                r = await http.get(f"/catalog/events/{event_id}")
                return r.json() if r.status_code == 200 else None
        found = await asyncio.gather(*(one(i) for i in range(1, limit + 1)))
    return [e for e in found if e]


async def ingest(limit: int, batch: int = 32) -> int:
    events = await fetch_events(limit)
    for start in range(0, len(events), batch):
        chunk = events[start:start + batch]
        texts = [event_text(e) for e in chunk]
        vectors = await embed([DOC_PREFIX + t for t in texts])
        async with pool.connection() as conn:
            async with conn.cursor() as cur:
                await cur.executemany(
                    """
                    INSERT INTO assistant.documents (id, kind, title, content, embedding)
                    VALUES (%s, 'event', %s, %s, %s)
                    ON CONFLICT (id) DO UPDATE
                    SET title = EXCLUDED.title, content = EXCLUDED.content,
                        embedding = EXCLUDED.embedding, updated_at = now()
                    """,
                    [(f"event:{e['id']}", e["title"], t, np.array(v, dtype=np.float32))
                     for e, t, v in zip(chunk, texts, vectors)],
                )
    return len(events)


async def search(query: str, k: int) -> list[Hit]:
    [vector] = await embed([QUERY_PREFIX + query])
    async with pool.connection() as conn:
        rows = await (await conn.execute(
            """
            SELECT id, title, content, 1 - (embedding <=> %s) AS score
            FROM assistant.documents
            ORDER BY embedding <=> %s
            LIMIT %s
            """,
            (np.array(vector, dtype=np.float32), np.array(vector, dtype=np.float32), k),
        )).fetchall()
    return [Hit(id=r[0], title=r[1], content=r[2], score=float(r[3])) for r in rows]

# /// script
# requires-python = ">=3.12"
# dependencies = ["httpx>=0.28"]
# ///
"""Retrieval evaluation: does the vector search find the right events? (hit@k and MRR)

The catalog seed of lab 1 is generated by a formula, so the right answers are known exactly:
event g is title ["Concert", "Match", "Festival", "Play"][g % 4] + " #g", at venue 1 + (g * 7) % 50.

    uv run scripts/eval_retrieval.py --k 5 --limit 300
"""
import argparse

import httpx

CITIES = ["São Paulo", "Rio de Janeiro", "Belo Horizonte", "Porto Alegre", "Curitiba",
          "Salvador", "Recife", "Fortaleza", "Brasília", "Manaus"]
KINDS = ["Stadium", "Arena", "Hall", "Theater", "Club"]
TITLES = ["Concert", "Match", "Festival", "Play"]
WORDS = {"Concert": "concert", "Match": "match", "Festival": "festival", "Play": "play"}


def event(g: int) -> dict:
    if g == 1:
        return {"id": 1, "kind": "Festival", "venue_kind": "Stadium", "city": "São Paulo"}
    v = 1 + (g * 7) % 50
    return {"id": g, "kind": TITLES[g % 4], "venue_kind": KINDS[(v - 1) % 5], "city": CITIES[(v - 1) // 5]}


def labeled(limit: int) -> list[tuple[str, set[str]]]:
    events = [event(g) for g in range(1, limit + 1)]
    cases = []
    for g in [7, 42, 118, 199, 256]:  # exact lookups: one right answer
        if g <= limit:
            e = event(g)
            cases.append((f"When is {e['kind']} #{g}?", {f"event:{g}"}))
    for kind, city in [("Festival", "Recife"), ("Play", "Curitiba"), ("Match", "Salvador"),
                       ("Concert", "Manaus"), ("Festival", "São Paulo")]:  # open questions: many right answers
        right = {f"event:{e['id']}" for e in events if e["kind"] == kind and e["city"] == city}
        if right:
            cases.append((f"Is there any {WORDS[kind]} in {city}?", right))
    cases.append(("the big festival night at the stadium in São Paulo", {"event:1"}))
    return cases


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--url", default="http://localhost:8088")
    ap.add_argument("--k", type=int, default=5)
    ap.add_argument("--limit", type=int, default=300, help="how many events were ingested")
    args = ap.parse_args()
    hits, rr = 0, 0.0
    print("| question | relevant | first relevant rank |\n|---|---|---|")
    for question, right in labeled(args.limit):
        found = httpx.post(f"{args.url}/assistant/search", json={"query": question, "k": args.k}, timeout=60).json()
        ranks = [i + 1 for i, h in enumerate(found) if h["id"] in right]
        rank = ranks[0] if ranks else None
        hits += rank is not None
        rr += 1 / rank if rank else 0
        print(f"| {question} | {len(right)} | {rank or '-'} |")
    n = len(labeled(args.limit))
    print(f"\nhit@{args.k} = {hits / n:.2f}   MRR = {rr / n:.2f}   ({n} questions)")


if __name__ == "__main__":
    main()