codeatrium 0.1.0__py3-none-any.whl

codeatrium/distiller.py

@@ -0,0 +1,290 @@
+"""蒸留モジュール: claude -p で exchange を palace object に変換する
+
+SPEC Section 6 DISTILLER フロー準拠:
+  ① files_touched を regex で抽出（LLM非使用）
+  ② claude -p で palace object 生成（--output-format json --json-schema）
+  ③ distill_text を embedding して vec_palace に登録
+  ④ files_touched を tree-sitter で解析してシンボルを symbols テーブルに登録
+"""
+
+from __future__ import annotations
+
+import datetime
+import hashlib
+import os
+import re
+import struct
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+
+from codeatrium.embedder import Embedder
+from codeatrium.llm import DISTILL_PROMPT_TEMPLATE, call_claude
+from codeatrium.models import PalaceObject
+
+# ---- ファイルパス抽出 ----
+
+_FILES_PATTERN = re.compile(
+    r"(/(?:[a-zA-Z0-9._\-]+/)*[a-zA-Z0-9._\-]+\.[a-zA-Z0-9]+)"  # 絶対パス
+    r"|([a-zA-Z0-9._\-]+(?:/[a-zA-Z0-9._\-]+)+\.[a-zA-Z0-9]+)"  # 相対パス（1段以上のディレクトリ）
+)
+
+
+# ---- 内部ヘルパー ----
+
+
+def _sha256(text: str) -> str:
+    return hashlib.sha256(text.encode()).hexdigest()
+
+
+# ---- 公開 API ----
+
+
+_EXTERNAL_PATH_MARKERS = (
+    "site-packages/",
+    "dist-packages/",
+    "/lib/python",
+    "/opt/",
+    "/usr/lib/",
+    "/usr/local/lib/",
+    ".venv/",
+    "/venv/",
+    "node_modules/",
+)
+
+
+def _is_external_path(path: str, project_root: str | None = None) -> bool:
+    """プロジェクト外のパスか判定する。
+
+    絶対パス: project_root が指定されていればその配下かチェック。
+    相対パス: ハードコードマーカーでフィルタ。
+    """
+    if path.startswith("/"):
+        # 絶対パス: project_root 配下でなければ外部
+        if project_root:
+            return not path.startswith(project_root)
+        # project_root 不明時はマーカーでフォールバック
+    return any(marker in path for marker in _EXTERNAL_PATH_MARKERS)
+
+
+def extract_files_touched(
+    user_content: str, agent_content: str, project_root: str | None = None
+) -> list[str]:
+    """user_content + agent_content から regex でファイルパスを抽出する（重複排除・順序維持）
+
+    project_root が指定された場合、絶対パスはその配下のもののみ残す。
+    相対パスはハードコードマーカー（node_modules 等）でフィルタする。
+    """
+    text = user_content + "\n" + agent_content
+    # project_root を末尾スラッシュ付きに正規化
+    root_prefix = (project_root.rstrip("/") + "/") if project_root else None
+    seen: set[str] = set()
+    result: list[str] = []
+    for m in _FILES_PATTERN.findall(text):
+        path = m[0] or m[1]
+        if path and path not in seen and not _is_external_path(path, root_prefix):
+            seen.add(path)
+            result.append(path)
+    return result
+
+
+def distill_exchange(
+    exchange_id: str,
+    user_content: str,
+    agent_content: str,
+    ply_start: int,
+    ply_end: int,
+    model: str | None = None,
+    project_root: str | None = None,
+) -> PalaceObject:
+    """1つの exchange を蒸留して PalaceObject を返す"""
+    messages_text = (user_content + "\n" + agent_content)[:4000]
+    prompt = DISTILL_PROMPT_TEMPLATE.format(
+        ply_start=ply_start,
+        ply_end=ply_end,
+        messages_text=messages_text,
+    )
+    raw = call_claude(prompt, model=model)
+    files_touched = extract_files_touched(
+        user_content, agent_content, project_root=project_root
+    )
+    return PalaceObject(
+        exchange_core=raw["exchange_core"],
+        specific_context=raw["specific_context"],
+        room_assignments=raw["room_assignments"],
+        files_touched=files_touched,
+    )
+
+
+def save_palace_object(
+    db_path: Path,
+    exchange_id: str,
+    palace: PalaceObject,
+    embedding: Any,  # np.ndarray
+) -> None:
+    """PalaceObject を DB に保存し exchange の distilled_at を更新する"""
+    import numpy as np
+
+    from codeatrium.db import get_connection
+
+    palace_id = _sha256(f"palace:{exchange_id}")
+    distill_text = palace.exchange_core + "\n" + palace.specific_context
+
+    con = get_connection(db_path)
+
+    con.execute(
+        """
+        INSERT OR IGNORE INTO palace_objects
+            (id, exchange_id, exchange_core, specific_context, distill_text)
+        VALUES (?, ?, ?, ?, ?)
+        """,
+        (
+            palace_id,
+            exchange_id,
+            palace.exchange_core,
+            palace.specific_context,
+            distill_text,
+        ),
+    )
+
+    for room in palace.room_assignments:
+        dedup = _sha256(f"{room['room_type']}:{room['room_key']}")
+        room_id = _sha256(f"{palace_id}:{dedup}")
+        con.execute(
+            """
+            INSERT OR IGNORE INTO rooms
+                (id, palace_object_id, room_type, room_key, room_label, relevance, dedup_hash)
+            VALUES (?, ?, ?, ?, ?, ?, ?)
+            """,
+            (
+                room_id,
+                palace_id,
+                room["room_type"],
+                room["room_key"],
+                room["room_label"],
+                room["relevance"],
+                dedup,
+            ),
+        )
+
+    arr = embedding.astype(np.float32)
+    blob = struct.pack(f"{len(arr)}f", *arr.tolist())
+    exists = con.execute(
+        "SELECT 1 FROM vec_palace WHERE palace_id = ?", (palace_id,)
+    ).fetchone()
+    if not exists:
+        con.execute(
+            "INSERT INTO vec_palace (palace_id, embedding) VALUES (?, ?)",
+            (palace_id, blob),
+        )
+
+    # ⑤ tree-sitter シンボル解決
+    from codeatrium.resolver import SymbolResolver
+
+    resolver = SymbolResolver()
+    for file_str in palace.files_touched:
+        for sym in resolver.extract(Path(file_str)):
+            sym_id = _sha256(f"{sym.symbol_name}:{sym.file_path}")
+            con.execute(
+                """
+                INSERT OR IGNORE INTO symbols
+                    (id, palace_object_id, symbol_name, symbol_kind,
+                     file_path, signature, line, dedup_hash)
+                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+                """,
+                (
+                    sym_id,
+                    palace_id,
+                    sym.symbol_name,
+                    sym.symbol_kind,
+                    sym.file_path,
+                    sym.signature,
+                    sym.line,
+                    sym_id,
+                ),
+            )
+
+    con.execute(
+        "UPDATE exchanges SET distilled_at = ? WHERE id = ?",
+        (datetime.datetime.utcnow().isoformat(), exchange_id),
+    )
+
+    con.commit()
+    con.close()
+
+
+def distill_all(
+    db_path: Path,
+    limit: int | None = None,
+    model: str | None = None,
+    on_progress: Callable[..., None] | None = None,
+    project_root: str | None = None,
+    distill_min_chars: int = 100,
+) -> int:
+    """未蒸留の exchange を処理する。
+
+    distill_min_chars: この文字数未満の exchange は蒸留スキップ（デフォルト100）
+    on_progress: (current, total, error=None) を受け取るコールバック
+    Returns: 処理した exchange 数
+    """
+    from codeatrium.db import get_connection
+
+    con = get_connection(db_path)
+
+    # 蒸留対象外の exchange を skipped にマーク:
+    # - 1-exchange セッション
+    # - distill_min_chars 未満（ワンフレーズ指示・システムメッセージ等）
+    con.execute("""
+        UPDATE exchanges SET distilled_at = 'skipped'
+        WHERE distilled_at IS NULL
+          AND ((SELECT COUNT(*) FROM exchanges e2
+                WHERE e2.conversation_id = exchanges.conversation_id) < 2
+               OR LENGTH(user_content) + LENGTH(agent_content) < ?)
+    """, (distill_min_chars,))
+    con.commit()
+
+    query = """
+        SELECT e.id, e.user_content, e.agent_content, e.ply_start, e.ply_end
+        FROM exchanges e
+        WHERE e.distilled_at IS NULL
+    """
+    params: list[int] = []
+    if limit is not None:
+        query += " LIMIT ?"
+        params.append(int(limit))
+    rows = con.execute(query, params).fetchall()
+    con.close()
+
+    if not rows:
+        return 0
+
+    total = len(rows)
+    embedder = Embedder()
+    count = 0
+    errors = 0
+    for row in rows:
+        try:
+            palace = distill_exchange(
+                row["id"],
+                row["user_content"],
+                row["agent_content"],
+                row["ply_start"],
+                row["ply_end"],
+                model=model,
+                project_root=project_root,
+            )
+            distill_text = palace.exchange_core + "\n" + palace.specific_context
+            vec = embedder.embed_passage(distill_text)
+            save_palace_object(db_path, row["id"], palace, vec)
+            count += 1
+        except Exception as e:
+            errors += 1
+            if on_progress is not None:
+                on_progress(count, total, error=str(e))
+            continue
+        if on_progress is not None:
+            on_progress(count, total)
+
+    return count

codeatrium/embedder.py

@@ -0,0 +1,168 @@
+"""
+sentence-transformers / multilingual-e5-small の embedding ラッパー
+
+モデル: intfloat/multilingual-e5-small（384次元・日本語+英語混在対応・CPU動作）
+
+ソケットサーバー方式:
+  - .codeatrium/embedder.sock が存在すれば Unix ソケット経由で embed（< 1秒）
+  - ソケットなければ直接モデルをロード（~7秒）＋バックグラウンドでサーバー起動
+  - 2回目以降は常にソケット経由になるため高速
+
+cold start が問題になった場合は SentenceTransformer(..., backend="onnx") でも対処可能
+"""
+
+from __future__ import annotations
+
+import json
+import socket
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+
+
+def _loci_python() -> str:
+    """venv の Python パスを返す（sys.executable はシステム Python の場合があるため）"""
+    # loci CLI と同じ venv の python を使う
+    python = Path(sys.executable).parent / "python3"
+    if python.exists():
+        return str(python)
+    return sys.executable
+
+
+MODEL_NAME = "intfloat/multilingual-e5-small"
+CONNECT_TIMEOUT = 2.0  # ソケット接続タイムアウト
+RECV_TIMEOUT = 30.0  # 埋め込み受信タイムアウト
+
+
+def _sock_path_from_env() -> Path | None:
+    """環境変数 CODEATRIUM_SOCK_PATH があれば優先使用（テスト用）。
+    CODEATRIUM_NO_SOCK=1 の場合はソケット無効（サーバー内自己接続デッドロック防止）。
+    """
+    import os
+
+    if os.environ.get("CODEATRIUM_NO_SOCK"):
+        return None
+    p = os.environ.get("CODEATRIUM_SOCK_PATH")
+    return Path(p) if p else None
+
+
+def _find_sock_path() -> Path | None:
+    """DB の親ディレクトリの embedder.sock を探す"""
+    import os
+
+    if os.environ.get("CODEATRIUM_NO_SOCK"):
+        return None
+    # .codeatrium/ の場所を git root から解決
+    try:
+        import subprocess as sp
+
+        result = sp.run(
+            ["git", "rev-parse", "--show-toplevel"],
+            capture_output=True,
+            text=True,
+        )
+        if result.returncode == 0:
+            return Path(result.stdout.strip()) / ".codeatrium" / "embedder.sock"
+    except Exception:
+        pass
+    return None
+
+
+def _try_socket_embed(sock_path: Path, req_type: str, text: str) -> np.ndarray | None:
+    """ソケットサーバーに接続して embedding を取得する。失敗時は None を返す。"""
+    if not sock_path.exists():
+        return None
+    try:
+        with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s:
+            s.settimeout(CONNECT_TIMEOUT)
+            s.connect(str(sock_path))
+            s.settimeout(RECV_TIMEOUT)
+            req = json.dumps({"type": req_type, "text": text}) + "\n"
+            s.sendall(req.encode())
+            buf = b""
+            while b"\n" not in buf:
+                chunk = s.recv(65536)
+                if not chunk:
+                    break
+                buf += chunk
+            resp = json.loads(buf.split(b"\n")[0])
+            if "embedding" in resp:
+                return np.array(resp["embedding"], dtype=np.float32)
+    except Exception:
+        pass
+    return None
+
+
+def _start_server_background(sock_path: Path) -> None:
+    """embedder_server をバックグラウンドで起動する"""
+    try:
+        subprocess.Popen(
+            [_loci_python(), "-m", "codeatrium.embedder_server", str(sock_path)],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+            start_new_session=True,
+        )
+        # サーバーが起動するまで少し待つ
+        for _ in range(20):
+            if sock_path.exists():
+                break
+            time.sleep(0.2)
+    except Exception:
+        pass
+
+
+class Embedder:
+    """multilingual-e5-small の薄いラッパー。
+
+    ソケットサーバーが起動済みなら高速パス（Unix ソケット）を使い、
+    なければ直接モデルをロードしてバックグラウンドでサーバーを起動する。
+    """
+
+    def __init__(self, sock_path: Path | None = None) -> None:
+        self._sock_path: Path | None = (
+            _sock_path_from_env() or sock_path or _find_sock_path()
+        )
+        self._model = None  # 遅延ロード
+
+    def _ensure_model(self) -> None:
+        """直接モデルをロードする（ソケット不使用時のフォールバック）"""
+        if self._model is None:
+            from sentence_transformers import SentenceTransformer
+
+            self._model = SentenceTransformer(MODEL_NAME)
+
+    def _embed_via_socket_or_direct(
+        self, text: str, req_type: str, prefix: str
+    ) -> np.ndarray:
+        """ソケット優先・なければ直接ロード＋サーバー起動"""
+        # ① ソケット経由を試みる
+        if self._sock_path is not None:
+            vec = _try_socket_embed(self._sock_path, req_type, text)
+            if vec is not None:
+                return vec
+
+        # ② 直接ロード
+        self._ensure_model()
+        assert self._model is not None
+        result = self._model.encode(
+            [f"{prefix}{text}"],
+            normalize_embeddings=True,
+        )
+        vec = result[0].astype(np.float32)
+
+        # ③ バックグラウンドでサーバーを起動（次回から高速化）
+        if self._sock_path is not None and not self._sock_path.exists():
+            _start_server_background(self._sock_path)
+
+        return vec
+
+    def embed(self, text: str) -> np.ndarray:
+        """クエリ用 embedding（query: プレフィックス）"""
+        return self._embed_via_socket_or_direct(text, "query", "query: ")
+
+    def embed_passage(self, text: str) -> np.ndarray:
+        """インデックス登録用 embedding（passage: プレフィックス）"""
+        return self._embed_via_socket_or_direct(text, "passage", "passage: ")

codeatrium/embedder_server.py

@@ -0,0 +1,172 @@
+"""
+embedder_server.py — multilingual-e5-small を常駐させる Unix ソケットサーバー
+
+プロトコル:
+  request (改行区切り JSON):  {"type": "query"|"passage", "text": "..."}
+  response (改行区切り JSON): {"embedding": [0.1, 0.2, ...]}
+  特殊コマンド:               {"type": "ping"} → {"status": "ok"}
+                              {"type": "stop"} → サーバー終了
+
+ライフサイクル:
+  - loci server start でバックグラウンド起動
+  - IDLE_TIMEOUT 秒間リクエストなし → 自動終了
+  - loci server stop / SIGTERM で即終了
+  - ソケットファイルは .codeatrium/embedder.sock
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import signal
+import socket
+import sys
+import threading
+import time
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from codeatrium.embedder import Embedder
+
+IDLE_TIMEOUT = 600  # 10分間無リクエストで自動終了
+BACKLOG = 8
+RECV_BUFSIZE = 65536
+
+
+def _load_embedder() -> Embedder:
+    # サーバー内では直接モデルを使う（ソケット経由にすると自己接続デッドロック）
+
+    os.environ["CODEATRIUM_NO_SOCK"] = "1"
+    from codeatrium.embedder import Embedder
+
+    embedder = Embedder()
+    del os.environ["CODEATRIUM_NO_SOCK"]
+    return embedder
+
+
+def _handle_client(
+    conn: socket.socket,
+    embedder: Embedder,
+    last_activity: list[float],
+    stop_event: threading.Event,
+) -> None:
+    """1クライアント接続を処理する"""
+    try:
+        buf = b""
+        while True:
+            chunk = conn.recv(RECV_BUFSIZE)
+            if not chunk:
+                break
+            buf += chunk
+            while b"\n" in buf:
+                line, buf = buf.split(b"\n", 1)
+                line = line.strip()
+                if not line:
+                    continue
+                try:
+                    req = json.loads(line)
+                except json.JSONDecodeError:
+                    conn.sendall(json.dumps({"error": "invalid json"}).encode() + b"\n")
+                    continue
+
+                req_type = req.get("type", "")
+
+                if req_type == "ping":
+                    conn.sendall(json.dumps({"status": "ok"}).encode() + b"\n")
+                    last_activity[0] = time.monotonic()
+                    continue
+
+                if req_type == "stop":
+                    conn.sendall(json.dumps({"status": "stopping"}).encode() + b"\n")
+                    stop_event.set()
+                    return
+
+                text = req.get("text", "")
+                if not text:
+                    conn.sendall(json.dumps({"error": "missing text"}).encode() + b"\n")
+                    continue
+
+                if req_type == "query":
+                    vec = embedder.embed(text)
+                elif req_type == "passage":
+                    vec = embedder.embed_passage(text)
+                else:
+                    conn.sendall(
+                        json.dumps({"error": f"unknown type: {req_type}"}).encode()
+                        + b"\n"
+                    )
+                    continue
+
+                resp = json.dumps({"embedding": vec.tolist()})
+                conn.sendall(resp.encode() + b"\n")
+                last_activity[0] = time.monotonic()
+    except (OSError, BrokenPipeError):
+        pass
+    finally:
+        try:
+            conn.close()
+        except OSError:
+            pass
+
+
+def run_server(sock_path: Path) -> None:
+    """ソケットサーバーを起動してリクエストを処理する（ブロッキング）"""
+    # 既存ソケットファイルを削除
+    sock_path.unlink(missing_ok=True)
+
+    embedder = _load_embedder()
+
+    server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+    old_umask = os.umask(0o177)  # enforce 0o600 on the socket, no TOCTOU
+    try:
+        server.bind(str(sock_path))
+    finally:
+        os.umask(old_umask)
+    server.listen(BACKLOG)
+    server.settimeout(1.0)  # accept の timeout（idle チェック用）
+
+    last_activity: list[float] = [time.monotonic()]
+    stop_event = threading.Event()
+
+    def _sigterm_handler(signum: int, frame: object) -> None:
+        stop_event.set()
+
+    signal.signal(signal.SIGTERM, _sigterm_handler)
+    signal.signal(signal.SIGINT, _sigterm_handler)
+
+    try:
+        while not stop_event.is_set():
+            # idle timeout チェック
+            if time.monotonic() - last_activity[0] > IDLE_TIMEOUT:
+                break
+
+            try:
+                conn, _ = server.accept()
+            except TimeoutError:
+                continue
+            except OSError:
+                break
+
+            # クライアントごとにスレッド
+            t = threading.Thread(
+                target=_handle_client,
+                args=(conn, embedder, last_activity, stop_event),
+                daemon=True,
+            )
+            t.start()
+    finally:
+        server.close()
+        sock_path.unlink(missing_ok=True)
+
+
+def main() -> None:
+    """CLI エントリポイント: python -m codeatrium.embedder_server <sock_path>"""
+    if len(sys.argv) < 2:
+        print("Usage: python -m codeatrium.embedder_server <sock_path>", file=sys.stderr)
+        sys.exit(1)
+    run_server(Path(sys.argv[1]))
+
+
+if __name__ == "__main__":
+    main()