diff --git a/CLAUDE.md b/CLAUDE.md index 34f87c0..cdf9215 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -90,10 +90,22 @@ query → BM25 (FTS5 verbatim) + HNSW (distilled embedding) ```toml [distill] -model = "claude-haiku-4-5-20251001" # 蒸留に使うモデル +provider = "claude" # 蒸留 backend: "claude" | "openai"（既定 "claude"） +model = "claude-haiku-4-5-20251001" # 蒸留に使うモデル（provider 共通） batch_limit = 20 # hook 1回あたりの蒸留上限 ``` +ローカル LLM（OpenAI 互換: Ollama / LM Studio / llama.cpp-server / vLLM）で蒸留する場合は `provider = "openai"` と `base_url` を指定する（新規依存なし・API キー不要＝`Authorization` ヘッダーは送らないローカル専用）: + +```toml +[distill] +provider = "openai" +model = "qwen2.5:7b" +base_url = "http://localhost:11434/v1" # Ollama（LM Studio は :1234/v1） +``` + +`provider = "openai"` のとき `base_url` は必須。未設定/空なら警告して `claude` にフォールバックする。`provider = "claude"`（既定）では `base_url` は無視され、従来どおり `claude --print` で蒸留する。実装は `llm.py` の `DistillBackend` + `call_claude` dispatcher（`_call_openai` は stdlib `urllib`、`response_format=json_object` + プロンプト内スキーマ + `_validate_palace` 検証 + 1 回リトライ）。 + ## CLI Commands ```bash diff --git a/README.md b/README.md index 59c459b..898083c 100644 --- a/README.md +++ b/README.md @@ -145,6 +145,7 @@ After `loci init` (or `loci hook install`), everything runs automatically: ```toml [distill] +provider = "claude" # Distillation backend: "claude" | "openai" (default "claude") model = "claude-haiku-4-5" # Model for distillation (default) batch_limit = 20 # Max distillations per hook run min_chars = 100 # Skip distillation for exchanges shorter than this @@ -155,6 +156,20 @@ min_chars = 50 # Skip indexing exchanges shorter than th There are two `min_chars` settings: `[index] min_chars` controls what gets indexed at all, while `[distill] min_chars` further skips distillation (the LLM cost) for short exchanges that were already indexed. +### Distilling with a local LLM + +Distillation is a small per-exchange structured-extraction task, so a local model is usually good enough. Any OpenAI-compatible endpoint (Ollama, LM Studio, llama.cpp-server, vLLM) works by setting `provider = "openai"` and `base_url` — no new dependencies, no API key (the `Authorization` header is never sent, so this is local-only): + +```toml +[distill] +provider = "openai" +model = "qwen2.5:7b" +base_url = "http://localhost:11434/v1" # Ollama +# base_url = "http://localhost:1234/v1" # LM Studio +``` + +`base_url` is required when `provider = "openai"`; if it is missing or empty the provider falls back to `claude` with a warning. With `provider = "claude"` (the default), `base_url` is ignored and distillation runs through `claude --print` as before. + ## Acknowledgments The palace object model, room-based topic grouping, and BM25+HNSW fusion search are based on: diff --git a/src/codeatrium/cli/__init__.py b/src/codeatrium/cli/__init__.py index bece7ac..93b8b56 100644 --- a/src/codeatrium/cli/__init__.py +++ b/src/codeatrium/cli/__init__.py @@ -130,9 +130,16 @@ def init( "# Codeatrium configuration\n" "\n" "[distill]\n" + '# provider = "claude" # 蒸留 LLM: "claude" | "openai"（既定 "claude"）\n' '# model = "claude-haiku-4-5-20251001"\n' "# batch_limit = 20\n" "# min_chars = 100 # この文字数未満の exchange は蒸留スキップ\n" + "#\n" + "# ローカル LLM（OpenAI 互換）で蒸留する場合:\n" + '# provider = "openai"\n' + '# model = "qwen2.5:7b"\n' + '# base_url = "http://localhost:11434/v1" # Ollama\n' + '# # base_url = "http://localhost:1234/v1" # LM Studio\n' "\n" "[index]\n" "# min_chars = 50 # trivial フィルタ閾値（文字数）\n" @@ -225,6 +232,7 @@ def init( try: from codeatrium.config import load_config from codeatrium.distiller import distill_all + from codeatrium.llm import DistillBackend cfg = load_config(root) typer.echo("Running distillation...") @@ -239,7 +247,7 @@ def _on_progress( count, err_count = distill_all( db, - model=cfg.distill_model, + backend=DistillBackend.from_config(cfg), on_progress=_on_progress, project_root=str(root), distill_min_chars=cfg.distill_min_chars, diff --git a/src/codeatrium/cli/distill_cmd.py b/src/codeatrium/cli/distill_cmd.py index 72ce819..be1da05 100644 --- a/src/codeatrium/cli/distill_cmd.py +++ b/src/codeatrium/cli/distill_cmd.py @@ -1,4 +1,4 @@ -"""loci distill コマンド""" +"""loci distill コマンド — backend/provider サポート""" from __future__ import annotations @@ -19,6 +19,7 @@ def distill( from codeatrium.config import load_config from codeatrium.distiller import distill_all + from codeatrium.llm import DistillBackend from codeatrium.paths import db_path, find_project_root root = find_project_root() @@ -59,7 +60,7 @@ def _on_progress(cur: int, tot: int, error: str | None = None) -> None: count, err_count = distill_all( db, limit=limit, - model=cfg.distill_model, + backend=DistillBackend.from_config(cfg), on_progress=_on_progress, project_root=str(root), distill_min_chars=cfg.distill_min_chars, diff --git a/src/codeatrium/config.py b/src/codeatrium/config.py index ae7f09a..0999419 100644 --- a/src/codeatrium/config.py +++ b/src/codeatrium/config.py @@ -1,4 +1,6 @@ -"""設定ファイルの読み込み — .codeatrium/config.toml""" +"""設定ファイルの読み込み — .codeatrium/config.toml + +Supports distill_provider (claude/openai) and distill_base_url for provider switching.""" from __future__ import annotations @@ -16,6 +18,8 @@ DEFAULT_DISTILL_BATCH_LIMIT = 20 DEFAULT_INDEX_MIN_CHARS = 50 DEFAULT_DISTILL_MIN_CHARS = 100 +DEFAULT_DISTILL_PROVIDER = "claude" +VALID_DISTILL_PROVIDERS = frozenset({"claude", "openai"}) @dataclass @@ -26,6 +30,8 @@ class Config: distill_batch_limit: int = DEFAULT_DISTILL_BATCH_LIMIT index_min_chars: int = DEFAULT_INDEX_MIN_CHARS distill_min_chars: int = DEFAULT_DISTILL_MIN_CHARS + distill_provider: str = "claude" + distill_base_url: str | None = None def load_config(project_root: Path) -> Config: @@ -79,9 +85,31 @@ def load_config(project_root: Path) -> Config: ) distill_min_chars = DEFAULT_DISTILL_MIN_CHARS + provider = distill.get("provider", DEFAULT_DISTILL_PROVIDER) + if not isinstance(provider, str) or provider not in VALID_DISTILL_PROVIDERS: + print( + "Warning: distill.provider must be one of {claude, openai}, using default.", + file=sys.stderr, + ) + provider = DEFAULT_DISTILL_PROVIDER + + base_url = distill.get("base_url") + if base_url is not None and not isinstance(base_url, str): + base_url = None + + if provider == "openai" and (base_url is None or not base_url.strip()): + print( + "Warning: distill.provider is 'openai' but base_url is not set, falling back to 'claude'.", + file=sys.stderr, + ) + provider = DEFAULT_DISTILL_PROVIDER + base_url = None + return Config( distill_model=model, distill_batch_limit=batch_limit, index_min_chars=min_chars, distill_min_chars=distill_min_chars, + distill_provider=provider, + distill_base_url=base_url, ) diff --git a/src/codeatrium/distiller.py b/src/codeatrium/distiller.py index 928fc79..edd4c51 100644 --- a/src/codeatrium/distiller.py +++ b/src/codeatrium/distiller.py @@ -5,6 +5,8 @@ ② claude -p で palace object 生成（--output-format json --json-schema） ③ distill_text を embedding して vec_palace に登録 ④ files_touched を tree-sitter で解析してシンボルを symbols テーブルに登録 + +backend パラメータで蒸留に使う DistillBackend を指定可能。 """ from __future__ import annotations @@ -19,7 +21,7 @@ os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") from codeatrium.embedder import Embedder, EmbedderSetupError -from codeatrium.llm import DISTILL_PROMPT_TEMPLATE, call_claude +from codeatrium.llm import DISTILL_PROMPT_TEMPLATE, DistillBackend, call_claude from codeatrium.models import PalaceObject from codeatrium.utils import sha256 @@ -92,7 +94,7 @@ def distill_exchange( agent_content: str, ply_start: int, ply_end: int, - model: str | None = None, + backend: DistillBackend | None = None, project_root: str | None = None, ) -> PalaceObject: """1つの exchange を蒸留して PalaceObject を返す @@ -104,7 +106,7 @@ def distill_exchange( agent_content: エージェントコンテンツ ply_start: 開始ply ply_end: 終了ply - model: 蒸留に使うモデル（デフォルトはconfig.toml から） + backend: 蒸留に使う DistillBackend（省略時はデフォルト claude） project_root: プロジェクトルート（ファイルパスフィルタ用） """ from codeatrium.db import get_connection @@ -115,7 +117,7 @@ def distill_exchange( ply_end=ply_end, messages_text=messages_text, ) - raw = call_claude(prompt, model=model) + raw = call_claude(prompt, backend=backend) # PRIMARY: exchange_files から読み込み con = get_connection(db_path) @@ -302,13 +304,14 @@ def save_palace_object( def distill_all( db_path: Path, limit: int | None = None, - model: str | None = None, + backend: DistillBackend | None = None, on_progress: Callable[..., None] | None = None, project_root: str | None = None, distill_min_chars: int = 100, ) -> tuple[int, int]: """未蒸留の exchange を処理する。 + backend: 蒸留に使う DistillBackend（省略時はデフォルト claude） distill_min_chars: この文字数未満の exchange は蒸留スキップ（デフォルト100） on_progress: (current, total, error=None) を受け取るコールバック Returns: (処理した exchange 数, エラー数) @@ -360,7 +363,7 @@ def distill_all( row["agent_content"], row["ply_start"], row["ply_end"], - model=model, + backend=backend, project_root=project_root, ) distill_text = palace.exchange_core + "\n" + palace.specific_context diff --git a/src/codeatrium/llm.py b/src/codeatrium/llm.py index 468252a..81c9d1d 100644 --- a/src/codeatrium/llm.py +++ b/src/codeatrium/llm.py @@ -1,10 +1,16 @@ -"""LLM 呼び出しラッパー: claude --print でプロンプトを実行し JSON を返す""" +"""LLM 呼び出しラッパー: claude --print または OpenAI API でプロンプトを実行し JSON を返す + +DistillBackend 抽象化により claude / openai プロバイダ両対応。 +会話エッセンス (exchange_core, specific_context, room_assignments) を蒸留する""" from __future__ import annotations import hashlib import json import subprocess +import urllib.error +import urllib.request +from dataclasses import dataclass from pathlib import Path from typing import Any @@ -33,7 +39,9 @@ JSONのみで回答してください。""" # プロンプト sha256 先頭8桁（B5: drift 検出用） -DISTILL_PROMPT_VERSION = hashlib.sha256(DISTILL_PROMPT_TEMPLATE.encode()).hexdigest()[:8] +DISTILL_PROMPT_VERSION = hashlib.sha256(DISTILL_PROMPT_TEMPLATE.encode()).hexdigest()[ + :8 +] JSON_SCHEMA = json.dumps( { @@ -68,6 +76,36 @@ ) +# ---- 例外 ---- + + +class LLMValidationError(Exception): + """LLM response validation failed""" + + pass + + +# ---- DistillBackend 抽象化 ---- + + +@dataclass(frozen=True) +class DistillBackend: + """LLM backend configuration for distillation (claude or openai)""" + + provider: str + model: str + base_url: str | None + + @classmethod + def from_config(cls, cfg) -> DistillBackend: + """Create DistillBackend from Config object""" + return DistillBackend( + provider=cfg.distill_provider, + model=cfg.distill_model, + base_url=cfg.distill_base_url, + ) + + # ---- 副作用制御 ---- @@ -94,11 +132,51 @@ def _cleanup_side_effect_jsonls(session_dir: Path, before: set[Path]) -> None: pass -# ---- LLM 呼び出し ---- +# ---- ヘルパー関数 ---- + + +def _strip_json_fence(text: str) -> str: + """Remove markdown code fence (```...```) if present, return raw string""" + if text.startswith("```"): + lines = text.splitlines() + text = "\n".join(lines[1:-1] if lines[-1].strip() == "```" else lines[1:]) + return text.strip() + + +def _build_distill_prompt(base_prompt: str) -> str: + """Append JSON schema and instruction to base prompt""" + return ( + base_prompt + + "\n" + + JSON_SCHEMA + + "\n" + + "このスキーマに厳密に従って、上記の指示通りにJSONのみで回答してください。" + ) + + +def _validate_palace(raw: dict[str, Any]) -> dict[str, Any]: + """Validate that raw dict has required palace object fields""" + if not isinstance(raw.get("exchange_core"), str): + raise LLMValidationError("exchange_core must be a string") + if not isinstance(raw.get("specific_context"), str): + raise LLMValidationError("specific_context must be a string") + if not isinstance(raw.get("room_assignments"), list): + raise LLMValidationError("room_assignments must be a list") + for i, room in enumerate(raw.get("room_assignments", [])): + if not isinstance(room, dict): + raise LLMValidationError(f"room_assignments[{i}] must be a dict") + required_keys = {"room_type", "room_key", "room_label", "relevance"} + if not all(k in room for k in required_keys): + raise LLMValidationError( + f"room_assignments[{i}] missing required keys: {required_keys}" + ) -def call_claude(prompt: str, model: str | None = None) -> dict[str, Any]: - """claude -p でプロンプトを実行し JSON を返す（テストでモック対象）""" + return raw + + +def _call_claude_cli(prompt: str, model: str | None = None) -> dict[str, Any]: + """Call claude --print CLI and return parsed JSON (unvalidated)""" import shutil from codeatrium.config import DEFAULT_DISTILL_MODEL @@ -142,11 +220,70 @@ def call_claude(prompt: str, model: str | None = None) -> dict[str, Any]: return outer["structured_output"] inner = outer.get("result", "") if isinstance(inner, str) and inner.strip(): - text = inner.strip() - if text.startswith("```"): - lines = text.splitlines() - text = "\n".join( - lines[1:-1] if lines[-1].strip() == "```" else lines[1:] - ) - return json.loads(text.strip()) + text = _strip_json_fence(inner) + return json.loads(text) return outer + + +def _call_openai(prompt: str, backend: DistillBackend) -> dict[str, Any]: + """Call OpenAI API with retry on validation failure, return parsed JSON""" + if backend.base_url is None: + raise RuntimeError("base_url is required for openai provider") + enhanced_prompt = _build_distill_prompt(prompt) + + body = { + "model": backend.model, + "messages": [{"role": "user", "content": enhanced_prompt}], + "response_format": {"type": "json_object"}, + "temperature": 0, + } + body_bytes = json.dumps(body).encode() + + for attempt in range(2): + try: + req = urllib.request.Request( + url=backend.base_url + "/chat/completions", + data=body_bytes, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req, timeout=300) as response: + response_text = response.read().decode() + except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError) as e: + raise RuntimeError(f"OpenAI API request failed: {e}") from e + + response_data = json.loads(response_text) + message_content = response_data["choices"][0]["message"]["content"] + stripped = _strip_json_fence(message_content) + result = json.loads(stripped) + + try: + return _validate_palace(result) + except LLMValidationError: + if attempt == 1: + raise + raise RuntimeError("retry exhausted") + + +# ---- LLM 呼び出し ---- + + +def call_claude( + prompt: str, model: str | None = None, *, backend: DistillBackend | None = None +) -> dict[str, Any]: + """Dispatch to configured backend (claude CLI or OpenAI API) and return validated palace object""" + if backend is None: + from codeatrium.config import DEFAULT_DISTILL_MODEL + + backend = DistillBackend( + provider="claude", + model=model or DEFAULT_DISTILL_MODEL, + base_url=None, + ) + + if backend.provider == "openai": + raw = _call_openai(prompt, backend) + else: + raw = _call_claude_cli(prompt, backend.model) + + return _validate_palace(raw) diff --git a/tests/test_config.py b/tests/test_config.py index 9550bfe..e23b69e 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -9,6 +9,7 @@ DEFAULT_DISTILL_BATCH_LIMIT, DEFAULT_DISTILL_MIN_CHARS, DEFAULT_DISTILL_MODEL, + DEFAULT_DISTILL_PROVIDER, DEFAULT_INDEX_MIN_CHARS, Config, load_config, @@ -40,7 +41,7 @@ def test_load_config_partial(tmp_path: Path) -> None: """一部だけ設定した場合、残りはデフォルト""" codeatrium_dir = tmp_path / ".codeatrium" codeatrium_dir.mkdir() - (codeatrium_dir / "config.toml").write_text('[distill]\nbatch_limit = 5\n') + (codeatrium_dir / "config.toml").write_text("[distill]\nbatch_limit = 5\n") cfg = load_config(tmp_path) assert cfg.distill_model == DEFAULT_DISTILL_MODEL assert cfg.distill_batch_limit == 5 @@ -50,7 +51,7 @@ def test_load_config_invalid_batch_limit_fallback(tmp_path: Path) -> None: """不正な batch_limit はデフォルトにフォールバック""" codeatrium_dir = tmp_path / ".codeatrium" codeatrium_dir.mkdir() - (codeatrium_dir / "config.toml").write_text('[distill]\nbatch_limit = -1\n') + (codeatrium_dir / "config.toml").write_text("[distill]\nbatch_limit = -1\n") cfg = load_config(tmp_path) assert cfg.distill_batch_limit == DEFAULT_DISTILL_BATCH_LIMIT @@ -139,3 +140,65 @@ def test_load_config_oserror_fallback(tmp_path: Path) -> None: with patch("pathlib.Path.open", side_effect=OSError("disk error")): cfg = load_config(tmp_path) assert cfg == Config() + + +def test_load_config_provider_default(tmp_path: Path) -> None: + """[distill] provider キーがないときはデフォルト""" + codeatrium_dir = tmp_path / ".codeatrium" + codeatrium_dir.mkdir() + (codeatrium_dir / "config.toml").write_text("[distill]\n") + cfg = load_config(tmp_path) + assert cfg.distill_provider == DEFAULT_DISTILL_PROVIDER + assert cfg.distill_base_url is None + + +def test_load_config_provider_invalid_fallback(tmp_path: Path) -> None: + """不正な provider はデフォルトにフォールバック""" + codeatrium_dir = tmp_path / ".codeatrium" + codeatrium_dir.mkdir() + (codeatrium_dir / "config.toml").write_text("[distill]\nprovider = 'badprovider'\n") + cfg = load_config(tmp_path) + assert cfg.distill_provider == DEFAULT_DISTILL_PROVIDER + + +def test_load_config_provider_openai_with_base_url(tmp_path: Path) -> None: + """provider = 'openai' で base_url が設定されている場合""" + codeatrium_dir = tmp_path / ".codeatrium" + codeatrium_dir.mkdir() + (codeatrium_dir / "config.toml").write_text( + "[distill]\nprovider = 'openai'\nbase_url = 'http://localhost:11434/v1'\n" + ) + cfg = load_config(tmp_path) + assert cfg.distill_provider == "openai" + assert cfg.distill_base_url == "http://localhost:11434/v1" + + +def test_load_config_provider_openai_missing_base_url_fallback(tmp_path: Path) -> None: + """provider = 'openai' だが base_url がない場合はデフォルトにフォールバック""" + codeatrium_dir = tmp_path / ".codeatrium" + codeatrium_dir.mkdir() + (codeatrium_dir / "config.toml").write_text("[distill]\nprovider = 'openai'\n") + cfg = load_config(tmp_path) + assert cfg.distill_provider == DEFAULT_DISTILL_PROVIDER + assert cfg.distill_base_url is None + + +def test_load_config_provider_openai_empty_base_url_fallback(tmp_path: Path) -> None: + """provider = 'openai' で base_url が空文字列の場合はデフォルトにフォールバック""" + codeatrium_dir = tmp_path / ".codeatrium" + codeatrium_dir.mkdir() + (codeatrium_dir / "config.toml").write_text( + "[distill]\nprovider = 'openai'\nbase_url = ''\n" + ) + cfg = load_config(tmp_path) + assert cfg.distill_provider == DEFAULT_DISTILL_PROVIDER + + +def test_load_config_provider_claude_no_base_url(tmp_path: Path) -> None: + """provider = 'claude' の場合は base_url は不要""" + codeatrium_dir = tmp_path / ".codeatrium" + codeatrium_dir.mkdir() + (codeatrium_dir / "config.toml").write_text("[distill]\nprovider = 'claude'\n") + cfg = load_config(tmp_path) + assert cfg.distill_provider == "claude" + assert cfg.distill_base_url is None diff --git a/tests/test_distiller.py b/tests/test_distiller.py index a55c927..6d43b7c 100644 --- a/tests/test_distiller.py +++ b/tests/test_distiller.py @@ -18,6 +18,7 @@ extract_files_touched, save_palace_object, ) +from codeatrium.llm import DistillBackend # --- フィクスチャ --- @@ -114,9 +115,7 @@ def test_extract_files_excludes_venv() -> None: def test_extract_files_excludes_node_modules() -> None: - result = extract_files_touched( - "node_modules/react/index.js", "" - ) + result = extract_files_touched("node_modules/react/index.js", "") assert result == [] @@ -145,7 +144,8 @@ def test_extract_files_project_root_filters_absolute() -> None: def test_extract_files_project_root_unknown_external() -> None: """ハードコードマーカーにない外部パスも git root で除外される""" result = extract_files_touched( - "/home/user/.local/lib/python3.11/foo/bar.py", "", + "/home/user/.local/lib/python3.11/foo/bar.py", + "", project_root="/home/user/myproject", ) assert result == [] @@ -154,7 +154,8 @@ def test_extract_files_project_root_unknown_external() -> None: def test_extract_files_relative_paths_unaffected_by_root() -> None: """相対パスは project_root に影響されずマーカーのみでフィルタ""" result = extract_files_touched( - "src/app.py node_modules/react/index.js", "", + "src/app.py node_modules/react/index.js", + "", project_root="/home/user/myproject", ) assert result == ["src/app.py"] @@ -167,7 +168,9 @@ def test_extract_files_relative_paths_unaffected_by_root() -> None: def test_distill_exchange_returns_palace(mock_call, tmp_path) -> None: db_path = tmp_path / "memory.db" init_db(db_path) - palace = distill_exchange("ex1", db_path, "pool の設定", "pool_size=5 を追加した", 0, 3) + palace = distill_exchange( + "ex1", db_path, "pool の設定", "pool_size=5 を追加した", 0, 3 + ) assert palace.exchange_core == "pool_size を 5 に設定した" assert palace.specific_context == "pool_size=5" assert len(palace.room_assignments) == 1 @@ -185,7 +188,9 @@ def test_distill_exchange_calls_claude_once(mock_call, tmp_path) -> None: def test_distill_exchange_extracts_files(mock_call, tmp_path) -> None: db_path = tmp_path / "memory.db" init_db(db_path) - palace = distill_exchange("ex1", db_path, "src/db/pool.py を修正", "pool_size=5", 0, 3) + palace = distill_exchange( + "ex1", db_path, "src/db/pool.py を修正", "pool_size=5", 0, 3 + ) assert "src/db/pool.py" in palace.files_touched @@ -211,7 +216,9 @@ def test_distill_exchange_merges_exchange_files(mock_call, tmp_path) -> None: def test_save_palace_object_symbol_id_uses_3part_hash(tmp_path) -> None: db_path = tmp_path / "memory.db" init_db(db_path) - _make_exchange(db_path, "ex1", user_text="Foo.bar method", agent_text="Foo.bar implementation") + _make_exchange( + db_path, "ex1", user_text="Foo.bar method", agent_text="Foo.bar implementation" + ) resolver = MagicMock() sym = MagicMock() @@ -228,7 +235,9 @@ def test_save_palace_object_symbol_id_uses_3part_hash(tmp_path) -> None: room_assignments=[], files_touched=["src/foo.py"], ) - save_palace_object(db_path, "ex1", palace, np.zeros(384, dtype=np.float32), resolver=resolver) + save_palace_object( + db_path, "ex1", palace, np.zeros(384, dtype=np.float32), resolver=resolver + ) con = get_connection(db_path) row = con.execute("SELECT id, dedup_hash FROM symbols").fetchone() @@ -274,7 +283,12 @@ def test_save_palace_object_stores_in_db(tmp_path) -> None: def test_save_palace_object_skips_symbol_not_in_body(tmp_path) -> None: db_path = tmp_path / "memory.db" init_db(db_path) - _make_exchange(db_path, "ex1", user_text="some unrelated text " * 5, agent_text="more text " * 5) + _make_exchange( + db_path, + "ex1", + user_text="some unrelated text " * 5, + agent_text="more text " * 5, + ) resolver = MagicMock() sym = MagicMock() @@ -288,7 +302,9 @@ def test_save_palace_object_skips_symbol_not_in_body(tmp_path) -> None: room_assignments=[], files_touched=["src/foo.py"], ) - save_palace_object(db_path, "ex1", palace, np.zeros(384, dtype=np.float32), resolver=resolver) + save_palace_object( + db_path, "ex1", palace, np.zeros(384, dtype=np.float32), resolver=resolver + ) con = get_connection(db_path) count = con.execute("SELECT COUNT(*) FROM symbols").fetchone()[0] @@ -300,7 +316,9 @@ def test_save_palace_object_skips_symbol_not_in_body(tmp_path) -> None: def test_save_palace_object_includes_symbol_in_body(tmp_path) -> None: db_path = tmp_path / "memory.db" init_db(db_path) - _make_exchange(db_path, "ex1", user_text="Foo.bar method " * 5, agent_text="more text " * 5) + _make_exchange( + db_path, "ex1", user_text="Foo.bar method " * 5, agent_text="more text " * 5 + ) resolver = MagicMock() sym = MagicMock() @@ -317,7 +335,9 @@ def test_save_palace_object_includes_symbol_in_body(tmp_path) -> None: room_assignments=[], files_touched=["src/foo.py"], ) - save_palace_object(db_path, "ex1", palace, np.zeros(384, dtype=np.float32), resolver=resolver) + save_palace_object( + db_path, "ex1", palace, np.zeros(384, dtype=np.float32), resolver=resolver + ) con = get_connection(db_path) count = con.execute("SELECT COUNT(*) FROM symbols").fetchone()[0] @@ -375,8 +395,18 @@ def test_save_palace_object_saves_rooms(tmp_path) -> None: def test_save_palace_object_two_palace_objects_same_symbol(tmp_path) -> None: db_path = tmp_path / "memory.db" init_db(db_path) - _make_exchange(db_path, "ex1", user_text="Foo.bar method " * 5, agent_text="implementation " * 5) - _make_exchange(db_path, "ex2", user_text="Foo.bar method " * 5, agent_text="implementation " * 5) + _make_exchange( + db_path, + "ex1", + user_text="Foo.bar method " * 5, + agent_text="implementation " * 5, + ) + _make_exchange( + db_path, + "ex2", + user_text="Foo.bar method " * 5, + agent_text="implementation " * 5, + ) resolver = MagicMock() sym = MagicMock() @@ -394,8 +424,12 @@ def test_save_palace_object_two_palace_objects_same_symbol(tmp_path) -> None: files_touched=["src/foo.py"], ) - save_palace_object(db_path, "ex1", palace, np.zeros(384, dtype=np.float32), resolver=resolver) - save_palace_object(db_path, "ex2", palace, np.zeros(384, dtype=np.float32), resolver=resolver) + save_palace_object( + db_path, "ex1", palace, np.zeros(384, dtype=np.float32), resolver=resolver + ) + save_palace_object( + db_path, "ex2", palace, np.zeros(384, dtype=np.float32), resolver=resolver + ) con = get_connection(db_path) all_rows = con.execute("SELECT id, dedup_hash FROM symbols").fetchall() @@ -451,7 +485,9 @@ def test_distill_all_skips_distilled(mock_call, tmp_path) -> None: _make_exchange(db_path, "ex1") con = get_connection(db_path) - con.execute("UPDATE exchanges SET distilled_at = '2026-01-01', distill_status = 'distilled' WHERE id = 'ex1'") + con.execute( + "UPDATE exchanges SET distilled_at = '2026-01-01', distill_status = 'distilled' WHERE id = 'ex1'" + ) con.commit() con.close() @@ -613,3 +649,47 @@ def test_save_palace_object_rollback_on_error(tmp_path) -> None: palace_rows = con.execute("SELECT * FROM palace_objects").fetchall() assert len(palace_rows) == 0 con.close() + + +@patch("codeatrium.distiller.call_claude", return_value=MOCK_PALACE_RESPONSE) +def test_distill_exchange_accepts_backend(mock_call, tmp_path) -> None: + """distill_exchange は backend パラメータを受け取れる""" + db_path = tmp_path / "memory.db" + init_db(db_path) + backend = DistillBackend( + provider="claude", model="claude-haiku-4-5-20251001", base_url=None + ) + palace = distill_exchange( + "ex1", db_path, "pool の設定", "pool_size=5", 0, 3, backend=backend + ) + assert palace.exchange_core == "pool_size を 5 に設定した" + mock_call.assert_called_once() + + +@patch("codeatrium.distiller.call_claude", return_value=MOCK_PALACE_RESPONSE) +def test_distill_all_accepts_backend(mock_call, tmp_path) -> None: + """distill_all は backend パラメータを受け取れる""" + db_path = tmp_path / "memory.db" + init_db(db_path) + _make_exchange(db_path, "ex1") + + backend = DistillBackend( + provider="claude", model="claude-haiku-4-5-20251001", base_url=None + ) + mock_embedder = MagicMock() + mock_embedder.embed_passage.return_value = np.zeros(384, dtype=np.float32) + + with patch("codeatrium.distiller.Embedder", return_value=mock_embedder): + count, _ = distill_all(db_path, backend=backend) + + assert count == 1 + + +@patch("codeatrium.distiller.call_claude", return_value=MOCK_PALACE_RESPONSE) +def test_distill_exchange_patch_point_unchanged(mock_call, tmp_path) -> None: + """patch('codeatrium.distiller.call_claude') がパッチ対象として機能し続ける（回帰防止）""" + db_path = tmp_path / "memory.db" + init_db(db_path) + palace = distill_exchange("ex1", db_path, "pool の設定", "pool_size=5", 0, 3) + assert palace.exchange_core == "pool_size を 5 に設定した" + mock_call.assert_called_once() diff --git a/tests/test_llm.py b/tests/test_llm.py index 1be016f..6e48a48 100644 --- a/tests/test_llm.py +++ b/tests/test_llm.py @@ -4,12 +4,21 @@ import json import subprocess +import urllib.error from pathlib import Path +from typing import Any from unittest.mock import MagicMock, patch import pytest -from codeatrium.llm import call_claude +from codeatrium.llm import ( + DistillBackend, + LLMValidationError, + _call_openai, + _strip_json_fence, + _validate_palace, + call_claude, +) # ---- テストデータ ---- @@ -110,9 +119,7 @@ def fake_run(*args, **kwargs): with patch("codeatrium.llm.subprocess.run", side_effect=fake_run): with patch("shutil.which", return_value="/usr/bin/claude"): - with patch( - "codeatrium.llm._session_dir", return_value=tmp_path - ): + with patch("codeatrium.llm._session_dir", return_value=tmp_path): call_claude("test prompt") # 正常終了後は副作用 .jsonl がクリーンアップされたことを確認 @@ -137,12 +144,269 @@ def fake_run(*args, **kwargs): side_effect=fake_run, ): with patch("shutil.which", return_value="/usr/bin/claude"): - with patch( - "codeatrium.llm._session_dir", return_value=tmp_path - ): + with patch("codeatrium.llm._session_dir", return_value=tmp_path): # TimeoutExpired が発生することを確認 with pytest.raises(subprocess.TimeoutExpired): call_claude("test prompt") # タイムアウト時にも副作用 .jsonl がクリーンアップされたことを確認 assert not side_jsonl.exists() + + +def test_call_claude_dispatches_to_openai_backend() -> None: + """ + call_claude が backend パラメータを受け取り、_call_openai に委譲することを確認する + """ + backend = DistillBackend( + provider="openai", model="llama3", base_url="http://localhost:11434/v1" + ) + + with patch("codeatrium.llm._call_openai") as mock_call_openai: + mock_call_openai.return_value = MOCK_JSON_RESPONSE["structured_output"] + call_claude("prompt", backend=backend) + + # _call_openai が呼ばれ、prompt と backend が渡されたことを確認 + assert mock_call_openai.called + call_args = mock_call_openai.call_args + assert call_args is not None + assert call_args[0][0] == "prompt" + assert call_args[0][1] == backend + + +def test_call_claude_dispatches_to_claude_by_default() -> None: + """ + call_claude が backend パラメータなしの場合、_call_claude_cli に委譲することを確認する + """ + mock_result = MagicMock() + mock_result.returncode = 0 + mock_result.stdout = json.dumps(MOCK_JSON_RESPONSE) + + with patch("codeatrium.llm._call_claude_cli") as mock_call_claude_cli: + mock_call_claude_cli.return_value = MOCK_JSON_RESPONSE["structured_output"] + call_claude("prompt") + + # _call_claude_cli が呼ばれたことを確認 + assert mock_call_claude_cli.called + + +def test_call_openai_sends_correct_url_and_body() -> None: + """ + _call_openai が正しい URL とボディで request を送信することを確認する + """ + backend = DistillBackend( + provider="openai", model="llama3", base_url="http://localhost:11434/v1" + ) + + mock_response = MagicMock() + response_body = json.dumps( + { + "choices": [ + { + "message": { + "content": json.dumps(MOCK_JSON_RESPONSE["structured_output"]) + } + } + ] + } + ) + mock_response.__enter__ = MagicMock(return_value=mock_response) + mock_response.__exit__ = MagicMock(return_value=None) + mock_response.read.return_value = response_body.encode() + + with patch("urllib.request.urlopen", return_value=mock_response) as mock_urlopen: + result = _call_openai("prompt", backend) + + # urlopen が呼ばれたことを確認 + assert mock_urlopen.called + + # Request オブジェクトを取得 + call_args = mock_urlopen.call_args + assert call_args is not None + request_obj = call_args[0][0] + + # URL が正しいことを確認 + assert request_obj.full_url == "http://localhost:11434/v1/chat/completions" + + # Authorization ヘッダーがないことを確認 + assert "Authorization" not in request_obj.headers + assert result == MOCK_JSON_RESPONSE["structured_output"] + + +def test_call_openai_no_authorization_header() -> None: + """ + _call_openai の Request に Authorization ヘッダーがないことを確認する + """ + backend = DistillBackend( + provider="openai", model="llama3", base_url="http://localhost:11434/v1" + ) + + mock_response = MagicMock() + response_body = json.dumps( + { + "choices": [ + { + "message": { + "content": json.dumps(MOCK_JSON_RESPONSE["structured_output"]) + } + } + ] + } + ) + mock_response.__enter__ = MagicMock(return_value=mock_response) + mock_response.__exit__ = MagicMock(return_value=None) + mock_response.read.return_value = response_body.encode() + + with patch("urllib.request.urlopen", return_value=mock_response) as mock_urlopen: + _call_openai("prompt", backend) + + # Request オブジェクトを取得 + call_args = mock_urlopen.call_args + assert call_args is not None + request_obj = call_args[0][0] + + # Authorization ヘッダーがないことを確認 + assert "Authorization" not in request_obj.headers + + +def test_strip_json_fence_with_fence() -> None: + """ + _strip_json_fence が markdown code fence を削除することを確認する + """ + input_text = '```json\n{"a":1}\n```' + result = _strip_json_fence(input_text) + assert result == '{"a":1}' + + +def test_strip_json_fence_without_fence() -> None: + """ + _strip_json_fence が fence なし JSON を返すことを確認する + """ + input_text = '{"a":1}' + result = _strip_json_fence(input_text) + assert result == '{"a":1}' + + +def test_validate_palace_valid() -> None: + """ + _validate_palace が有効な palace dict を受け入れることを確認する + """ + palace = { + "exchange_core": "c", + "specific_context": "s", + "room_assignments": [ + { + "room_type": "concept", + "room_key": "k", + "room_label": "l", + "relevance": 0.5, + } + ], + } + result = _validate_palace(palace) + assert result == palace + + +def test_validate_palace_missing_key_raises() -> None: + """ + _validate_palace が必須キーなしの dict で LLMValidationError を発生させることを確認する + """ + palace = {"exchange_core": "c"} + with pytest.raises(LLMValidationError): + _validate_palace(palace) + + +def test_call_openai_retry_on_validation_failure() -> None: + """ + _call_openai が validation 失敗時にリトライし、2回目に成功することを確認する + """ + backend = DistillBackend( + provider="openai", model="llama3", base_url="http://localhost:11434/v1" + ) + + mock_response = MagicMock() + response_body = json.dumps( + { + "choices": [ + { + "message": { + "content": json.dumps(MOCK_JSON_RESPONSE["structured_output"]) + } + } + ] + } + ) + mock_response.__enter__ = MagicMock(return_value=mock_response) + mock_response.__exit__ = MagicMock(return_value=None) + mock_response.read.return_value = response_body.encode() + + call_count = {"count": 0} + + def validate_side_effect(palace: dict[str, Any]) -> dict[str, Any]: + call_count["count"] += 1 + if call_count["count"] == 1: + raise LLMValidationError("validation failed") + return palace + + with patch("urllib.request.urlopen", return_value=mock_response) as mock_urlopen: + with patch( + "codeatrium.llm._validate_palace", side_effect=validate_side_effect + ) as mock_validate: + result = _call_openai("prompt", backend) + + # urlopen が2回呼ばれたことを確認 + assert mock_urlopen.call_count == 2 + + # _validate_palace が2回呼ばれたことを確認 + assert mock_validate.call_count == 2 + + # 結果が返されたことを確認 + assert result == MOCK_JSON_RESPONSE["structured_output"] + + +def test_call_openai_raises_after_two_validation_failures() -> None: + """ + _call_openai が validation が2回失敗した場合、LLMValidationError を発生させることを確認する + """ + backend = DistillBackend( + provider="openai", model="m", base_url="http://localhost:11434/v1" + ) + + mock_response = MagicMock() + response_body = json.dumps( + { + "choices": [ + { + "message": { + "content": json.dumps(MOCK_JSON_RESPONSE["structured_output"]) + } + } + ] + } + ) + mock_response.__enter__ = MagicMock(return_value=mock_response) + mock_response.__exit__ = MagicMock(return_value=None) + mock_response.read.return_value = response_body.encode() + + with patch("urllib.request.urlopen", return_value=mock_response): + with patch( + "codeatrium.llm._validate_palace", + side_effect=LLMValidationError("always fails"), + ): + with pytest.raises(LLMValidationError): + _call_openai("prompt", backend) + + +def test_call_openai_raises_on_connection_error() -> None: + """ + _call_openai が urllib.error.URLError を RuntimeError に変換することを確認する + """ + backend = DistillBackend( + provider="openai", model="m", base_url="http://localhost:11434/v1" + ) + + with patch( + "urllib.request.urlopen", + side_effect=urllib.error.URLError("connection refused"), + ): + with pytest.raises(RuntimeError): + _call_openai("prompt", backend)