【TypeSafe】Jev で書類項目の突合はどこまでできるか？精度・速度・コストを測ってみた

#!/usr/bin/env python3
"""JEV(TypeSafe AI System One)で横突合の精度・レイテンシ・トークンを測る。

各ケースについて値A・値Bを state として渡し、noul(yes/no確率)で
「完全一致か」「意味として同じ値か」を同時に聞く。1リクエストで両方answerが返る。

API: POST https://api.typesafe.ai/v1/systemone
     https://docs.typesafe.ai/api.md
"""

from __future__ import annotations

import argparse
import json
import os
import statistics
import sys
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor

ENDPOINT = "https://api.typesafe.ai/v1/systemone"
DEFAULT_MODEL = "jev-latest"

# 入力トークンのみ課金。出力トークンは無料
USD_PER_INPUT_MTOK = 0.042

# state のキー名。質問文と同じ言語に揃える。
STATE_KEYS = {"ja": ("値A", "値B"), "en": ("value_a", "value_b")}

# 質問文の言語だけを変えた2セット。中身の意味は揃えてある。
QUESTIONS = {
    "ja": {
        "exact": {
            "type": "noul",
            "instructions": (
                "値Aと値Bは、文字列として完全に同一か。"
                "空白・記号・大文字小文字・全角半角の違いも「同一ではない」と扱う。"
            ),
            "criteria": {
                "true": "1文字も違わず完全に同じ",
                "false": "1文字でも違う",
            },
        },
        "semantic": {
            "type": "noul",
            "instructions": (
                "値Aと値Bは、表記が違っていても同じ値を指しているか。"
                "これは2つの書類の同じ項目を突合している場面での判断である。"
                "片方が明細の羅列で、もう片方が合計値の場合は、"
                "明細をすべて足し合わせてから合計値と比べること。"
            ),
            "criteria": {
                "true": ("書式・言語・単位・略記の違いだけで、指している値は同じ。"
                         "または明細の合計が合計値と一致する"),
                "false": ("指している値そのものが違う、または片方にしか値がない。"
                          "または明細の合計が合計値と一致しない"),
            },
        },
    },
    "en": {
        "exact": {
            "type": "noul",
            "instructions": (
                "Are value_a and value_b identical as strings? "
                "Any difference in whitespace, punctuation, letter case, or "
                "full-width versus half-width characters counts as not identical."
            ),
            "criteria": {
                "true": "Identical character for character.",
                "false": "They differ by at least one character.",
            },
        },
        "semantic": {
            "type": "noul",
            "instructions": (
                "Do value_a and value_b refer to the same underlying value, "
                "even when they are written differently? This is a cross-check "
                "of the same field on two business documents. When one side "
                "lists itemized values and the other states a total, add the "
                "items up and compare that sum against the stated total."
            ),
            "criteria": {
                "true": (
                    "Only the format, language, unit, or abbreviation differs "
                    "and the underlying value is the same, or the itemized "
                    "values add up to the stated total."
                ),
                "false": (
                    "The underlying value itself differs, one side has no "
                    "value, or the itemized values do not add up to the "
                    "stated total."
                ),
            },
        },
    },
}

def build_questions(lang: str) -> dict:
    """1リクエストで聞く質問セットを組み立てる。"""
    src = QUESTIONS[lang]
    return {"exact": src["exact"], "semantic": src["semantic"]}

def call_jev(case: dict, questions: dict, api_key: str, model: str,
             timeout: float, lang: str) -> dict:
    """1ケースをJEVに投げ、応答とレイテンシを返す。"""
    key_a, key_b = STATE_KEYS[lang]
    body = json.dumps({
        "state": {key_a: case["a"], key_b: case["b"]},
        "model": model,
        "questions": questions,
    }, ensure_ascii=False).encode("utf-8")

    req = urllib.request.Request(
        ENDPOINT,
        data=body,
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        },
        method="POST",
    )

    started = time.perf_counter()
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            payload = json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")[:500]
        return {"error": f"HTTP {exc.code}: {detail}",
                "latency_ms": (time.perf_counter() - started) * 1000}
    except Exception as exc:  # ネットワーク断・タイムアウト
        return {"error": f"{type(exc).__name__}: {exc}",
                "latency_ms": (time.perf_counter() - started) * 1000}

    return {"payload": payload,
            "latency_ms": (time.perf_counter() - started) * 1000}

def judge(case: dict, result: dict, thresholds: dict[str, float]) -> dict:
    """応答を期待値と突き合わせて1行分のレコードにする。"""
    row = {
        "id": case["id"],
        "category": case["category"],
        "note": case["note"],
        "latency_ms": round(result["latency_ms"], 1),
    }
    if "error" in result:
        row["error"] = result["error"]
        return row

    payload = result["payload"]
    answers = payload.get("answers", {})
    usage = payload.get("usage") or {}

    row["model"] = payload.get("model")
    row["input_tokens"] = usage.get("input_tokens")
    row["output_tokens"] = usage.get("output_tokens")

    for key in ("exact", "semantic"):
        prob = answers.get(key, {}).get("noul")
        expected = case[f"expect_{key}"]
        got = None if prob is None else prob >= thresholds[key]
        row[f"{key}_prob"] = prob
        row[f"{key}_threshold"] = thresholds[key]
        row[f"{key}_got"] = got
        row[f"{key}_expected"] = expected
        # 期待値 None は正解が決まらないケースなので正答率から外す
        row[f"{key}_ok"] = None if expected is None else (got == expected)

    return row

def summarize(rows: list[dict]) -> None:
    errors = [r for r in rows if "error" in r]
    ok_rows = [r for r in rows if "error" not in r]

    print("\n=== 結果サマリ ===")
    print(f"ケース数: {len(rows)}  成功: {len(ok_rows)}  失敗: {len(errors)}")

    for key in ("exact", "semantic"):
        scored = [r for r in ok_rows if r.get(f"{key}_ok") is not None]
        hit = sum(1 for r in scored if r[f"{key}_ok"])
        rate = f"{hit / len(scored):.1%}" if scored else "-"
        th = ok_rows[0].get(f"{key}_threshold") if ok_rows else None
        print(f"{key:9s} 正答 {hit}/{len(scored)} ({rate})  [閾値 {th}]")

    print("\n--- カテゴリ別(semantic) ---")
    for cat in sorted({r["category"] for r in ok_rows}):
        scored = [r for r in ok_rows
                  if r["category"] == cat and r.get("semantic_ok") is not None]
        if not scored:
            continue
        hit = sum(1 for r in scored if r["semantic_ok"])
        print(f"{cat:12s} {hit}/{len(scored)} ({hit / len(scored):.0%})")

    for key in ("exact", "semantic"):
        misses = [r for r in ok_rows if r.get(f"{key}_ok") is False]
        if not misses:
            continue
        print(f"\n--- {key} を外したケース ---")
        for r in misses:
            print(f"{r['id']:5s} {r['note']}  期待={r[f'{key}_expected']} "
                  f"(p={r[f'{key}_prob']:.3f})")

    if ok_rows:
        lat = sorted(r["latency_ms"] for r in ok_rows)
        p = lambda q: lat[min(int(len(lat) * q), len(lat) - 1)]  # noqa: E731
        print("\n--- レイテンシ(ms) ---")
        print(f"平均 {statistics.mean(lat):.0f} / 中央値 {p(0.5):.0f} / "
              f"p95 {p(0.95):.0f} / 最小 {lat[0]:.0f} / 最大 {lat[-1]:.0f}")

        tin = sum(r.get("input_tokens") or 0 for r in ok_rows)
        tout = sum(r.get("output_tokens") or 0 for r in ok_rows)
        print("\n--- トークン ---")
        print(f"入力 合計 {tin} (平均 {tin / len(ok_rows):.0f}/件) / 出力 合計 {tout}")
        print(f"概算コスト: ${tin / 1_000_000 * USD_PER_INPUT_MTOK:.6f} "
              f"(入力 ${USD_PER_INPUT_MTOK}/1Mtok、出力は課金対象外)")

    for r in errors:
        print(f"[ERROR] {r['id']}: {r['error']}", file=sys.stderr)

def main() -> int:
    ap = argparse.ArgumentParser(description="JEVで横突合の精度・速度・トークンを測る")
    ap.add_argument("--cases", default="cases.jsonl")
    ap.add_argument("--out", help="デフォルトは results_<lang>.jsonl")
    ap.add_argument("--lang", choices=("ja", "en"), default="en",
                    help="質問文とstateキーの言語")
    ap.add_argument("--model", default=DEFAULT_MODEL)
    ap.add_argument("--threshold-exact", type=float, default=0.8,
                    help="exactのnoul確率を一致とみなす閾値")
    ap.add_argument("--threshold-semantic", type=float, default=0.8,
                    help="semanticのnoul確率を一致とみなす閾値")
    ap.add_argument("--concurrency", type=int, default=1,
                    help="並列数。レイテンシを正しく測るなら1のまま")
    ap.add_argument("--timeout", type=float, default=30.0)
    args = ap.parse_args()
    out_path = args.out or f"results_{args.lang}.jsonl"

    with open(args.cases, encoding="utf-8") as f:
        cases = [json.loads(line) for line in f if line.strip()]

    questions = build_questions(args.lang)

    api_key = os.environ.get("TYPESAFE_API_KEY")
    if not api_key:
        print("TYPESAFE_API_KEY が未設定です。export してください。", file=sys.stderr)
        return 1

    print(f"{len(cases)} ケースを {args.model} に投げます "
          f"(質問文 {args.lang} / 並列 {args.concurrency})...", file=sys.stderr)

    thresholds = {"exact": args.threshold_exact,
                  "semantic": args.threshold_semantic}

    run = lambda c: judge(  # noqa: E731
        c, call_jev(c, questions, api_key, args.model, args.timeout, args.lang),
        thresholds)

    wall = time.perf_counter()
    if args.concurrency > 1:
        with ThreadPoolExecutor(max_workers=args.concurrency) as pool:
            rows = list(pool.map(run, cases))
    else:
        rows = [run(c) for c in cases]
    wall = time.perf_counter() - wall

    with open(out_path, "w", encoding="utf-8") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")

    summarize(rows)
    print(f"\n全体所要: {wall:.1f}s ({wall / len(rows):.2f}s/件)")
    print(f"明細: {out_path}")
    return 0

if __name__ == "__main__":
    raise SystemExit(main())