Gemini API の動的 Few-Shot 設計 — ベクター検索で実例を選び続ける自己改善型プロンプト

# pip install google-genai numpy
import os
import json
import numpy as np
from google import genai
from google.genai import types
 
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
 
# 初期 Few-Shot 例（人手で用意した起点）
SEED_EXAMPLES = [
    {
        "input": "請求書を発行してください。宛先は株式会社サンプル様、件名は5月分支援費です。",
        "output": json.dumps({
            "intent": "invoice_create",
            "client": "株式会社サンプル",
            "subject": "5月分支援費",
        }, ensure_ascii=False),
        "tags": ["invoice", "form-fill"],
    },
    {
        "input": "今月の経費レポートを CSV でダウンロードしたい",
        "output": json.dumps({
            "intent": "expense_export",
            "format": "csv",
            "period": "this_month",
        }, ensure_ascii=False),
        "tags": ["expense", "export"],
    },
    {
        "input": "先月発行した請求書一覧を見せて",
        "output": json.dumps({
            "intent": "invoice_list",
            "period": "last_month",
        }, ensure_ascii=False),
        "tags": ["invoice", "list"],
    },
]
 
def embed_text(text: str) -> np.ndarray:
    """gemini-embedding-001 で 768次元ベクトルを取得"""
    resp = client.models.embed_content(
        model="gemini-embedding-001",
        contents=text,
        config=types.EmbedContentConfig(
            task_type="RETRIEVAL_DOCUMENT",
            output_dimensionality=768,
        ),
    )
    vec = np.array(resp.embeddings[0].values, dtype=np.float32)
    # 正規化しておくと後段のコサイン類似度が単純な内積で計算できる
    return vec / np.linalg.norm(vec)
 
def build_store(examples):
    """例ストアを構築する。ベクトルとメタデータを揃えて返す。"""
    vectors = np.stack([embed_text(ex["input"]) for ex in examples])
    return {"vectors": vectors, "examples": examples}
 
EXAMPLE_STORE = build_store(SEED_EXAMPLES)
print(f"[init] stored {len(EXAMPLE_STORE['examples'])} examples")
# 期待出力: [init] stored 3 examples

def embed_query(text: str) -> np.ndarray:
    resp = client.models.embed_content(
        model="gemini-embedding-001",
        contents=text,
        config=types.EmbedContentConfig(
            task_type="RETRIEVAL_QUERY",
            output_dimensionality=768,
        ),
    )
    vec = np.array(resp.embeddings[0].values, dtype=np.float32)
    return vec / np.linalg.norm(vec)
 
def select_examples(store, query: str, k: int = 3, min_score: float = 0.55):
    """類似度上位 K件を返す。min_score 未満の例は除外する。"""
    qvec = embed_query(query)
    scores = store["vectors"] @ qvec  # 正規化済みなので内積=コサイン類似度
    order = np.argsort(-scores)
    selected = []
    for idx in order[:k]:
        score = float(scores[idx])
        if score < min_score:
            break
        selected.append({**store["examples"][idx], "score": score})
    return selected
 
# 動作確認
candidates = select_examples(
    EXAMPLE_STORE,
    "前月の請求書をまとめて出してほしい",
    k=2,
)
for c in candidates:
    print(f"score={c['score']:.3f} input={c['input'][:30]}")
# 期待出力（順位は意味の近さで決まる）:
# score=0.812 input=先月発行した請求書一覧を見せて
# score=0.612 input=請求書を発行してください。宛先は…

def build_prompt(query: str, examples: list[dict]) -> str:
    blocks = []
    for ex in examples:
        blocks.append(f"# 入力\n{ex['input']}\n# 出力\n{ex['output']}")
    examples_text = "\n\n".join(blocks)
    return (
        "あなたは社内ツールの自然言語コマンドを JSON へ変換するアシスタントです。\n"
        "以下の例と同じ形式で、最後の入力に対する JSON を返してください。\n\n"
        f"{examples_text}\n\n"
        f"# 入力\n{query}\n# 出力\n"
    )
 
def classify(query: str, k: int = 3):
    examples = select_examples(EXAMPLE_STORE, query, k=k)
    if not examples:
        # 例が見つからない場合は zero-shot にフォールバック
        # ストアが薄い初期や未知ドメインで重要
        return _zero_shot(query)
    prompt = build_prompt(query, examples)
    response = client.models.generate_content(
        model="gemini-2.5-flash",
        contents=prompt,
        config=types.GenerateContentConfig(
            temperature=0.1,
            response_mime_type="application/json",
        ),
    )
    return {
        "result": json.loads(response.text),
        "selected_examples": [ex["input"] for ex in examples],
    }
 
def _zero_shot(query: str):
    response = client.models.generate_content(
        model="gemini-2.5-flash",
        contents=(
            "次の自然言語コマンドを JSON に変換してください。"
            "intent と必要なパラメータだけを含めてください。\n\n"
            f"入力: {query}"
        ),
        config=types.GenerateContentConfig(
            temperature=0.1,
            response_mime_type="application/json",
        ),
    )
    return {"result": json.loads(response.text), "selected_examples": []}
 
# 動作確認
out = classify("前月の請求書をまとめて出してほしい")
print(json.dumps(out, ensure_ascii=False, indent=2))
# 期待出力:
# {
#   "result": {"intent": "invoice_list", "period": "last_month"},
#   "selected_examples": ["先月発行した請求書一覧を見せて", "請求書を発行してください…"]
# }

import uuid
from datetime import datetime, timezone
 
# 推論結果と例の追加候補をひとまずバッファに溜める運用
PENDING_BUFFER = {}
 
def classify_with_logging(query: str, k: int = 3):
    examples = select_examples(EXAMPLE_STORE, query, k=k)
    prompt = build_prompt(query, examples) if examples else None
    response = client.models.generate_content(
        model="gemini-2.5-flash",
        contents=prompt or query,
        config=types.GenerateContentConfig(
            temperature=0.1,
            response_mime_type="application/json",
        ),
    )
    rec_id = str(uuid.uuid4())
    PENDING_BUFFER[rec_id] = {
        "input": query,
        "output": response.text,
        "selected_examples": [ex["input"] for ex in examples],
        "ts": datetime.now(timezone.utc).isoformat(),
    }
    return rec_id, json.loads(response.text)
 
def submit_feedback(rec_id: str, score: int, store=EXAMPLE_STORE):
    """score は 1（良い）/ 0（中立）/ -1（悪い）。1 のみ昇格候補にする"""
    rec = PENDING_BUFFER.get(rec_id)
    if rec is None:
        raise KeyError(f"unknown rec_id: {rec_id}")
    if score == 1:
        promote_to_store(store, rec)
    elif score == -1:
        # 失敗例は別ストアに送って後でキュレーションする
        log_failure(rec)
 
def promote_to_store(store, rec, dedup_threshold=0.93):
    """昇格時は重複に近い例を弾く（コサイン類似度しきい値）"""
    vec = embed_text(rec["input"])
    if len(store["vectors"]) > 0:
        sims = store["vectors"] @ vec
        if float(sims.max()) >= dedup_threshold:
            return False  # 既存例と意味が近すぎるので追加しない
    new_example = {
        "input": rec["input"],
        "output": rec["output"],
        "tags": ["promoted"],
    }
    store["vectors"] = np.vstack([store["vectors"], vec[np.newaxis, :]])
    store["examples"].append(new_example)
    return True
 
def log_failure(rec):
    # 実装はストレージに任せる。SQLite/BigQuery など
    print(f"[failure] {rec['input'][:40]}…")
 
# 動作確認
rec_id, result = classify_with_logging("当月の経費を CSV で取り出して")
print("classified:", result)
submit_feedback(rec_id, score=1)
print(f"[grow] store size: {len(EXAMPLE_STORE['examples'])}")
# 期待出力（だいたい）:
# classified: {"intent": "expense_export", "format": "csv", "period": "this_month"}
# [grow] store size: 4

import random
 
STATIC_EXAMPLES_TEXT = "\n\n".join(
    f"# 入力\n{e['input']}\n# 出力\n{e['output']}" for e in SEED_EXAMPLES
)
 
def classify_static(query: str):
    prompt = (
        "次の自然言語コマンドを JSON へ変換してください。\n\n"
        f"{STATIC_EXAMPLES_TEXT}\n\n# 入力\n{query}\n# 出力\n"
    )
    resp = client.models.generate_content(
        model="gemini-2.5-flash",
        contents=prompt,
        config=types.GenerateContentConfig(
            temperature=0.1,
            response_mime_type="application/json",
        ),
    )
    return json.loads(resp.text)
 
def classify_ab(query: str):
    arm = "dynamic" if random.random() < 0.5 else "static"
    if arm == "dynamic":
        rec_id, result = classify_with_logging(query)
    else:
        result = classify_static(query)
        rec_id = None
    return arm, rec_id, result