knowledge base search

kb.pyexcerpt

def sync(conn: sqlite3.Connection, articles: list[Article],
         source: str) -> SyncResult:
    """Make the index match `articles`, re-embedding only new or changed articles."""
    existing = {aid: (h, title) for aid, h, title in conn.execute(
        "SELECT article_id, content_hash, title FROM articles")}
    incoming = {a.article_id: a for a in articles}
    res = SyncResult([], [], [], 0)

    with conn:
        for aid in existing.keys() - incoming.keys():
            conn.execute("DELETE FROM chunks WHERE article_id = ?", (aid,))
            conn.execute("DELETE FROM articles WHERE article_id = ?", (aid,))
            res.removed.append(existing[aid][1])
        for aid, art in incoming.items():
            h = art.content_hash
            if aid not in existing:
                res.added.append(art.title)
            elif existing[aid][0] == h:
                res.unchanged += 1
                continue
            else:
                conn.execute("DELETE FROM chunks WHERE article_id = ?", (aid,))
                res.updated.append(art.title)
            row = {**asdict(art), "content_hash": h}
            conn.execute(f"INSERT OR REPLACE INTO articles ({', '.join(row)}) "
                         f"VALUES ({', '.join(':' + c for c in row)})", row)
            conn.executemany(
                "INSERT INTO chunks (article_id, section, text, vector) VALUES (?,?,?,?)",
                [(aid, sec, txt, embed(txt).tobytes()) for sec, txt in art.chunks()])
        conn.execute(
            "INSERT INTO sync_log (source, added, updated, removed, unchanged) "
            "VALUES (?,?,?,?,?)",
            (source, len(res.added), len(res.updated), len(res.removed), res.unchanged))
    return res


def search(conn: sqlite3.Connection, query: str, k: int = 3) -> list[Hit]:
    """Top k articles: an exact error code first, then by best chunk similarity."""
    qvec = embed(query)
    codes = {re.sub(r"^[Ee]-?", "E-", c) for c in ERROR_CODE_RE.findall(query)}

    # brute force; fine for a few hundred articles (see README for timings)
    best: dict[int, tuple[float, str]] = {}
    for aid, section, blob in conn.execute(
            "SELECT article_id, section, vector FROM chunks"):
        vec = array.array("f")
        vec.frombytes(blob)
        score = cosine(qvec, vec)
        if aid not in best or score > best[aid][0]:
            best[aid] = (score, section)

    if codes:  # exact error-code match beats similarity
        marks = ",".join("?" * len(codes))
        for (aid,) in conn.execute(
                f"SELECT article_id FROM articles WHERE error_code IN ({marks})",
                sorted(codes)):
            best[aid] = (1.0, "error code")

    top = sorted(best.items(), key=lambda item: item[1][0], reverse=True)[:k]
    if not top:
        return []
    rows = {r[0]: r for r in conn.execute(
        f"SELECT article_id, title, product, category, error_code, body "
        f"FROM articles WHERE article_id IN ({','.join('?' * len(top))})",
        [aid for aid, _ in top])}
    return [Hit(*rows[aid], score=max(score, 0.0), matched=section)
            for aid, (score, section) in top]