Semantic search over your own files with sqlite-vec

#!/usr/bin/env python3
# busca.py - busca por similaridade no seu acervo, com sqlite-vec quando existe.
#
# Uso:
#   python3 busca.py                demo com documentos de exemplo
#   python3 busca.py /minha/pasta   indexa os .txt/.md da pasta e responde 3 consultas
#
# O vetor de cada texto sai de um hash de trigramas de caractere (sem modelo, sem
# dependencia). Com o sqlite-vec instalado, os vetores ficam numa tabela vec0 e a
# busca e por KNN dentro do SQLite. Sem ele, o mesmo resultado sai de uma varredura
# em Python puro. Troque o embed() por um modelo de verdade e o resto nao muda.
import hashlib
import math
import os
import sqlite3
import struct
import sys

DIM = 256
TOP = 3

def embed(texto):
    texto = texto.lower()
    v = [0.0] * DIM
    for i in range(len(texto) - 2):
        h = int.from_bytes(hashlib.md5(texto[i:i + 3].encode("utf-8")).digest()[:4], "little")
        v[h % DIM] += 1.0
    norma = math.sqrt(sum(x * x for x in v))
    if norma == 0.0:
        return v
    return [x / norma for x in v]

def pack(vec):
    return struct.pack("%df" % DIM, *vec)

def carregar_vec(conn):
    try:
        conn.enable_load_extension(True)
    except Exception:
        return False
    try:
        so = os.environ.get("SQLITE_VEC_SO", "")
        if so and os.path.exists(so):
            conn.load_extension(so)
            return True
        import sqlite_vec
        sqlite_vec.load(conn)
        return True
    except Exception:
        return False
    finally:
        try:
            conn.enable_load_extension(False)
        except Exception:
            pass

DOCS = [
    ("backup.txt", "como fazer backup com rsync, guardar versoes com hardlink e testar a restauracao antes de precisar dela"),
    ("busca.txt", "indice de vetores dentro do proprio sqlite para busca por similaridade e consulta knn"),
    ("docker.txt", "rodar aplicacoes em contenedores isolados, com imagem, volume e rede propria"),
    ("rede.txt", "tunel seguro para expor um servico local sem abrir porta no roteador nem liberar firewall"),
    ("dados.txt", "consulta sql para agregar milhoes de linhas, montar relatorio e contar ocorrencias"),
]

def trecho(texto, n=64):
    return texto if len(texto) <= n else texto[:n] + "..."

def indexar(docs):
    conn = sqlite3.connect(":memory:")
    tem_vec = carregar_vec(conn)
    vetores = [embed(t) for _c, t in docs]
    meta = [(c, trecho(t)) for c, t in docs]
    if tem_vec:
        conn.execute("create virtual table docs using vec0(e float[256] distance_metric=cosine)")
        for i, v in enumerate(vetores):
            conn.execute("insert into docs(rowid, e) values (?, ?)", (i + 1, pack(v)))
    return conn, meta, vetores, tem_vec

def consultar(conn, meta, vetores, tem_vec, consulta):
    q = embed(consulta)
    if tem_vec:
        rows = conn.execute(
            "select rowid, distance from docs where e match ? and k = ? order by distance",
            (pack(q), TOP),
        ).fetchall()
        return [(meta[r - 1][0], meta[r - 1][1], d) for r, d in rows]
    pares = []
    for i, v in enumerate(vetores):
        dist = 1.0 - sum(a * b for a, b in zip(q, v))
        pares.append((dist, i))
    pares.sort()
    return [(meta[i][0], meta[i][1], d) for d, i in pares[:TOP]]

def main():
    if len(sys.argv) >= 2:
        pasta = sys.argv[1]
        nomes = sorted(n for n in os.listdir(pasta) if n.endswith((".txt", ".md")))
        if not nomes:
            print("nenhum .txt/.md em", pasta)
            return 1
        docs = []
        for nome in nomes:
            with open(os.path.join(pasta, nome), encoding="utf-8", errors="replace") as f:
                docs.append((nome, f.read()))
        print("indexados %d documentos de %s" % (len(docs), pasta))
    else:
        docs = DOCS
        print("sem argumentos: demo com %d documentos de exemplo" % len(docs))

    conn, meta, vetores, tem_vec = indexar(docs)
    print("motor:", "sqlite-vec" if tem_vec else "python puro (sqlite-vec nao encontrado)")

    consultas = [
        "como fazer backup e restaurar arquivos",
        "buscar documentos por similaridade",
        "expor um servico com tunel seguro",
    ]
    for c in consultas:
        print()
        print("consulta: %s" % c)
        for caminho, trecho_, d in consultar(conn, meta, vetores, tem_vec, c):
            print("  %.3f  %-14s %s" % (d, caminho, trecho_))
    return 0

if __name__ == "__main__":
    sys.exit(main())

motor: sqlite-vec

consulta: como fazer backup e restaurar arquivos
  0.430  backup.txt     como fazer backup com rsync, guardar versoes...

consulta: buscar documentos por similaridade
  0.514  busca.txt      indice de vetores dentro do proprio sqlite...

consulta: expor um servico com tunel seguro
  0.437  rede.txt       tunel seguro para expor um servico local...