自分のコードベースを Gemini 埋め込みと sqlite-vec で意味検索する — git 差分だけを再インデックスする増分パイプライン

import ast
from pathlib import Path
 
def chunk_python(path: str):
    src = Path(path).read_text(encoding="utf-8")
    try:
        tree = ast.parse(src)
    except SyntaxError:
        # 構文エラーのファイルは先頭を1チャンクとして退避
        return [(path, 1, src[:2000])]
    lines = src.splitlines()
    chunks = []
    for node in ast.walk(tree):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
            start = node.lineno
            end = getattr(node, "end_lineno", start)
            body = "\n".join(lines[start - 1:end])
            chunks.append((path, start, body))
    return chunks or [(path, 1, src[:2000])]

import numpy as np
from google import genai
 
client = genai.Client()
MODEL = "gemini-embedding-001"
DIM = 768
 
def normalize(v):
    a = np.array(v, dtype="float32")
    return (a / np.linalg.norm(a)).tolist()
 
def embed_batch(texts, task_type):
    vectors = []
    # 100件ずつに分割してリクエストサイズの上限を避ける
    for i in range(0, len(texts), 100):
        batch = texts[i:i + 100]
        resp = client.models.embed_content(
            model=MODEL,
            contents=batch,
            config={"task_type": task_type, "output_dimensionality": DIM},
        )
        vectors.extend(normalize(e.values) for e in resp.embeddings)
    return vectors

import sqlite3, sqlite_vec, struct
 
def serialize_f32(vec):
    return struct.pack(f"{len(vec)}f", *vec)
 
def open_db(path="code_index.db"):
    db = sqlite3.connect(path)
    db.enable_load_extension(True)
    sqlite_vec.load(db)
    db.enable_load_extension(False)
    db.execute("""
      CREATE TABLE IF NOT EXISTS chunks(
        id INTEGER PRIMARY KEY,
        file TEXT, line INTEGER, blob TEXT, body TEXT)""")
    db.execute("""
      CREATE VIRTUAL TABLE IF NOT EXISTS vec_chunks
      USING vec0(id INTEGER PRIMARY KEY, embedding FLOAT[768])""")
    return db

import subprocess
 
def blob_hash(path: str) -> str:
    return subprocess.run(
        ["git", "hash-object", path],
        capture_output=True, text=True).stdout.strip()
 
def changed_files(db, files):
    known = dict(db.execute(
        "SELECT file, blob FROM chunks GROUP BY file").fetchall())
    changed, unchanged = [], 0
    for f in files:
        h = blob_hash(f)
        if known.get(f) != h:
            changed.append((f, h))
        else:
            unchanged += 1
    return changed, unchanged

def reindex(db, changed):
    for f, h in changed:
        ids = [r[0] for r in db.execute(
            "SELECT id FROM chunks WHERE file = ?", (f,)).fetchall()]
        for i in ids:
            db.execute("DELETE FROM vec_chunks WHERE id = ?", (i,))
        db.execute("DELETE FROM chunks WHERE file = ?", (f,))
        pieces = chunk_python(f)
        vectors = embed_batch([p[2] for p in pieces], "RETRIEVAL_DOCUMENT")
        for (path, line, body), vec in zip(pieces, vectors):
            cur = db.execute(
                "INSERT INTO chunks(file, line, blob, body) VALUES (?,?,?,?)",
                (path, line, h, body))
            db.execute("INSERT INTO vec_chunks(id, embedding) VALUES (?, ?)",
                       (cur.lastrowid, serialize_f32(vec)))
    db.commit()

def search(db, query, k=8):
    qv = embed_batch([query], "RETRIEVAL_QUERY")[0]
    hits = db.execute("""
      SELECT id, distance FROM vec_chunks
      WHERE embedding MATCH ? AND k = ?
      ORDER BY distance""", (serialize_f32(qv), k)).fetchall()
    out = []
    for cid, dist in hits:
        f, line, body = db.execute(
            "SELECT file, line, body FROM chunks WHERE id = ?",
            (cid,)).fetchone()
        out.append((f, line, dist, body[:200]))
    return out