ADK アシスタントが三日前の約束を静かに忘れていた — 会話圧縮の記憶欠落を再現テストで炙り出す

"""recall_probe.py — 圧縮後セッションの記憶保持率を測る"""
import asyncio
import re
from dataclasses import dataclass, field
 
from google import genai
from google.genai import types
 
client = genai.Client()  # 環境変数 GOOGLE_API_KEY を参照
 
@dataclass
class Fact:
    """会話中に登場し、後から必ず想起されてほしい事実"""
    key: str                       # 台帳上の識別子
    question: str                  # 想起を促す質問
    patterns: list[str]            # 許容表記（正規表現）
    critical: bool = False         # 落ちたら即失格とみなすか
 
@dataclass
class ProbeResult:
    strategy: str
    recalled: set[str] = field(default_factory=set)
    lost: set[str] = field(default_factory=set)
    lost_critical: set[str] = field(default_factory=set)
    prompt_tokens: int = 0
 
    @property
    def recall_rate(self) -> float:
        total = len(self.recalled) + len(self.lost)
        return len(self.recalled) / total if total else 0.0
 
 
def _matched_by_pattern(answer: str, fact: Fact) -> bool:
    """第一段: 表記ゆれを正規表現で吸収する。ここで通れば LLM 判定を呼ばない"""
    return any(re.search(p, answer) for p in fact.patterns)
 
 
async def _matched_by_judge(answer: str, fact: Fact) -> bool:
    """第二段: 正規表現で拾えなかった言い換えだけを Gemini に判定させる"""
    verdict = await client.aio.models.generate_content(
        model="gemini-flash-latest",
        contents=(
            f"質問: {fact.question}\n"
            f"回答: {answer}\n\n"
            f"この回答は質問に対する具体的な事実を含んでいますか。"
            f"「わからない」「情報がない」等は含んでいないと判定します。"
            f"YES または NO の一語のみで答えてください。"
        ),
        config=types.GenerateContentConfig(
            temperature=0.0,       # 判定器のブレを測定に混ぜない
            max_output_tokens=4,
        ),
    )
    return (verdict.text or "").strip().upper().startswith("YES")
 
 
async def probe(ask, facts: list[Fact], strategy: str) -> ProbeResult:
    """ask(question) -> (answer, prompt_tokens) を受け取り、事実ごとに想起可否を記録する"""
    result = ProbeResult(strategy=strategy)
    for fact in facts:
        answer, tokens = await ask(fact.question)
        result.prompt_tokens += tokens
 
        ok = _matched_by_pattern(answer, fact)
        if not ok:
            ok = await _matched_by_judge(answer, fact)
 
        if ok:
            result.recalled.add(fact.key)
        else:
            result.lost.add(fact.key)
            if fact.critical:
                result.lost_critical.add(fact.key)
    return result

FACTS = [
    Fact(
        key="deadline",
        question="私が話した締切はいつでしたか",
        patterns=[r"金曜", r"18:00", r"18時", r"午後6時", r"午後六時"],
        critical=True,
    ),
    Fact(
        key="decision_screenshot_tool",
        question="スクリーンショットの生成にはどの方法で進めることにしましたか",
        patterns=[r"スクリプト", r"自動化", r"CLI"],
        critical=True,
    ),
    Fact(
        key="user_level",
        question="私の Python の習熟度をどう理解していますか",
        patterns=[r"初学者", r"学び始め", r"入門", r"beginner"],
    ),
    Fact(
        key="budget",
        question="月あたりの API 予算はいくらと伝えましたか",
        patterns=[r"3,?000\s*円", r"三千円"],
        critical=True,
    ),
    Fact(
        key="project_name",
        question="私が進めているアプリの名前は何でしたか",
        patterns=[r"Wallpaper\s*Studio", r"ウォールペーパー"],
    ),
]

"""fact_extractor.py — 会話ターンから構造化事実を抜き、圧縮の外側に置く"""
from google.genai import types
 
FACT_SCHEMA = {
    "type": "object",
    "properties": {
        "facts": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "kind": {
                        "type": "string",
                        "enum": ["deadline", "decision", "constraint",
                                 "entity", "user_attribute"],
                    },
                    "value": {"type": "string"},
                    "source_turn": {"type": "integer"},
                },
                "required": ["kind", "value", "source_turn"],
            },
        }
    },
    "required": ["facts"],
}
 
 
async def extract_facts(client, turn_text: str, turn_index: int) -> list[dict]:
    """1ターン分のテキストから事実を抽出する。抽出漏れより過剰抽出を許容する設計"""
    resp = await client.aio.models.generate_content(
        model="gemini-flash-latest",     # 抽出は軽量モデルで十分。Pro を使う場所ではない
        contents=(
            f"以下の発言から、後の会話で参照される可能性のある事実だけを抽出してください。\n"
            f"日時・決定事項・数値制約・固有名詞・ユーザー属性のいずれかに該当するものに限ります。\n"
            f"該当がなければ空配列を返してください。\n\n"
            f"発言（ターン {turn_index}）:\n{turn_text}"
        ),
        config=types.GenerateContentConfig(
            response_mime_type="application/json",
            response_schema=FACT_SCHEMA,   # 自由文パースを挟まないことが安定性の要
            temperature=0.0,
        ),
    )
    return resp.parsed["facts"]
 
 
def merge_into_state(state: dict, new_facts: list[dict]) -> dict:
    """同一 kind + value は上書きせず、より新しい source_turn を採用する"""
    ledger: dict[tuple[str, str], dict] = {
        (f["kind"], f["value"]): f for f in state.get("fact_ledger", [])
    }
    for f in new_facts:
        k = (f["kind"], f["value"])
        if k not in ledger or f["source_turn"] > ledger[k]["source_turn"]:
            ledger[k] = f
    state["fact_ledger"] = list(ledger.values())
    return state
 
 
def render_ledger(state: dict) -> str:
    """毎ターンの system 側に原文のまま差し込む。要約を通さないことが唯一の要件"""
    facts = state.get("fact_ledger", [])
    if not facts:
        return ""
    lines = [f"- [{f['kind']}] {f['value']}" for f in facts]
    return "確定済みの事実（会話の要約より優先して参照すること）:\n" + "\n".join(lines)

def gate(result: ProbeResult, baseline_tokens: int) -> tuple[bool, list[str]]:
    problems: list[str] = []
    if result.lost_critical:
        problems.append(f"critical欠落: {sorted(result.lost_critical)}")
    if result.recall_rate < 0.95:
        problems.append(f"リコール率 {result.recall_rate:.0%} < 95%")
 
    growth = (result.prompt_tokens - baseline_tokens) / baseline_tokens
    if growth > 0.20:
        # 警告どまり。コスト増は品質劣化より緊急度が低い
        print(f"⚠️ プロンプトトークン +{growth:.0%}")
 
    return (not problems), problems