tokenjam 0.2.0__py3-none-any.whl

tokenjam/cli/cmd_export.py

@@ -0,0 +1,200 @@
+from __future__ import annotations
+
+import csv
+import io
+import json
+from pathlib import Path
+
+import click
+
+from tokenjam.core.models import TraceFilters
+from tokenjam.utils.formatting import console
+from tokenjam.utils.time_parse import parse_since
+
+
+@click.command("export")
+@click.option("--agent", default=None, help="Filter to specific agent_id")
+@click.option("--since", default="7d", help="Time window (e.g. 1h, 7d)")
+@click.option("--format", "fmt",
+              type=click.Choice(["json", "csv", "otlp", "openevals"]),
+              default="json")
+@click.option("--output", "output_path", default=None, help="Output file path (stdout if omitted)")
+@click.pass_context
+def cmd_export(ctx: click.Context, agent: str | None, since: str,
+               fmt: str, output_path: str | None) -> None:
+    """Export spans in various formats."""
+    db = ctx.obj["db"]
+    agent_filter = agent or ctx.obj.get("agent")
+    filters = TraceFilters(
+        agent_id=agent_filter,
+        since=parse_since(since),
+        limit=10000,
+    )
+    traces = db.get_traces(filters)
+
+    if fmt == "json":
+        output = _export_json(db, traces)
+    elif fmt == "csv":
+        output = _export_csv(db, traces)
+    elif fmt == "otlp":
+        _export_otlp(ctx, db, traces)
+        return
+    elif fmt == "openevals":
+        output = _export_openevals(db, traces)
+    else:
+        console.print(f"[red]Unknown format: {fmt}[/red]")
+        return
+
+    if output_path:
+        Path(output_path).write_text(output)
+        console.print(f"[green]Exported to {output_path}[/green]")
+    else:
+        click.echo(output)
+
+
+def _export_json(db: object, traces: list) -> str:
+    lines = []
+    seen_traces: set[str] = set()
+    for t in traces:
+        if t.trace_id in seen_traces:
+            continue
+        seen_traces.add(t.trace_id)
+        spans = db.get_trace_spans(t.trace_id)
+        for s in spans:
+            lines.append(json.dumps({
+                "span_id": s.span_id,
+                "trace_id": s.trace_id,
+                "agent_id": s.agent_id,
+                "name": s.name,
+                "start_time": s.start_time.isoformat() if s.start_time else None,
+                "end_time": s.end_time.isoformat() if s.end_time else None,
+                "duration_ms": s.duration_ms,
+                "cost_usd": s.cost_usd,
+                "input_tokens": s.input_tokens,
+                "output_tokens": s.output_tokens,
+                "status_code": s.status_code.value,
+                "provider": s.provider,
+                "model": s.model,
+                "tool_name": s.tool_name,
+                "attributes": s.attributes,
+            }, default=str))
+    return "\n".join(lines)
+
+
+def _export_csv(db: object, traces: list) -> str:
+    buf = io.StringIO()
+    writer = csv.writer(buf)
+    writer.writerow([
+        "span_id", "trace_id", "agent_id", "name", "start_time",
+        "duration_ms", "cost_usd", "input_tokens", "output_tokens", "status_code",
+    ])
+    seen_traces: set[str] = set()
+    for t in traces:
+        if t.trace_id in seen_traces:
+            continue
+        seen_traces.add(t.trace_id)
+        spans = db.get_trace_spans(t.trace_id)
+        for s in spans:
+            writer.writerow([
+                s.span_id, s.trace_id, s.agent_id, s.name,
+                s.start_time.isoformat() if s.start_time else "",
+                s.duration_ms, s.cost_usd,
+                s.input_tokens, s.output_tokens, s.status_code.value,
+            ])
+    return buf.getvalue()
+
+
+_KIND_MAP = {"internal": 1, "server": 2, "client": 3, "producer": 4, "consumer": 5}
+
+
+def _export_otlp(ctx: click.Context, db: object, traces: list) -> None:
+    config = ctx.obj["config"]
+    if not config.export.otlp.enabled:
+        console.print("[red]OTLP export is not enabled. Set export.otlp.enabled = true "
+                      "in config.[/red]")
+        return
+
+    import httpx
+    endpoint = config.export.otlp.endpoint.rstrip("/") + "/v1/traces"
+    headers = dict(config.export.otlp.headers)
+    headers.setdefault("Content-Type", "application/json")
+
+    seen_traces: set[str] = set()
+    succeeded = 0
+    failed = 0
+    for t in traces:
+        if t.trace_id in seen_traces:
+            continue
+        seen_traces.add(t.trace_id)
+        spans = db.get_trace_spans(t.trace_id)
+        payload = {
+            "resourceSpans": [{
+                "resource": {"attributes": [
+                    {"key": "service.name", "value": {"stringValue": "tokenjam"}},
+                ]},
+                "scopeSpans": [{
+                    "spans": [
+                        {
+                            "traceId": s.trace_id,
+                            "spanId": s.span_id,
+                            "name": s.name,
+                            "kind": _KIND_MAP.get(s.kind.value, 1) if s.kind else 1,
+                            "startTimeUnixNano": str(int(s.start_time.timestamp() * 1e9))
+                            if s.start_time else "0",
+                            "endTimeUnixNano": str(int(s.end_time.timestamp() * 1e9))
+                            if s.end_time else "0",
+                        }
+                        for s in spans
+                    ],
+                }],
+            }],
+        }
+        resp = httpx.post(endpoint, json=payload, headers=headers)
+        if resp.status_code >= 400:
+            console.print(f"[red]OTLP export failed for trace {t.trace_id}: "
+                          f"{resp.status_code}[/red]")
+            failed += 1
+        else:
+            succeeded += 1
+
+    console.print(f"[green]Exported {succeeded} traces to {endpoint}[/green]"
+                  + (f" ({failed} failed)" if failed else ""))
+
+
+def _export_openevals(db: object, traces: list) -> str:
+    results = []
+    seen_traces: set[str] = set()
+    for t in traces:
+        if t.trace_id in seen_traces:
+            continue
+        seen_traces.add(t.trace_id)
+        spans = db.get_trace_spans(t.trace_id)
+        messages = []
+        for s in spans:
+            if s.attributes.get("gen_ai.prompt.content"):
+                messages.append({
+                    "role": "user",
+                    "content": s.attributes["gen_ai.prompt.content"],
+                })
+            if s.tool_name:
+                messages.append({
+                    "role": "assistant",
+                    "content": "",
+                    "tool_calls": [{"name": s.tool_name}],
+                })
+                if s.attributes.get("gen_ai.tool.output"):
+                    messages.append({
+                        "role": "tool",
+                        "content": s.attributes["gen_ai.tool.output"],
+                    })
+            if s.attributes.get("gen_ai.completion.content"):
+                messages.append({
+                    "role": "assistant",
+                    "content": s.attributes["gen_ai.completion.content"],
+                })
+        results.append({
+            "trace_id": t.trace_id,
+            "agent_id": t.agent_id,
+            "messages": messages,
+        })
+    return json.dumps(results, default=str, indent=2)

tokenjam/cli/cmd_mcp.py

@@ -0,0 +1,78 @@
+"""tj mcp — start the stdio MCP server."""
+from __future__ import annotations
+
+import shutil
+import socket
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+import click
+import duckdb
+
+from tokenjam.core.config import find_config_file, load_config
+
+
+def _port_open(host: str, port: int) -> bool:
+    """Return True if something is listening on host:port."""
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.settimeout(0.5)
+        try:
+            s.connect((host, port))
+            return True
+        except (ConnectionRefusedError, OSError):
+            return False
+
+
+def _start_and_wait(host: str, port: int, timeout: float = 10.0) -> bool:
+    """Start tj serve in the background and wait up to *timeout* seconds for it
+    to accept connections. Returns True if the server is ready in time."""
+    ocw_bin = shutil.which("tj") or sys.argv[0]
+    popen_kwargs: dict = {
+        "stdout": subprocess.DEVNULL,
+        "stderr": subprocess.DEVNULL,
+    }
+    if sys.platform == "win32":
+        popen_kwargs["creationflags"] = subprocess.DETACHED_PROCESS
+    else:
+        popen_kwargs["start_new_session"] = True
+
+    try:
+        subprocess.Popen([ocw_bin, "serve"], **popen_kwargs)
+    except (FileNotFoundError, OSError):
+        return False
+
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        time.sleep(0.25)
+        if _port_open(host, port):
+            return True
+    return False
+
+
+@click.command("mcp")
+@click.pass_context
+def cmd_mcp(ctx: click.Context) -> None:
+    """Start the OCW MCP server (stdio transport for Claude Code)."""
+    from tokenjam.mcp.server import mcp, init
+
+    config_path = find_config_file()
+    if config_path is not None:
+        config = load_config(str(config_path))
+        host = config.api.host
+        port = config.api.port
+
+        if _port_open(host, port) or _start_and_wait(host, port):
+            # tj serve is running (already up or we just started it)
+            serve_url = f"http://{host}:{port}"
+            init(ro_conn=None, config=config, serve_url=serve_url)
+        else:
+            # Could not reach or start tj serve — fall back to read-only DuckDB
+            # so MCP read tools still work, though live ingest won't be available.
+            db_path = str(Path(config.storage.path).expanduser())
+            ro_conn = duckdb.connect(db_path, read_only=True)
+            init(ro_conn=ro_conn, config=config, serve_url=None)
+    # If no config: init is not called; tools return the no-config sentinel.
+
+    mcp.run()