Hybrid Search

# script_id: day_027_hybrid_search/fusion
# fusion.py: reciprocal rank fusion (RRF). Pure Python: no installs, no API key.
from typing import Sequence


def rrf(rankings: Sequence[Sequence[str]], k: int = 60) -> list[tuple[str, float]]:
    """Each ranking votes 1/(k + rank) for every id in it; the highest total wins."""
    scores: dict[str, float] = {}
    for ranking in rankings:
        for rank, item in enumerate(ranking, start=1):
            scores[item] = scores.get(item, 0.0) + 1.0 / (k + rank)
    return sorted(scores.items(), key=lambda kv: (-kv[1], kv[0]))  # ties break by id, so runs repeat


if __name__ == "__main__":
    vector = ["concepts-jobs", "configuration-reference", "results-backend", "flows-steps-and-state"]
    keyword = ["error-codes", "configuration-reference", "broker-redis", "concepts-jobs"]
    for doc_id, score in rrf([vector, keyword]):
        print(f"{score:.5f}  {doc_id}")

0.03226  configuration-reference
0.03202  concepts-jobs
0.01639  error-codes
0.01587  broker-redis
0.01587  results-backend
0.01562  flows-steps-and-state

# script_id: day_027_hybrid_search/keyword_sql
for text in ["ORB-1202 PayloadTooLarge", "heartbeat_seconds", "orbit.toml", "X-Forwarded-Proto"]:
    lexemes = conn.execute("SELECT to_tsvector('english', %s)::text", (text,)).fetchone()[0]
    print(f"{text!r:28} -> {lexemes}")
query = conn.execute("SELECT websearch_to_tsquery('english', %s)::text", ("what does ORB-1202 mean?",))
print("websearch:", query.fetchone()[0])
try:
    conn.execute("SELECT to_tsquery('english', %s)", ("lease seconds",))
except psycopg.errors.SyntaxError as err:
    print("to_tsquery:", err)
score = conn.execute("SELECT 1.0 / (60 + 1) + 1.0 / (60 + 3)").fetchone()[0]
try:
    json.dumps({"score": score})
except TypeError as err:
    print(f"{type(score).__name__}: {err}")
print(json.dumps({"score": float(score)}))

'ORB-1202 PayloadTooLarge'   -> '-1202':2 'orb':1 'payloadtoolarg':3
'heartbeat_seconds'          -> 'heartbeat':1 'second':2
'orbit.toml'                 -> 'orbit.toml':1
'X-Forwarded-Proto'          -> 'forward':3 'proto':4 'x':2 'x-forwarded-proto':1
websearch: 'orb' <-> '-1202' & 'mean'
to_tsquery: syntax error in tsquery: "lease seconds"
Decimal: Object of type Decimal is not JSON serializable
{"score": 0.032266458495966696}

# script_id: day_027_hybrid_search/atlas_base
# atlas_base.py: compact copies of earlier days' helpers, so today's files stay self-contained.
# built on Day 15 (issues), Day 23 (embed, hit_rate), Day 24 (Document), Day 25 (chunk), Day 26 (Hit)
# Install: uv add "openai>=3.46,<4" numpy
# Env: ORBIT_DATA (path to datasets/orbit); OPENAI_API_KEY only when something calls embed().
import json
import os
import re
from dataclasses import dataclass, field
from functools import cache
from pathlib import Path
from typing import Callable, Sequence

import numpy as np
from openai import OpenAI

DATA_DIR = Path(os.environ.get("ORBIT_DATA", "datasets/orbit"))
EMBED_MODEL = os.environ.get("ATLAS_EMBED_MODEL", "text-embedding-3-small")
EMBED_DIM = 1536


@dataclass
class Document:  # Day 24's contract
    id: str
    title: str
    text: str
    source: str
    metadata: dict = field(default_factory=dict)


@dataclass
class DocChunk:  # Day 15 fields + Day 25 metadata; today metadata also carries source + allowed_groups
    id: str
    doc_id: str
    section: str
    applies_to: str
    text: str
    metadata: dict = field(default_factory=dict)


@dataclass
class Hit:  # Day 26: higher score = better
    chunk: DocChunk
    score: float

# script_id: day_027_hybrid_search/atlas_base
def parse_front_matter(raw: str) -> tuple[dict[str, str], str]:
    """Flat `key: value` front matter only (Day 24 uses a YAML parser)."""
    meta, body = {}, raw
    if raw.startswith("---\n"):
        head, _, body = raw[4:].partition("\n---\n")
        for line in head.splitlines():
            key, _, value = line.partition(":")
            meta[key.strip()] = value.split("#")[0].strip().strip('"')
    return meta, body


def load_docs(root: Path = DATA_DIR) -> list[Document]:
    docs = []
    for path in sorted((root / "docs").glob("*.md")):
        meta, body = parse_front_matter(path.read_text(encoding="utf-8"))
        docs.append(Document(meta.get("id", path.stem), meta.get("title", path.stem), body, "docs", meta))
    return docs


def load_issues(root: Path = DATA_DIR) -> list[dict]:
    """Day 15 returns `Issue` models; plain dicts keep this copy short."""
    lines = (root / "issues" / "issues.jsonl").read_text(encoding="utf-8").splitlines()
    return [json.loads(line) for line in lines if line.strip()]


def chunk(doc: Document, max_chars: int = 6000) -> list[DocChunk]:
    """One chunk per `## ` section, title + heading as a header, long sections cut by characters."""
    chunks: list[DocChunk] = []
    for part in (p.strip() for p in re.split(r"(?m)^(?=## )", doc.text) if p.strip()):
        heading = part.splitlines()[0].lstrip("# ")
        for start in range(0, len(part), max_chars):
            text = f"{doc.title} > {heading}\n{part[start:start + max_chars]}"
            chunks.append(DocChunk(f"{doc.id}#{len(chunks)}", doc.id, doc.metadata.get("section", ""),
                                   doc.metadata.get("applies_to", "all"), text,
                                   {**doc.metadata, "source": doc.source}))
    return chunks

# script_id: day_027_hybrid_search/atlas_base
@cache
def _client() -> OpenAI:
    return OpenAI()  # reads OPENAI_API_KEY from the environment


def embed(texts: Sequence[str], *, model: str = EMBED_MODEL, batch_size: int = 100) -> np.ndarray:
    """Day 23's embed() without its cache: float32 (n, d), unit length, input order."""
    vectors: list[list[float]] = []
    for start in range(0, len(texts), batch_size):
        resp = _client().embeddings.create(model=model, input=list(texts[start:start + batch_size]))
        vectors += [item.embedding for item in resp.data]
    return np.asarray(vectors, dtype=np.float32)  # OpenAI embeddings are unit length: dot = cosine


@dataclass
class RetrievalCase:
    id: str
    query: str
    relevant_doc_ids: list[str]
    kind: str


def load_retrieval_cases(split: str = "dev") -> list[RetrievalCase]:
    """Day 23's deterministic split: every other case by id within each kind."""
    lines = (DATA_DIR / "eval" / "retrieval_queries.jsonl").read_text(encoding="utf-8").splitlines()
    cases = sorted((RetrievalCase(**json.loads(line)) for line in lines if line.strip()), key=lambda c: c.id)
    if split == "all":
        return cases
    kinds = dict.fromkeys(c.kind for c in cases)
    return [c for kind in kinds for c in [x for x in cases if x.kind == kind][(split != "dev")::2]]


def ranked_doc_ids(hits: Sequence[Hit]) -> list[str]:
    return list(dict.fromkeys(h.chunk.doc_id for h in hits))  # Day 25: dedupe, keep order


def hit_rate(cases: Sequence[RetrievalCase], search_fn: Callable[[str, int], list[str]],
             k: int = 5) -> dict[str, float]:
    results: dict[str, list[bool]] = {}
    for case in cases:
        found = bool(set(search_fn(case.query, k)[:k]) & set(case.relevant_doc_ids))
        for key in ("overall", case.kind):
            results.setdefault(key, []).append(found)
    return {key: sum(v) / len(v) for key, v in results.items()}

# script_id: day_027_hybrid_search/acl
# acl.py: metadata + permission filters, defined once and used by both arms of both indexes.
# Install: nothing beyond atlas_base. Env: ORBIT_DATA.
from typing import Sequence

from atlas_base import Document, DocChunk, chunk, load_docs, load_issues

FILTER_KEYS = ("applies_to", "doc_id", "source")  # allowlist: these keys become SQL column names
Filters = dict[str, str | list[str]]  # str = equals, list = any of


def acl_for(doc: Document) -> list[str]:
    # Issues quote reporter emails, hosts and pasted tokens, so only maintainers may read them.
    # An unknown source gets no groups: invisible to everyone until someone decides (fail closed).
    return {"docs": ["public"], "issue": ["maintainers"]}.get(doc.source, [])


def issue_to_document(issue: dict) -> Document:
    comments = "\n\n".join(f"{c['author']}: {c['body']}" for c in issue.get("comments", []))
    version = issue.get("orbit_version") or ""
    meta = {"labels": ",".join(issue["labels"]), "state": issue["state"],
            "applies_to": f"{version.split('.')[0]}.x" if version else "all"}
    return Document(f"issue-{issue['id']}", issue["title"],
                    f"{issue['title']}\n\n{issue['body']}\n\n{comments}", "issue", meta)


def load_corpus(include_issues: bool = False) -> list[DocChunk]:
    docs = load_docs() + ([issue_to_document(i) for i in load_issues()] if include_issues else [])
    for doc in docs:
        doc.metadata["allowed_groups"] = acl_for(doc)  # decided at ingest, stored on every chunk
    return [c for doc in docs for c in chunk(doc)]

# script_id: day_027_hybrid_search/acl
def _check(filters: Filters | None, groups: Sequence[str]) -> Filters:
    if not groups or isinstance(groups, str):
        raise ValueError("groups must be a non-empty sequence such as ('public',)")  # fail closed
    unknown = set(filters or {}) - set(FILTER_KEYS)
    if unknown:
        raise ValueError(f"unknown filter keys {sorted(unknown)}; allowed: {FILTER_KEYS}")
    return filters or {}


def where_clause(filters: Filters | None, groups: Sequence[str]) -> tuple[str, dict]:
    """SQL version: one fragment + params, pasted into BOTH arms before their LIMIT."""
    filters = _check(filters, groups)
    parts, params = ["allowed_groups && %(acl_groups)s::text[]"], {"acl_groups": list(groups)}
    for i, (key, value) in enumerate(sorted(filters.items())):
        if isinstance(value, str):
            parts.append(f"{key} = %(f{i})s")
            params[f"f{i}"] = value
        else:
            parts.append(f"{key} = ANY(%(f{i})s::text[])")
            params[f"f{i}"] = list(value)
    return " AND ".join(parts), params


def matches(chunk: DocChunk, filters: Filters | None, groups: Sequence[str]) -> bool:
    """In-memory version of the same rule, for MemoryIndex."""
    filters = _check(filters, groups)
    if not set(chunk.metadata.get("allowed_groups", [])) & set(groups):
        return False  # a row with no groups is visible to nobody
    values = {"applies_to": chunk.applies_to, "doc_id": chunk.doc_id, "source": chunk.metadata.get("source")}
    return all(values[k] == v if isinstance(v, str) else values[k] in v for k, v in filters.items())

# script_id: day_027_hybrid_search/index
# index.py: Day 26's pgvector Index (compact) + today's keyword arm, RRF and ACL filters.
# Install: uv add "psycopg[binary]" pgvector numpy "openai>=3.46,<4"
# Env: DATABASE_URL, OPENAI_API_KEY, ORBIT_DATA. Needs Postgres + pgvector (Day 26's compose.yaml).
import os
from typing import Callable, Sequence

import numpy as np
import psycopg
from pgvector.psycopg import register_vector

from acl import Filters, load_corpus, where_clause
from atlas_base import EMBED_DIM, DocChunk, Hit, embed, hit_rate, load_retrieval_cases, ranked_doc_ids

DATABASE_URL = os.environ.get("DATABASE_URL", "postgresql://atlas:atlas@localhost:5432/atlas")
SCHEMA = """CREATE TABLE IF NOT EXISTS {t} (id text PRIMARY KEY, doc_id text NOT NULL, section text,
    applies_to text, text text NOT NULL, embedding vector({dim}));
CREATE INDEX IF NOT EXISTS {t}_hnsw ON {t} USING hnsw (embedding vector_cosine_ops);"""
MIGRATION = """
ALTER TABLE {t} ADD COLUMN IF NOT EXISTS tsv tsvector
    GENERATED ALWAYS AS (to_tsvector('english', text)) STORED;
ALTER TABLE {t} ADD COLUMN IF NOT EXISTS source text NOT NULL DEFAULT 'docs';
ALTER TABLE {t} ADD COLUMN IF NOT EXISTS allowed_groups text[] NOT NULL DEFAULT '{{}}';
CREATE INDEX IF NOT EXISTS {t}_tsv_gin ON {t} USING gin (tsv);
CREATE INDEX IF NOT EXISTS {t}_acl_gin ON {t} USING gin (allowed_groups);
UPDATE {t} SET allowed_groups = '{{public}}' WHERE source = 'docs' AND allowed_groups = '{{}}';"""
COLS = "id, doc_id, section, applies_to, text, source, allowed_groups"


def connect(url: str = DATABASE_URL) -> psycopg.Connection:
    conn = psycopg.connect(url, autocommit=True)
    conn.execute("CREATE EXTENSION IF NOT EXISTS vector")  # must run before register_vector
    register_vector(conn)
    return conn

# script_id: day_027_hybrid_search/index
class Index:
    def __init__(self, conn: psycopg.Connection, *, table: str = "chunks", dim: int = EMBED_DIM,
                 embed_fn: Callable[[Sequence[str]], np.ndarray] = embed):
        if not table.isidentifier():
            raise ValueError(f"bad table name {table!r}")  # it is formatted into SQL below
        self.conn, self.t, self.dim, self.embed_fn = conn, table, dim, embed_fn

    def create(self) -> None:
        self.conn.execute(SCHEMA.format(t=self.t, dim=self.dim))
        self.conn.execute(MIGRATION.format(t=self.t))  # today: tsv, source, allowed_groups

    def upsert(self, chunks: Sequence[DocChunk]) -> int:
        """Compact Day 26 upsert. Day 26's version embeds only new content-hash ids."""
        vectors = self.embed_fn([c.text for c in chunks])
        rows = [(c.id, c.doc_id, c.section, c.applies_to, c.text, v, c.metadata["source"],
                 c.metadata.get("allowed_groups", [])) for c, v in zip(chunks, vectors)]
        with self.conn.cursor() as cur:
            cur.executemany(f"""INSERT INTO {self.t}
                (id, doc_id, section, applies_to, text, embedding, source, allowed_groups)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
                ON CONFLICT (id) DO UPDATE SET doc_id = EXCLUDED.doc_id, section = EXCLUDED.section,
                  applies_to = EXCLUDED.applies_to, text = EXCLUDED.text, embedding = EXCLUDED.embedding,
                  source = EXCLUDED.source, allowed_groups = EXCLUDED.allowed_groups""", rows)
        return len(rows)  # tsv is GENERATED, so the update above recomputes it

    def _hits(self, sql: str, params: dict) -> list[Hit]:
        return [Hit(DocChunk(r[0], r[1], r[2], r[3], r[4], {"source": r[5], "allowed_groups": r[6]}),
                    float(r[7])) for r in self.conn.execute(sql, params).fetchall()]

# script_id: day_027_hybrid_search/index
    def search(self, query: str, k: int = 5, *, filters: Filters | None = None,
               groups: Sequence[str] = ("public",)) -> list[Hit]:
        where, params = where_clause(filters, groups)
        return self._hits(f"""SELECT {COLS}, 1 - (embedding <=> %(q)s) FROM {self.t}
            WHERE {where} ORDER BY embedding <=> %(q)s LIMIT %(k)s""",
            {**params, "q": self.embed_fn([query])[0], "k": k})

    def keyword_search(self, query: str, k: int = 5, *, filters: Filters | None = None,
                       groups: Sequence[str] = ("public",)) -> list[Hit]:
        where, params = where_clause(filters, groups)
        return self._hits(f"""SELECT {COLS}, ts_rank_cd(tsv, query)
            FROM {self.t}, websearch_to_tsquery('english', %(text)s) AS query
            WHERE tsv @@ query AND {where} ORDER BY ts_rank_cd(tsv, query) DESC, id LIMIT %(k)s""",
            {**params, "text": query, "k": k})

# script_id: day_027_hybrid_search/index
    def hybrid_search(self, query: str, k: int = 5, *, filters: Filters | None = None,
                      groups: Sequence[str] = ("public",), candidates: int = 20,
                      k_rrf: int = 60) -> list[Hit]:
        where, params = where_clause(filters, groups)  # ONE fragment, pasted into both arms
        sql = f"""
        WITH semantic AS (
            SELECT id, RANK() OVER (ORDER BY dist) AS rank FROM (
                SELECT id, embedding <=> %(q)s AS dist FROM {self.t}
                WHERE {where} ORDER BY embedding <=> %(q)s LIMIT %(n)s) AS nearest),
        keyword AS (
            SELECT id, RANK() OVER (ORDER BY ts_rank_cd(tsv, query) DESC) AS rank
            FROM {self.t}, websearch_to_tsquery('english', %(text)s) AS query
            WHERE tsv @@ query AND {where}
            ORDER BY ts_rank_cd(tsv, query) DESC LIMIT %(n)s),
        fused AS (
            SELECT COALESCE(s.id, kw.id) AS id,
                   COALESCE(1.0 / (%(k_rrf)s + s.rank), 0.0)
                 + COALESCE(1.0 / (%(k_rrf)s + kw.rank), 0.0) AS rrf
            FROM semantic s FULL OUTER JOIN keyword kw ON s.id = kw.id)
        SELECT {COLS}, rrf::float8 FROM fused JOIN {self.t} USING (id)
        ORDER BY rrf DESC, id LIMIT %(k)s"""
        params |= {"q": self.embed_fn([query])[0], "text": query, "n": candidates, "k_rrf": k_rrf, "k": k}
        return self._hits(sql, params)

# script_id: day_027_hybrid_search/index
if __name__ == "__main__":
    index = Index(connect())
    index.create()
    print(f"upserted {index.upsert(load_corpus())} doc chunks")  # issues join in the ACL demo
    cases = load_retrieval_cases("dev")
    arms = {"vector": index.search, "keyword": index.keyword_search, "hybrid": index.hybrid_search}
    rates = {name: hit_rate(cases, lambda q, k, fn=fn: ranked_doc_ids(fn(q, 20))[:k])
             for name, fn in arms.items()}
    print(f"{'hit@5':11}" + "".join(f"{name:>9}" for name in arms))
    for kind in ("exact_term", "semantic", "multi_hop", "overall"):
        print(f"{kind:11}" + "".join(f"{rates[name][kind]:9.2f}" for name in arms))

upserted 349 doc chunks
hit@5         vector  keyword   hybrid
exact_term      0.xx     0.xx     0.xx
semantic        0.xx     0.xx     0.xx
multi_hop       0.xx     0.xx     0.xx
overall         0.xx     0.xx     0.xx

# script_id: day_027_hybrid_search/memory_index
# memory_index.py: no-Postgres fallback. Stdlib BM25 keyword arm + numpy vector arm + RRF + ACL filters.
# Install: uv add numpy "openai>=3.46,<4". Env: ORBIT_DATA; OPENAI_API_KEY optional (no key = keyword-only).
import math
import os
import re
from collections import Counter
from typing import Callable, Sequence

import numpy as np

from acl import Filters, load_corpus, matches
from atlas_base import DocChunk, Hit, embed, hit_rate, load_retrieval_cases, ranked_doc_ids
from fusion import rrf


def tokenize(text: str) -> list[str]:
    """Keep `orb-1202`, `orbit.toml`, `max_payload_kb` whole AND add their parts."""
    tokens: list[str] = []
    for word in re.findall(r"\w+(?:[-.]\w+)*", text.lower()):
        parts = [p for p in re.split(r"[-._]", word) if p]
        tokens += [word] + (parts if len(parts) > 1 else [])
    return tokens


class BM25:
    """Okapi BM25 with Lucene-style IDF (never negative); k1 and b are the usual defaults."""

    def __init__(self, texts: Sequence[str], k1: float = 1.5, b: float = 0.75):
        self.docs = [Counter(tokenize(t)) for t in texts]
        self.lens = [sum(d.values()) for d in self.docs]
        self.avg = sum(self.lens) / max(len(self.lens), 1)
        df = Counter(term for d in self.docs for term in d)
        n = len(self.docs)
        self.idf = {t: math.log(1 + (n - f + 0.5) / (f + 0.5)) for t, f in df.items()}
        self.k1, self.b = k1, b

    def scores(self, query: str) -> list[float]:
        terms = set(tokenize(query))
        return [sum(self.idf[t] * d[t] * (self.k1 + 1)
                    / (d[t] + self.k1 * (1 - self.b + self.b * n / self.avg)) for t in terms & d.keys())
                for d, n in zip(self.docs, self.lens)]

# script_id: day_027_hybrid_search/memory_index
Embedder = Callable[[Sequence[str]], np.ndarray]


class MemoryIndex:
    """Same search API as the Postgres Index. embed_fn=None gives a keyword-only index."""

    def __init__(self, embed_fn: Embedder | None = embed):
        self.embed_fn = embed_fn
        self.chunks: dict[str, DocChunk] = {}
        self.vectors: dict[str, np.ndarray] = {}
        self.rows: list[DocChunk] = []
        self.bm25 = BM25([])

    def upsert(self, chunks: Sequence[DocChunk]) -> int:
        changed = [c for c in chunks if c.id not in self.chunks or self.chunks[c.id].text != c.text]
        if self.embed_fn and changed:  # embed only new or edited text
            self.vectors.update(zip((c.id for c in changed), self.embed_fn([c.text for c in changed])))
        self.chunks.update({c.id: c for c in chunks})
        self.rows = list(self.chunks.values())
        self.bm25 = BM25([c.text for c in self.rows])  # IDF is corpus-wide, so rebuild after writes
        return len(chunks)

    def search(self, query: str, k: int = 5, *, filters: Filters | None = None,
               groups: Sequence[str] = ("public",)) -> list[Hit]:
        allowed = [c for c in self.rows if matches(c, filters, groups)]  # filter BEFORE ranking
        if self.embed_fn is None or not allowed:
            return []
        sims = np.stack([self.vectors[c.id] for c in allowed]) @ self.embed_fn([query])[0]
        return [Hit(allowed[i], float(sims[i])) for i in np.argsort(-sims)[:k]]

    def keyword_search(self, query: str, k: int = 5, *, filters: Filters | None = None,
                       groups: Sequence[str] = ("public",)) -> list[Hit]:
        scores = self.bm25.scores(query)
        ranked = sorted((i for i, c in enumerate(self.rows) if matches(c, filters, groups) and scores[i] > 0),
                        key=lambda i: -scores[i])
        return [Hit(self.rows[i], scores[i]) for i in ranked[:k]]

    def hybrid_search(self, query: str, k: int = 5, *, filters: Filters | None = None,
                      groups: Sequence[str] = ("public",), candidates: int = 20, k_rrf: int = 60) -> list[Hit]:
        arms = [self.search(query, candidates, filters=filters, groups=groups),
                self.keyword_search(query, candidates, filters=filters, groups=groups)]
        fused = rrf([[h.chunk.id for h in arm] for arm in arms], k=k_rrf)
        return [Hit(self.chunks[cid], score) for cid, score in fused[:k]]

# script_id: day_027_hybrid_search/memory_index
if __name__ == "__main__":
    use_vectors = bool(os.environ.get("OPENAI_API_KEY"))
    index = MemoryIndex(embed if use_vectors else None)
    index.upsert(load_corpus())
    arms: dict[str, Callable] = {"keyword": index.keyword_search}
    if use_vectors:
        arms |= {"vector": index.search, "hybrid": index.hybrid_search}
    else:
        print("OPENAI_API_KEY not set: BM25 keyword arm only")
    cases = load_retrieval_cases("dev")
    rates = {n: hit_rate(cases, lambda q, k, fn=fn: ranked_doc_ids(fn(q, 20))[:k]) for n, fn in arms.items()}
    print(f"{len(index.rows)} chunks · {len(cases)} dev cases")
    print(f"{'hit@5':11}" + "".join(f"{n:>9}" for n in rates))
    for kind in ("exact_term", "semantic", "multi_hop", "overall"):
        print(f"{kind:11}" + "".join(f"{r[kind]:9.2f}" for r in rates.values()))

OPENAI_API_KEY not set: BM25 keyword arm only
349 chunks · 20 dev cases
hit@5        keyword
exact_term      1.00
semantic        0.86
multi_hop       1.00
overall         0.95

# script_id: day_027_hybrid_search/acl_demo
# Two users, one query, and a leak count. Keyword-only without OPENAI_API_KEY.
# Env: ORBIT_DATA; OPENAI_API_KEY optional (embedding 240 issues ≈ $0.01, as of 2026-11, verify).
import os

from acl import load_corpus
from atlas_base import embed, load_retrieval_cases, ranked_doc_ids
from memory_index import MemoryIndex

index = MemoryIndex(embed if os.environ.get("OPENAI_API_KEY") else None)
index.upsert(load_corpus(include_issues=True))
USERS = {"anonymous": ("public",), "maintainer": ("public", "maintainers")}
QUERY = "redis reconnect storm"
for user, groups in USERS.items():
    pages = ranked_doc_ids(index.hybrid_search(QUERY, k=10, groups=groups))[:5]
    print(f"{user:10} {pages}")

# Leak test: no hit may carry only groups the caller lacks.
queries = [c.query for c in load_retrieval_cases("all")[:20]] + [QUERY]
anon = USERS["anonymous"]
leaks = sum(1 for q in queries for h in index.hybrid_search(q, k=10, groups=anon)
            if not set(h.chunk.metadata["allowed_groups"]) & set(anon))
print(f"anonymous leak test over {len(queries)} queries: leaks = {leaks}")

# Anti-pattern: rank everything, then filter. The caller asked for 5 rows.
top = index.hybrid_search(QUERY, k=5, groups=USERS["maintainer"])
visible = [h for h in top if "public" in h.chunk.metadata["allowed_groups"]]
print(f"filter after fusion: {len(visible)} of 5 rows left for an anonymous user")
try:
    index.hybrid_search(QUERY, groups=())
except ValueError as err:
    print("empty groups:", err)

anonymous  ['troubleshooting-broker', 'broker-redis', 'changelog', 'deployment-kubernetes', 'contributing']
maintainer ['troubleshooting-broker', 'broker-redis', 'changelog', 'issue-143', 'issue-161']
anonymous leak test over 21 queries: leaks = 0
filter after fusion: 4 of 5 rows left for an anonymous user
empty groups: groups must be a non-empty sequence such as ('public',)