Gemini API × sqlite-vec で作る超軽量 RAG — 個人開発向けの本番設計と運用パターン

# app/indexer.py
# 役割: ノート文書を Gemini で埋め込み、sqlite-vec に保存する。
# 依存: pip install google-genai sqlite-vec numpy python-dotenv
 
import os
import sqlite3
import struct
from pathlib import Path
 
import sqlite_vec
from google import genai
from dotenv import load_dotenv
 
load_dotenv()
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
 
DB_PATH = "notes.db"
EMBED_MODEL = "text-embedding-001"
EMBED_DIM = 768  # text-embedding-001 の次元
 
def serialize(vec: list[float]) -> bytes:
    """sqlite-vec 用のバイナリ形式にシリアライズする。"""
    return struct.pack(f"{len(vec)}f", *vec)
 
def open_db(path: str = DB_PATH) -> sqlite3.Connection:
    """sqlite-vec 拡張をロードして接続を返す。"""
    con = sqlite3.connect(path)
    con.enable_load_extension(True)
    sqlite_vec.load(con)
    con.enable_load_extension(False)
    con.execute("PRAGMA journal_mode=WAL")  # 読み書き並行性を確保
    return con
 
def init_schema(con: sqlite3.Connection) -> None:
    """ノート本体と埋め込みテーブルを作成する。"""
    con.executescript(f"""
        CREATE TABLE IF NOT EXISTS notes (
            id INTEGER PRIMARY KEY,
            title TEXT NOT NULL,
            body  TEXT NOT NULL,
            updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
        );
 
        CREATE VIRTUAL TABLE IF NOT EXISTS note_vec USING vec0(
            id INTEGER PRIMARY KEY,
            embedding float[{EMBED_DIM}]
        );
    """)
 
def embed(text: str) -> list[float]:
    """Gemini の埋め込み API を呼び出す。空文字はゼロベクトルで回避。"""
    if not text.strip():
        return [0.0] * EMBED_DIM
    resp = client.models.embed_content(model=EMBED_MODEL, contents=text)
    return resp.embeddings[0].values
 
def upsert_note(con: sqlite3.Connection, title: str, body: str) -> int:
    """ノートと埋め込みを1トランザクションで書き込む。"""
    vec = embed(f"{title}\n{body}")
    with con:
        cur = con.execute(
            "INSERT INTO notes(title, body) VALUES (?, ?)", (title, body)
        )
        note_id = cur.lastrowid
        con.execute(
            "INSERT INTO note_vec(id, embedding) VALUES (?, ?)",
            (note_id, serialize(vec)),
        )
    return note_id
 
if __name__ == "__main__":
    con = open_db()
    init_schema(con)
    # サンプル投入
    upsert_note(con, "朝のランニング", "5km を 28 分。呼吸が落ち着いてきた。")
    upsert_note(con, "Gemini API メモ", "embed_content の次元は 768。正規化は不要。")
    print("OK")

# app/search.py
# 役割: 問い合わせ文を Gemini で埋め込み、類似するノートを返す。
 
import sqlite3
from typing import Any
from app.indexer import open_db, embed, serialize
 
def search(query: str, limit: int = 5) -> list[dict[str, Any]]:
    q_vec = embed(query)
    con = open_db()
    # vec0 テーブルに対する MATCH クエリ。distance は L2 距離。
    rows = con.execute(
        """
        SELECT n.id, n.title, n.body, v.distance
        FROM note_vec v
        JOIN notes n ON n.id = v.id
        WHERE v.embedding MATCH ? AND k = ?
        ORDER BY v.distance
        """,
        (serialize(q_vec), limit),
    ).fetchall()
    return [
        {"id": r[0], "title": r[1], "body": r[2], "distance": r[3]} for r in rows
    ]
 
if __name__ == "__main__":
    for hit in search("朝のトレーニング記録"):
        print(f"{hit['distance']:.3f}  {hit['title']}")

# app/chunking.py
from dataclasses import dataclass
 
@dataclass
class Chunk:
    text: str
    start: int
    end: int
 
def chunk_by_length(text: str, size: int = 500, overlap: int = 100) -> list[Chunk]:
    if not text:
        return []
    step = max(1, size - overlap)
    out: list[Chunk] = []
    i = 0
    while i < len(text):
        j = min(len(text), i + size)
        out.append(Chunk(text=text[i:j], start=i, end=j))
        if j == len(text):
            break
        i += step
    return out

CREATE TABLE IF NOT EXISTS chunks (
    id INTEGER PRIMARY KEY,
    note_id INTEGER NOT NULL REFERENCES notes(id) ON DELETE CASCADE,
    ord INTEGER NOT NULL,        -- ノート内の順序
    text TEXT NOT NULL,
    created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
);
CREATE INDEX IF NOT EXISTS chunks_note_id ON chunks(note_id);
 
CREATE VIRTUAL TABLE IF NOT EXISTS chunk_vec USING vec0(
    id INTEGER PRIMARY KEY,
    embedding float[768]
);

CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
    text,
    content='chunks',
    content_rowid='id',
    tokenize='trigram'
);
 
-- 日本語は tokenize='trigram' が実用上は最も安定します。
-- FTS5 と実テーブルを同期させるトリガーも合わせて仕掛けておきます。
CREATE TRIGGER IF NOT EXISTS chunks_ai AFTER INSERT ON chunks BEGIN
    INSERT INTO chunks_fts(rowid, text) VALUES (new.id, new.text);
END;
CREATE TRIGGER IF NOT EXISTS chunks_ad AFTER DELETE ON chunks BEGIN
    DELETE FROM chunks_fts WHERE rowid = old.id;
END;
CREATE TRIGGER IF NOT EXISTS chunks_au AFTER UPDATE ON chunks BEGIN
    DELETE FROM chunks_fts WHERE rowid = old.id;
    INSERT INTO chunks_fts(rowid, text) VALUES (new.id, new.text);
END;

# app/hybrid_search.py
from collections import defaultdict
from app.indexer import open_db, embed, serialize
 
def hybrid_search(query: str, k: int = 10, pool: int = 50) -> list[dict]:
    con = open_db()
    # 1) ベクトル側: 距離の昇順で上位 pool 件
    vec_rows = con.execute(
        """
        SELECT c.id, c.note_id, c.text, v.distance
        FROM chunk_vec v
        JOIN chunks c ON c.id = v.id
        WHERE v.embedding MATCH ? AND k = ?
        ORDER BY v.distance
        """,
        (serialize(embed(query)), pool),
    ).fetchall()
 
    # 2) FTS 側: BM25 スコアの降順で上位 pool 件
    fts_rows = con.execute(
        """
        SELECT c.id, c.note_id, c.text, bm25(chunks_fts) AS score
        FROM chunks_fts
        JOIN chunks c ON c.id = chunks_fts.rowid
        WHERE chunks_fts MATCH ?
        ORDER BY score
        LIMIT ?
        """,
        (query, pool),
    ).fetchall()
 
    # 3) 逆順位スコアで合算 (RRF, Reciprocal Rank Fusion)
    rr: dict[int, float] = defaultdict(float)
    meta: dict[int, dict] = {}
    for rank, row in enumerate(vec_rows):
        cid = row[0]
        rr[cid] += 1 / (60 + rank)
        meta[cid] = {"id": cid, "note_id": row[1], "text": row[2]}
    for rank, row in enumerate(fts_rows):
        cid = row[0]
        rr[cid] += 1 / (60 + rank)
        meta.setdefault(cid, {"id": cid, "note_id": row[1], "text": row[2]})
 
    ranked = sorted(rr.items(), key=lambda x: x[1], reverse=True)[:k]
    return [meta[c] | {"score": s} for c, s in ranked]

def tune(con: sqlite3.Connection) -> None:
    con.executescript("""
        PRAGMA journal_mode=WAL;
        PRAGMA synchronous=NORMAL;
        PRAGMA busy_timeout=5000;
        PRAGMA mmap_size=268435456;
        PRAGMA temp_store=MEMORY;
    """)

def backup(src_path: str = "notes.db", dst_path: str = "backup.db") -> None:
    con = open_db(src_path)
    con.execute(f"VACUUM INTO ?", (dst_path,))
    con.close()

def reindex(new_dim: int, new_table: str = "chunk_vec_v2") -> None:
    con = open_db()
    con.execute(f"CREATE VIRTUAL TABLE IF NOT EXISTS {new_table} USING vec0(id INTEGER PRIMARY KEY, embedding float[{new_dim}])")
    rows = con.execute("SELECT id, text FROM chunks").fetchall()
    for cid, text in rows:
        vec = embed(text)  # 新しいモデルで再計算
        con.execute(f"INSERT INTO {new_table}(id, embedding) VALUES (?, ?)", (cid, serialize(vec)))
    con.commit()

# eval/dataset.py
EVAL_CASES = [
    {"query": "3月にキャッシュの方針をどう決めたか", "expected_ids": [104]},
    {"query": "Stripe Webhook の署名エラー対処", "expected_ids": [218, 221]},
    {"query": "iPhone 壁紙アプリの売上急増", "expected_ids": [33]},
    # ... 合計 30〜60 ケース
]
 
def evaluate(cases, k=5):
    hits = 0
    rr = 0.0
    for case in cases:
        results = hybrid_search(case["query"], k=k)
        rank = None
        for i, r in enumerate(results, 1):
            if r["note_id"] in case["expected_ids"]:
                rank = i; break
        if rank:
            hits += 1
            rr += 1 / rank
    n = len(cases)
    return {"recall@k": hits/n, "mrr": rr/n}

import hashlib
 
def chunk_hash(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()
 
def upsert_chunk(con, note_id: int, ord: int, text: str) -> int:
    exists = con.execute(
        "SELECT id FROM chunks WHERE note_id=? AND ord=? AND text=?",
        (note_id, ord, text),
    ).fetchone()
    if exists:
        return exists[0]
    vec = embed(text)
    cur = con.execute(
        "INSERT INTO chunks(note_id, ord, text) VALUES (?, ?, ?)",
        (note_id, ord, text),
    )
    cid = cur.lastrowid
    con.execute("INSERT INTO chunk_vec(id, embedding) VALUES (?, ?)", (cid, serialize(vec)))
    return cid