opensmith 0.1.0__py3-none-any.whl

opensmith/__init__.py

@@ -0,0 +1,14 @@
+from opensmith.tracer import TraceCallable, trace
+from opensmith.patcher import autopatch, unpatch
+from opensmith.storage import Storage
+from opensmith.models import Trace, Step, Run
+
+__all__ = [
+    "trace",
+    "autopatch",
+    "unpatch",
+    "Storage",
+    "Trace",
+    "Step",
+    "Run",
+]

opensmith/cli.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+import time
+from typing import Any
+
+import click
+import uvicorn
+from rich.console import Console
+from rich.table import Table
+
+from opensmith.storage import Storage
+
+
+console = Console()
+
+
+@click.group(name="opensmith")
+def cli() -> None:
+    """Local-first LLM pipeline tracer."""
+
+
+@cli.command()
+@click.option("--port", default=7823, show_default=True, type=int)
+@click.option("--host", default="127.0.0.1", show_default=True)
+def ui(port: int, host: str) -> None:
+    """Start the local dashboard."""
+    click.echo(f"opensmith UI running at http://{host}:{port}")
+    click.echo("Press Ctrl+C to stop")
+    uvicorn.run("opensmith.server:app", host=host, port=port)
+
+
+@cli.command()
+def clear() -> None:
+    """Clear all traces."""
+    if not click.confirm("Clear all traces?", default=False):
+        return
+
+    Storage().delete_all()
+    click.echo("Cleared all traces.")
+
+
+@cli.command()
+def stats() -> None:
+    """Show trace statistics."""
+    data = Storage().get_stats()
+
+    table = Table(title="opensmith stats")
+    table.add_column("Metric")
+    table.add_column("Value", justify="right")
+
+    table.add_row("Total traces", str(data["total_traces"]))
+    table.add_row("Total steps", str(data["total_steps"]))
+    table.add_row("Total tokens", str(data["total_tokens"]))
+    table.add_row("Total cost (USD)", f"${float(data['total_cost_usd']):.6f}")
+
+    console.print(table)
+
+
+@cli.command()
+@click.option("--limit", default=20, show_default=True, type=int)
+def traces(limit: int) -> None:
+    """List recent traces."""
+    rows = Storage().get_traces(limit=limit)
+
+    table = Table(title="opensmith traces")
+    table.add_column("id")
+    table.add_column("name")
+    table.add_column("latency_ms", justify="right")
+    table.add_column("error")
+    table.add_column("created_at")
+
+    for row in rows:
+        table.add_row(
+            str(row.get("id", ""))[:8],
+            str(row.get("name") or ""),
+            _format_latency(row.get("latency_ms")),
+            "yes" if row.get("error") else "no",
+            _format_timestamp(row.get("created_at")),
+        )
+
+    console.print(table)
+
+
+def _format_latency(value: Any) -> str:
+    if value is None:
+        return ""
+    try:
+        return f"{float(value):.2f}"
+    except (TypeError, ValueError):
+        return ""
+
+
+def _format_timestamp(value: Any) -> str:
+    if value is None:
+        return ""
+    try:
+        return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(float(value)))
+    except (TypeError, ValueError, OSError):
+        return ""

opensmith/models.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+import time
+from typing import Any, Literal
+from uuid import uuid4
+
+from pydantic import BaseModel, Field
+
+
+StepType = Literal["llm", "retrieval", "tool", "custom"]
+
+
+class Step(BaseModel):
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    trace_id: str | None = None
+    name: str
+    input: Any | None = None
+    output: Any | None = None
+    error: str | None = None
+    start_time: float | None = None
+    end_time: float | None = None
+    latency_ms: float | None = None
+    tokens_input: int | None = None
+    tokens_output: int | None = None
+    tokens_total: int | None = None
+    model: str | None = None
+    cost_usd: float | None = None
+    step_type: StepType | None = None
+    metadata: dict[str, Any] | None = None
+
+
+class Trace(BaseModel):
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    name: str
+    input: Any | None = None
+    output: Any | None = None
+    error: str | None = None
+    start_time: float | None = None
+    end_time: float | None = None
+    latency_ms: float | None = None
+    parent_id: str | None = None
+    run_id: str | None = None
+    metadata: dict[str, Any] | None = None
+    steps: list[Step] = Field(default_factory=list)
+
+
+class Run(BaseModel):
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    name: str | None = None
+    tags: list[str] = Field(default_factory=list)
+    created_at: float = Field(default_factory=lambda: time.time())

opensmith/patcher.py

@@ -0,0 +1,263 @@
+from __future__ import annotations
+
+import functools
+import importlib
+import time
+from collections.abc import Callable
+from typing import Any
+
+from opensmith.models import Step
+from opensmith.storage import Storage
+from opensmith.tracer import _current_trace_id
+
+
+PATCH_TARGETS: dict[str, list[dict[str, Any]]] = {
+    "openai": [
+        {
+            "module": "openai",
+            "path": ["chat", "completions", "create"],
+            "name": "chat.completions.create",
+            "step_type": "llm",
+        },
+        {
+            "module": "openai",
+            "path": ["embeddings", "create"],
+            "name": "embeddings.create",
+            "step_type": "llm",
+        },
+    ],
+    "anthropic": [
+        {
+            "module": "anthropic",
+            "path": ["messages", "create"],
+            "name": "messages.create",
+            "step_type": "llm",
+        },
+    ],
+    "litellm": [
+        {
+            "module": "litellm",
+            "path": ["completion"],
+            "name": "completion",
+            "step_type": "llm",
+        },
+        {
+            "module": "litellm",
+            "path": ["embedding"],
+            "name": "embedding",
+            "step_type": "llm",
+        },
+    ],
+    "qdrant": [
+        {
+            "module": "qdrant_client",
+            "path": ["QdrantClient", "search"],
+            "name": "QdrantClient.search",
+            "step_type": "retrieval",
+        },
+        {
+            "module": "qdrant_client",
+            "path": ["QdrantClient", "upsert"],
+            "name": "QdrantClient.upsert",
+            "step_type": "retrieval",
+        },
+    ],
+    "chromadb": [
+        {
+            "module": "chromadb.api.models.Collection",
+            "path": ["Collection", "query"],
+            "name": "Collection.query",
+            "step_type": "retrieval",
+        },
+        {
+            "module": "chromadb.api.models.Collection",
+            "path": ["Collection", "add"],
+            "name": "Collection.add",
+            "step_type": "retrieval",
+        },
+    ],
+    "pinecone": [
+        {
+            "module": "pinecone",
+            "path": ["Index", "query"],
+            "name": "Index.query",
+            "step_type": "retrieval",
+        },
+        {
+            "module": "pinecone",
+            "path": ["Index", "upsert"],
+            "name": "Index.upsert",
+            "step_type": "retrieval",
+        },
+    ],
+}
+
+_patched: dict[tuple[str, tuple[str, ...]], tuple[Any, str, Callable[..., Any]]] = {}
+
+
+def autopatch(
+    only: list[str] | None = None,
+    exclude: list[str] | None = None,
+    storage: Storage | None = None,
+) -> None:
+    selected = set(only) if only is not None else set(PATCH_TARGETS)
+    excluded = set(exclude or [])
+    active_storage = storage or Storage()
+
+    for client_name in selected:
+        if client_name in excluded:
+            continue
+
+        for target in PATCH_TARGETS.get(client_name, []):
+            _patch_target(client_name, target, active_storage)
+
+
+def unpatch() -> None:
+    for owner, attr_name, original in _patched.values():
+        setattr(owner, attr_name, original)
+    _patched.clear()
+
+
+def _patch_target(
+    client_name: str,
+    target: dict[str, Any],
+    storage: Storage,
+) -> None:
+    module_name = target["module"]
+    path = tuple(target["path"])
+    key = (module_name, path)
+
+    if key in _patched:
+        return
+
+    try:
+        module = importlib.import_module(module_name)
+        owner, attr_name, original = _resolve_target(module, path)
+    except Exception:
+        return
+
+    if not callable(original):
+        return
+
+    wrapped = _wrap_method(
+        original=original,
+        storage=storage,
+        step_name=f"{client_name}.{target['name']}",
+        step_type=target["step_type"],
+    )
+
+    try:
+        setattr(owner, attr_name, wrapped)
+    except Exception:
+        return
+
+    _patched[key] = (owner, attr_name, original)
+
+
+def _resolve_target(module: Any, path: tuple[str, ...]) -> tuple[Any, str, Callable[..., Any]]:
+    owner = module
+    for part in path[:-1]:
+        owner = getattr(owner, part)
+
+    attr_name = path[-1]
+    original = getattr(owner, attr_name)
+    return owner, attr_name, original
+
+
+def _wrap_method(
+    original: Callable[..., Any],
+    storage: Storage,
+    step_name: str,
+    step_type: str,
+) -> Callable[..., Any]:
+    @functools.wraps(original)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        start_time = time.time()
+        step = Step(
+            trace_id=_current_trace_id(),
+            name=step_name,
+            input={
+                "args": [str(arg) for arg in args],
+                "kwargs": {key: str(value) for key, value in kwargs.items()},
+            },
+            start_time=start_time,
+            model=str(kwargs["model"]) if "model" in kwargs else None,
+            step_type=step_type,  # type: ignore[arg-type]
+        )
+
+        try:
+            result = original(*args, **kwargs)
+            step.output = _serialize_output(result)
+            _apply_usage(step, result)
+            return result
+        except Exception as exc:
+            step.error = repr(exc)
+            raise
+        finally:
+            step.end_time = time.time()
+            step.latency_ms = (step.end_time - start_time) * 1000
+            storage.save_step(step)
+
+    return wrapper
+
+
+def _serialize_output(value: Any) -> Any:
+    if isinstance(value, str | int | float | bool | list | tuple | dict) or value is None:
+        return value
+    if hasattr(value, "model_dump"):
+        try:
+            return value.model_dump()
+        except Exception:
+            pass
+    if hasattr(value, "dict"):
+        try:
+            return value.dict()
+        except Exception:
+            pass
+    return str(value)
+
+
+def _apply_usage(step: Step, result: Any) -> None:
+    usage = _get_value(result, "usage")
+    if usage is None:
+        return
+
+    tokens_input = (
+        _get_value(usage, "prompt_tokens")
+        or _get_value(usage, "input_tokens")
+        or _get_value(usage, "tokens_input")
+    )
+    tokens_output = (
+        _get_value(usage, "completion_tokens")
+        or _get_value(usage, "output_tokens")
+        or _get_value(usage, "tokens_output")
+    )
+    tokens_total = _get_value(usage, "total_tokens") or _get_value(
+        usage,
+        "tokens_total",
+    )
+
+    step.tokens_input = _to_int(tokens_input)
+    step.tokens_output = _to_int(tokens_output)
+    step.tokens_total = _to_int(tokens_total)
+
+    if step.tokens_total is None:
+        input_count = step.tokens_input or 0
+        output_count = step.tokens_output or 0
+        total = input_count + output_count
+        step.tokens_total = total if total else None
+
+
+def _get_value(value: Any, key: str) -> Any:
+    if isinstance(value, dict):
+        return value.get(key)
+    return getattr(value, key, None)
+
+
+def _to_int(value: Any) -> int | None:
+    if value is None:
+        return None
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        return None

opensmith/server.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import HTMLResponse
+from fastapi.staticfiles import StaticFiles
+
+from opensmith.storage import Storage
+
+
+UI_DIR = Path(__file__).parent / "ui"
+INDEX_HTML = UI_DIR / "index.html"
+
+storage = Storage()
+
+app = FastAPI(title="opensmith", docs_url=None, redoc_url=None)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=[
+        "http://localhost:7823",
+        "http://127.0.0.1:7823",
+    ],
+    allow_credentials=False,
+    allow_methods=["GET", "DELETE"],
+    allow_headers=["*"],
+)
+
+if UI_DIR.exists():
+    app.mount("/static", StaticFiles(directory=UI_DIR), name="static")
+
+
+@app.get("/", response_class=HTMLResponse)
+def index() -> HTMLResponse:
+    if INDEX_HTML.exists():
+        return HTMLResponse(INDEX_HTML.read_text(encoding="utf-8"))
+
+    return HTMLResponse(
+        """
+        <!doctype html>
+        <html>
+            <head>
+                <title>opensmith</title>
+            </head>
+            <body>
+                <h1>opensmith UI loading...</h1>
+            </body>
+        </html>
+        """
+    )
+
+
+@app.get("/api/traces")
+def get_traces(limit: int = 50) -> list[dict[str, Any]]:
+    return storage.get_traces(limit=limit)
+
+
+@app.get("/api/traces/{trace_id}")
+def get_trace(trace_id: str) -> dict[str, Any]:
+    trace, steps = storage.get_trace(trace_id)
+    return {
+        "trace": trace,
+        "steps": steps,
+    }
+
+
+@app.delete("/api/traces")
+def delete_traces() -> dict[str, str]:
+    storage.delete_all()
+    return {"status": "cleared"}
+
+
+@app.get("/api/stats")
+def get_stats() -> dict[str, Any]:
+    return storage.get_stats()