codeatrium 0.1.0__py3-none-any.whl

codeatrium/cli/server_cmd.py

@@ -0,0 +1,122 @@
+"""loci server start/stop/status コマンド"""
+
+from __future__ import annotations
+
+import typer
+
+server_app = typer.Typer(help="embedding サーバー管理")
+
+
+@server_app.command("start")
+def server_start() -> None:
+    """embedding サーバーをバックグラウンドで起動する"""
+    import json as _json
+    import socket as _socket
+    import subprocess
+
+    from codeatrium.embedder import _loci_python
+    from codeatrium.paths import db_path, find_project_root, server_pid_path, sock_path
+
+    root = find_project_root()
+    if not db_path(root).exists():
+        typer.echo("Not initialized. Run `loci init` first.", err=True)
+        raise typer.Exit(1)
+
+    sock = sock_path(root)
+
+    if sock.exists():
+        try:
+            with _socket.socket(_socket.AF_UNIX, _socket.SOCK_STREAM) as s:
+                s.settimeout(1.0)
+                s.connect(str(sock))
+                s.sendall((_json.dumps({"type": "ping"}) + "\n").encode())
+                resp = s.recv(256)
+                if b"ok" in resp:
+                    typer.echo("Server is already running.")
+                    return
+        except Exception:
+            sock.unlink(missing_ok=True)
+
+    pid_path = server_pid_path(root)
+    proc = subprocess.Popen(
+        [_loci_python(), "-m", "codeatrium.embedder_server", str(sock)],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+        start_new_session=True,
+    )
+    pid_path.write_text(str(proc.pid))
+
+    import time
+
+    for i in range(150):
+        if sock.exists():
+            typer.echo(f"Server started (PID {proc.pid})")
+            return
+        time.sleep(0.2)
+        if i % 25 == 24:
+            typer.echo("  Loading model...", err=True)
+
+    typer.echo("Server failed to start.", err=True)
+    raise typer.Exit(1)
+
+
+@server_app.command("stop")
+def server_stop() -> None:
+    """embedding サーバーを停止する"""
+    import json as _json
+    import socket as _socket
+
+    from codeatrium.paths import find_project_root, server_pid_path, sock_path
+
+    root = find_project_root()
+    sock = sock_path(root)
+
+    if not sock.exists():
+        typer.echo("Server is not running.")
+        return
+
+    try:
+        with _socket.socket(_socket.AF_UNIX, _socket.SOCK_STREAM) as s:
+            s.settimeout(2.0)
+            s.connect(str(sock))
+            s.sendall((_json.dumps({"type": "stop"}) + "\n").encode())
+        typer.echo("Server stopped.")
+    except Exception as e:
+        typer.echo(f"Could not connect to server: {e}", err=True)
+        sock.unlink(missing_ok=True)
+
+    server_pid_path(root).unlink(missing_ok=True)
+
+
+@server_app.command("status")
+def server_status() -> None:
+    """embedding サーバーの状態を確認する"""
+    import json as _json
+    import socket as _socket
+
+    from codeatrium.paths import find_project_root, server_pid_path, sock_path
+
+    root = find_project_root()
+    sock = sock_path(root)
+
+    if not sock.exists():
+        typer.echo("Server: stopped")
+        return
+
+    try:
+        with _socket.socket(_socket.AF_UNIX, _socket.SOCK_STREAM) as s:
+            s.settimeout(1.0)
+            s.connect(str(sock))
+            s.sendall((_json.dumps({"type": "ping"}) + "\n").encode())
+            resp = s.recv(256)
+            if b"ok" in resp:
+                pid_path = server_pid_path(root)
+                pid = pid_path.read_text().strip() if pid_path.exists() else "unknown"
+                typer.echo(f"Server: running (PID {pid})")
+                typer.echo(f"Socket: {sock}")
+                return
+    except Exception:
+        pass
+
+    typer.echo("Server: socket exists but not responding")
+    sock.unlink(missing_ok=True)

codeatrium/cli/show_cmd.py

@@ -0,0 +1,151 @@
+"""loci show / loci dump コマンド"""
+
+from __future__ import annotations
+
+import json
+from typing import Annotated, Any
+
+import typer
+
+
+def show(
+    ref: Annotated[str, typer.Argument(help="verbatim_ref (path:ply=N)")],
+    json_output: Annotated[bool, typer.Option("--json", help="JSON で出力")] = False,
+) -> None:
+    """verbatim_ref から exchange の原文を取得する"""
+    from codeatrium.db import get_connection
+    from codeatrium.paths import db_path, find_project_root
+
+    if ":ply=" not in ref:
+        typer.echo("Invalid ref format. Expected: <path>:ply=<N>", err=True)
+        raise typer.Exit(1)
+    path_part, ply_part = ref.rsplit(":ply=", 1)
+    try:
+        ply = int(ply_part)
+    except ValueError:
+        typer.echo(f"Invalid ply value: {ply_part}", err=True)
+        raise typer.Exit(1)
+
+    root = find_project_root()
+    db = db_path(root)
+    if not db.exists():
+        typer.echo("Not initialized. Run `loci init` first.", err=True)
+        raise typer.Exit(1)
+
+    con = get_connection(db)
+    row = con.execute(
+        """
+        SELECT e.user_content, e.agent_content, e.ply_start, e.ply_end
+        FROM exchanges e
+        JOIN conversations c ON c.id = e.conversation_id
+        WHERE c.source_path = ? AND e.ply_start = ?
+        """,
+        (path_part, ply),
+    ).fetchone()
+    con.close()
+
+    if row is None:
+        typer.echo("Exchange not found.")
+        return
+
+    if json_output:
+        typer.echo(
+            json.dumps(
+                {
+                    "user_content": row["user_content"],
+                    "agent_content": row["agent_content"],
+                    "ply_start": row["ply_start"],
+                    "ply_end": row["ply_end"],
+                },
+                ensure_ascii=False,
+                indent=2,
+            )
+        )
+    else:
+        typer.echo(f"[User] (ply {row['ply_start']}-{row['ply_end']})")
+        typer.echo(row["user_content"])
+        typer.echo("\n[Agent]")
+        typer.echo(row["agent_content"])
+
+
+def dump(
+    distilled: Annotated[
+        bool, typer.Option("--distilled", help="蒸留済み palace objects を出力")
+    ] = False,
+    limit: Annotated[int, typer.Option("--limit", "-n", help="最大件数")] = 1000,
+    json_output: Annotated[bool, typer.Option("--json", help="JSON で出力")] = False,
+) -> None:
+    """蒸留済み palace objects を新しい順に出力する（セッション開始時の in-context ロード用）"""
+    from codeatrium.db import get_connection
+    from codeatrium.paths import db_path, find_project_root
+
+    if not distilled:
+        typer.echo("Use --distilled to dump palace objects.", err=True)
+        raise typer.Exit(1)
+
+    root = find_project_root()
+    db = db_path(root)
+    if not db.exists():
+        typer.echo("Not initialized. Run `loci init` first.", err=True)
+        raise typer.Exit(1)
+
+    con = get_connection(db)
+    rows = con.execute(
+        """
+        SELECT p.id, p.exchange_id, p.exchange_core, p.specific_context,
+               e.distilled_at
+        FROM palace_objects p
+        JOIN exchanges e ON e.id = p.exchange_id
+        ORDER BY e.distilled_at DESC
+        LIMIT ?
+        """,
+        (limit,),
+    ).fetchall()
+
+    if not rows:
+        typer.echo("No distilled objects found.")
+        con.close()
+        return
+
+    palace_ids = [r["id"] for r in rows]
+    placeholders = ",".join("?" * len(palace_ids))
+    room_rows = con.execute(
+        f"""
+        SELECT palace_object_id, room_type, room_key, room_label
+        FROM rooms
+        WHERE palace_object_id IN ({placeholders})
+        ORDER BY relevance DESC
+        """,
+        palace_ids,
+    ).fetchall()
+    con.close()
+
+    rooms_map: dict[str, list[Any]] = {}
+    for r in room_rows:
+        rooms_map.setdefault(r["palace_object_id"], []).append(
+            {
+                "room_type": r["room_type"],
+                "room_key": r["room_key"],
+                "room_label": r["room_label"],
+            }
+        )
+
+    if json_output:
+        output = [
+            {
+                "exchange_core": r["exchange_core"],
+                "specific_context": r["specific_context"],
+                "rooms": rooms_map.get(r["id"], []),
+                "date": (r["distilled_at"] or "")[:10],
+            }
+            for r in rows
+        ]
+        typer.echo(json.dumps(output, ensure_ascii=False, indent=2))
+    else:
+        for r in rows:
+            date = (r["distilled_at"] or "")[:10]
+            typer.echo(f"\n[{date}] {r['exchange_core']}")
+            if r["specific_context"]:
+                typer.echo(f"  {r['specific_context']}")
+            for rm in rooms_map.get(r["id"], [])[:2]:
+                typer.echo(f"  #{rm['room_key']}")

codeatrium/cli/status_cmd.py

@@ -0,0 +1,59 @@
+"""loci status コマンド"""
+
+from __future__ import annotations
+
+import json
+from typing import Annotated
+
+import typer
+
+
+def status(
+    json_output: Annotated[bool, typer.Option("--json", help="JSON で出力")] = False,
+) -> None:
+    """インデックス状態（exchange 数・蒸留済み数・DB サイズ）を表示する"""
+    from codeatrium.db import get_connection
+    from codeatrium.paths import db_path, find_project_root
+
+    root = find_project_root()
+    db = db_path(root)
+
+    if not db.exists():
+        typer.echo("Not initialized. Run `loci init` first.", err=True)
+        raise typer.Exit(1)
+
+    con = get_connection(db)
+    total = con.execute("SELECT COUNT(*) FROM exchanges").fetchone()[0]
+    distilled = con.execute(
+        "SELECT COUNT(*) FROM exchanges WHERE distilled_at IS NOT NULL"
+    ).fetchone()[0]
+    palace_count = con.execute("SELECT COUNT(*) FROM palace_objects").fetchone()[0]
+    symbol_count = con.execute("SELECT COUNT(*) FROM symbols").fetchone()[0]
+    con.close()
+
+    db_size_bytes = db.stat().st_size
+    db_size_kb = db_size_bytes / 1024
+
+    if json_output:
+        typer.echo(
+            json.dumps(
+                {
+                    "db_path": str(db),
+                    "exchanges": total,
+                    "distilled": distilled,
+                    "undistilled": total - distilled,
+                    "palace_objects": palace_count,
+                    "symbols": symbol_count,
+                    "db_size_kb": round(db_size_kb, 1),
+                },
+                ensure_ascii=False,
+                indent=2,
+            )
+        )
+    else:
+        typer.echo(f"DB: {db} ({db_size_kb:.1f} KB)")
+        typer.echo(
+            f"Exchanges : {total} total, {distilled} distilled, {total - distilled} pending"
+        )
+        typer.echo(f"Palace    : {palace_count}")
+        typer.echo(f"Symbols   : {symbol_count}")

codeatrium/config.py

@@ -0,0 +1,96 @@
+"""設定ファイルの読み込み — .codeatrium/config.toml"""
+
+from __future__ import annotations
+
+import tomllib
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+CONFIG_FILENAME = "config.toml"
+
+# ---- デフォルト値 ----
+
+DEFAULT_DISTILL_MODEL = "claude-haiku-4-5-20251001"
+DEFAULT_DISTILL_BATCH_LIMIT = 20
+DEFAULT_INDEX_MIN_CHARS = 50
+DEFAULT_DISTILL_MIN_CHARS = 100
+
+
+@dataclass
+class Config:
+    """ユーザー設定"""
+
+    distill_model: str = DEFAULT_DISTILL_MODEL
+    distill_batch_limit: int = DEFAULT_DISTILL_BATCH_LIMIT
+    index_min_chars: int = DEFAULT_INDEX_MIN_CHARS
+    distill_min_chars: int = DEFAULT_DISTILL_MIN_CHARS
+
+
+def load_config(project_root: Path) -> Config:
+    """project_root/.codeatrium/config.toml を読んで Config を返す。
+    ファイルがなければデフォルト。不正な値は警告してデフォルトにフォールバック。
+    """
+    config_path = project_root / ".codeatrium" / CONFIG_FILENAME
+    if not config_path.exists():
+        return Config()
+
+    try:
+        with config_path.open("rb") as f:
+            data = tomllib.load(f)
+    except Exception as e:
+        import sys
+
+        print(f"Warning: failed to parse {config_path}: {e}", file=sys.stderr)
+        return Config()
+
+    distill: dict[str, Any] = data.get("distill", {})
+
+    model = distill.get("model", DEFAULT_DISTILL_MODEL)
+    if not isinstance(model, str) or not model.strip():
+        import sys
+
+        print(
+            "Warning: distill.model must be a non-empty string, using default.",
+            file=sys.stderr,
+        )
+        model = DEFAULT_DISTILL_MODEL
+
+    batch_limit = distill.get("batch_limit", DEFAULT_DISTILL_BATCH_LIMIT)
+    if not isinstance(batch_limit, int) or batch_limit < 1:
+        import sys
+
+        print(
+            "Warning: distill.batch_limit must be a positive integer, using default.",
+            file=sys.stderr,
+        )
+        batch_limit = DEFAULT_DISTILL_BATCH_LIMIT
+
+    index: dict[str, Any] = data.get("index", {})
+
+    min_chars = index.get("min_chars", DEFAULT_INDEX_MIN_CHARS)
+    if not isinstance(min_chars, int) or min_chars < 1:
+        import sys
+
+        print(
+            "Warning: index.min_chars must be a positive integer, using default.",
+            file=sys.stderr,
+        )
+        min_chars = DEFAULT_INDEX_MIN_CHARS
+
+    distill_min_chars = distill.get("min_chars", DEFAULT_DISTILL_MIN_CHARS)
+    if not isinstance(distill_min_chars, int) or distill_min_chars < 1:
+        import sys
+
+        print(
+            "Warning: distill.min_chars must be a positive integer, using default.",
+            file=sys.stderr,
+        )
+        distill_min_chars = DEFAULT_DISTILL_MIN_CHARS
+
+    return Config(
+        distill_model=model,
+        distill_batch_limit=batch_limit,
+        index_min_chars=min_chars,
+        distill_min_chars=distill_min_chars,
+    )

codeatrium/db.py

@@ -0,0 +1,135 @@
+"""
+SQLite DB の初期化・スキーマ定義・接続管理
+
+テーブル構成:
+  conversations  - .jsonl ファイル単位の会話記録（重複排除キャッシュ）
+  exchanges      - exchange 単位の verbatim テキスト
+  exchanges_fts  - exchanges の FTS5 仮想テーブル（BM25 verbatim 検索用）
+  vec_exchanges  - sqlite-vec HNSW インデックス（Phase1 verbatim ベクトル検索用）
+  palace_objects - 蒸留済み palace object（exchange_core + specific_context）
+  rooms          - palace object の room_assignments
+  vec_palace     - sqlite-vec HNSW インデックス（Phase2 distilled ベクトル検索用）
+  symbols        - tree-sitter 解決済みシンボル（Phase3 コード逆引き用）
+"""
+
+import sqlite3
+from pathlib import Path
+
+import sqlite_vec
+
+
+def get_connection(db_path: Path) -> sqlite3.Connection:
+    """sqlite-vec 拡張をロードした接続を返す"""
+    con = sqlite3.connect(db_path)
+    con.enable_load_extension(True)
+    sqlite_vec.load(con)
+    con.enable_load_extension(False)
+    con.row_factory = sqlite3.Row
+    return con
+
+
+def init_db(db_path: Path) -> None:
+    """DB を初期化してスキーマを作成する（冪等）"""
+    db_path.parent.mkdir(parents=True, exist_ok=True)
+    con = get_connection(db_path)
+
+    con.executescript("""
+        CREATE TABLE IF NOT EXISTS conversations (
+            id            TEXT PRIMARY KEY,   -- sha256(source_path)
+            source_path   TEXT NOT NULL UNIQUE,
+            started_at    TIMESTAMP,
+            last_ply_end  INT  NOT NULL DEFAULT -1  -- 最後にインデックスした ply_end（差分用）
+        );
+
+        CREATE TABLE IF NOT EXISTS exchanges (
+            id              TEXT PRIMARY KEY,  -- sha256(conversation_id + ":" + user_uuid)
+            conversation_id TEXT NOT NULL,
+            ply_start       INT  NOT NULL,
+            ply_end         INT  NOT NULL,
+            user_content    TEXT NOT NULL,
+            agent_content   TEXT NOT NULL,
+            distilled_at    TIMESTAMP          -- NULL = 未蒸留
+        );
+
+        CREATE VIRTUAL TABLE IF NOT EXISTS exchanges_fts USING fts5(
+            user_content,
+            agent_content,
+            content=exchanges,
+            content_rowid=rowid
+        );
+
+        CREATE TRIGGER IF NOT EXISTS exchanges_ai
+        AFTER INSERT ON exchanges BEGIN
+            INSERT INTO exchanges_fts(rowid, user_content, agent_content)
+            VALUES (new.rowid, new.user_content, new.agent_content);
+        END;
+
+        CREATE TRIGGER IF NOT EXISTS exchanges_ad
+        AFTER DELETE ON exchanges BEGIN
+            INSERT INTO exchanges_fts(exchanges_fts, rowid, user_content, agent_content)
+            VALUES ('delete', old.rowid, old.user_content, old.agent_content);
+        END;
+
+        CREATE TRIGGER IF NOT EXISTS exchanges_au
+        AFTER UPDATE ON exchanges BEGIN
+            INSERT INTO exchanges_fts(exchanges_fts, rowid, user_content, agent_content)
+            VALUES ('delete', old.rowid, old.user_content, old.agent_content);
+            INSERT INTO exchanges_fts(rowid, user_content, agent_content)
+            VALUES (new.rowid, new.user_content, new.agent_content);
+        END;
+
+        CREATE TABLE IF NOT EXISTS palace_objects (
+            id               TEXT PRIMARY KEY,
+            exchange_id      TEXT NOT NULL,
+            exchange_core    TEXT NOT NULL,
+            specific_context TEXT NOT NULL,
+            distill_text     TEXT NOT NULL    -- exchange_core + newline + specific_context
+        );
+
+        CREATE TABLE IF NOT EXISTS rooms (
+            id               TEXT PRIMARY KEY,
+            palace_object_id TEXT NOT NULL,
+            room_type        TEXT NOT NULL,   -- "file" / "concept" / "workflow"
+            room_key         TEXT NOT NULL,
+            room_label       TEXT NOT NULL,
+            relevance        REAL NOT NULL,
+            dedup_hash       TEXT NOT NULL    -- hash(room_type, room_key)
+        );
+
+        CREATE TABLE IF NOT EXISTS symbols (
+            id               TEXT PRIMARY KEY,   -- sha256(symbol_name + file_path)
+            palace_object_id TEXT NOT NULL,
+            symbol_name      TEXT NOT NULL,       -- "AuthMiddleware.validate"
+            symbol_kind      TEXT NOT NULL,       -- "function" / "class" / "method"
+            file_path        TEXT NOT NULL,
+            signature        TEXT NOT NULL,
+            line             INT  NOT NULL,
+            dedup_hash       TEXT NOT NULL        -- sha256(symbol_name + file_path)
+        );
+    """)
+
+    # マイグレーション: last_ply_end カラムが無い既存 DB に追加
+    try:
+        con.execute("ALTER TABLE conversations ADD COLUMN last_ply_end INT NOT NULL DEFAULT -1")
+        con.commit()
+    except Exception:
+        pass  # カラムが既に存在する場合は無視
+
+    # sqlite-vec の仮想テーブル（HNSW, Phase1 verbatim embedding 用）
+    con.execute("""
+        CREATE VIRTUAL TABLE IF NOT EXISTS vec_exchanges USING vec0(
+            exchange_id TEXT PRIMARY KEY,
+            embedding   FLOAT[384]
+        )
+    """)
+
+    # sqlite-vec の仮想テーブル（HNSW, Phase2 distilled embedding 用）
+    con.execute("""
+        CREATE VIRTUAL TABLE IF NOT EXISTS vec_palace USING vec0(
+            palace_id TEXT PRIMARY KEY,
+            embedding FLOAT[384]
+        )
+    """)
+
+    con.commit()
+    con.close()