touchstone-bench 0.1.0__py3-none-any.whl

touchstone/__init__.py

@@ -0,0 +1,106 @@
+"""Touchstone public API — the one-liners an instrumented app imports.
+
+    import touchstone; touchstone.trace()
+
+`trace()` wires capture to the SQLite db and patches whichever of openai/anthropic is importable.
+It never raises if neither is installed.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+
+from . import capture
+from .capture import episode, tool
+from .capture.patch_anthropic import patch as _patch_anthropic
+from .capture.patch_openai import patch as _patch_openai
+from .config import load_settings
+from .messages import canonical
+
+__all__ = ["trace", "episode", "outcome", "tool", "record_llm_call"]
+
+_log = logging.getLogger("touchstone")
+_traced = False
+
+
+def trace(db: str | None = None, otel: bool = False) -> dict:
+    """Idempotent. Point capture at `db` (or configured/default) and patch installed SDKs.
+
+    With `otel=True`, also register the OpenInference OTel span exporter (needs the
+    `touchstone[otel]` extra) so existing OpenInference spans land in the store.
+    """
+    global _traced
+    settings = load_settings()
+    db_path = db or settings.db_path
+    capture.configure(db_path)
+    from .survey.netshim import install_from_env
+
+    install_from_env()  # no-op in prod; rewrites a hardcoded host if TOUCHSTONE_SIMULATORS is set
+    patched = []
+    if _patch_openai():
+        patched.append("openai")
+    if _patch_anthropic():
+        patched.append("anthropic")
+    if otel:
+        from .capture.openinference import register
+
+        if register():
+            patched.append("openinference")
+    _traced = True
+    return {"db": db_path, "patched": patched}
+
+
+def outcome(score: float | None, label: str | None) -> None:
+    try:  # resolving the current/untracked episode can touch the db; never raise into the app
+        capture.current_episode().outcome(score, label)
+    except Exception as exc:
+        _log.warning("touchstone outcome failed: %r", exc)
+
+
+def _normalize_calls(items) -> list[dict]:
+    calls = []
+    for tc in items:
+        args = tc.get("arguments")
+        if not isinstance(args, str):
+            args = json.dumps(args, ensure_ascii=False)
+        calls.append({"id": tc.get("id"), "name": tc.get("name"), "arguments": args})
+    return calls
+
+
+def _normalize_reply(reply) -> tuple[str, list[dict], dict | None]:
+    from .llm.base import Reply
+
+    if isinstance(reply, Reply):
+        reply = {"content": reply.content, "tool_calls": reply.tool_calls, "usage": reply.usage}
+    if isinstance(reply, str):
+        return reply, [], None
+    if isinstance(reply, dict):
+        calls = _normalize_calls(reply.get("tool_calls") or [])
+        return reply.get("content", "") or "", calls, reply.get("usage")
+    return str(reply), [], None
+
+
+def record_llm_call(model, messages, reply, tools=None, usage=None) -> None:
+    """Record an llm call made outside a patched SDK (e.g. via a Provider)."""
+    try:  # capture must never raise into the app (db unwritable, bad reply shape, …)
+        _record_llm_call(model, messages, reply, tools, usage)
+    except Exception as exc:
+        _log.warning("touchstone record_llm_call failed: %r", exc)
+
+
+def _record_llm_call(model, messages, reply, tools, usage) -> None:
+    content, tool_calls, reply_usage = _normalize_reply(reply)
+    u = usage or reply_usage or {}
+    convo = canonical(
+        list(messages) + [{"role": "assistant", "content": content, "tool_calls": tool_calls}]
+    )
+    capture.add_span(
+        "model",
+        model or "model",
+        model=model,
+        input={"messages": convo[:-1], "tools": tools or [], "params": {}},
+        output={"message": convo[-1]},
+        tokens_in=u.get("tokens_in"),
+        tokens_out=u.get("tokens_out"),
+    )

touchstone/capture/__init__.py

@@ -0,0 +1,21 @@
+from .context import (
+    EpisodeHandle,
+    add_span,
+    configure,
+    current_episode,
+    episode,
+    get_conn,
+    is_configured,
+    tool,
+)
+
+__all__ = [
+    "episode",
+    "tool",
+    "add_span",
+    "configure",
+    "current_episode",
+    "get_conn",
+    "is_configured",
+    "EpisodeHandle",
+]

touchstone/capture/context.py

@@ -0,0 +1,241 @@
+"""Episode context + span recording. The instrumented app, the SDK patches, the litellm callback,
+and the public helpers all funnel through `add_span` here.
+
+A connection is thread-local and keyed to the configured db path, so concurrent writers (server +
+app) each hold their own connection and rely on WAL + busy_timeout for safety.
+"""
+
+from __future__ import annotations
+
+import inspect
+import json
+import logging
+import threading
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from functools import wraps
+
+from .. import store
+
+_log = logging.getLogger("touchstone")
+_db_path: str | None = None
+_local = threading.local()
+_current: ContextVar[EpisodeHandle | None] = ContextVar("touchstone_episode", default=None)
+_untracked_lock = threading.Lock()
+_untracked_ids: dict[str, str] = {}
+
+
+def configure(db_path: str) -> None:
+    global _db_path
+    _db_path = db_path
+    _local.__dict__.pop("conn", None)
+
+
+def is_configured() -> bool:
+    return _db_path is not None
+
+
+def _safe(action: str, fn, default=None):
+    """Run `fn`; if recording fails (e.g. a non-writable db), log one warning and return `default`.
+    Capture must never raise into the user's application — the explicit helpers (`episode`, `tool`,
+    `outcome`) funnel through here just as the SDK patches funnel through `spans._recorder`."""
+    try:
+        return fn()
+    except Exception as exc:
+        _log.warning("touchstone %s failed: %r", action, exc)
+        return default
+
+
+def get_conn():
+    if _db_path is None:
+        raise RuntimeError("touchstone.trace() has not been called")
+    conn = getattr(_local, "conn", None)
+    if conn is None or getattr(_local, "path", None) != _db_path:
+        conn = store.connect(_db_path)
+        _local.conn = conn
+        _local.path = _db_path
+    return conn
+
+
+@dataclass
+class EpisodeHandle:
+    id: str
+
+    def outcome(self, score: float | None, label: str | None) -> None:
+        if not self.id:  # detached handle from a failed episode start; nothing to record
+            return
+        _safe("outcome", lambda: store.outcome(get_conn(), self.id, score, label))
+
+
+def _start_episode(name: str, meta: dict | None, source: str) -> EpisodeHandle:
+    conn = get_conn()
+    ep = store.insert_episode(conn, store.Episode(name=name, source=source, meta=meta or {}))
+    return EpisodeHandle(ep.id)
+
+
+@contextmanager
+def episode(name: str, meta: dict | None = None, source: str = "app"):
+    handle = _safe("episode start", lambda: _start_episode(name, meta, source))
+    if handle is None:  # capture unavailable; the app must still run its `with` body
+        yield EpisodeHandle("")
+        return
+    token = _current.set(handle)
+    try:
+        yield handle
+    finally:
+        _safe("episode end", lambda: store.end_episode(get_conn(), handle.id))
+        _current.reset(token)
+
+
+def _find_untracked(conn, name: str) -> str | None:
+    for e in store.list_episodes(conn):
+        if e.name == name and e.source == "untracked":
+            return e.id
+    return None
+
+
+def _untracked() -> EpisodeHandle:
+    conn = get_conn()
+    date = datetime.now(UTC).date().isoformat()
+    name = f"untracked-{date}"
+    with _untracked_lock:
+        eid = _untracked_ids.get(date)
+        if eid and store.get_episode(conn, eid):
+            return EpisodeHandle(eid)
+        found = _find_untracked(conn, name)
+        if found is None:
+            found = store.insert_episode(conn, store.Episode(name=name, source="untracked")).id
+        _untracked_ids[date] = found
+        return EpisodeHandle(found)
+
+
+def current_episode() -> EpisodeHandle:
+    return _current.get() or _untracked()
+
+
+def add_span(
+    kind: str,
+    name: str,
+    *,
+    model: str | None = None,
+    input: dict | None = None,
+    output: dict | None = None,
+    tokens_in: int | None = None,
+    tokens_out: int | None = None,
+    cost_usd: float | None = None,
+    error: str | None = None,
+    started_at: str | None = None,
+    parent_id: str | None = None,
+    tool_call_id: str | None = None,
+) -> store.Span:
+    conn = get_conn()
+    ep = current_episode()
+    span = store.Span(
+        episode_id=ep.id,
+        kind=kind,
+        name=name,
+        model=model,
+        parent_id=parent_id,
+        input=input or {},
+        output=output or {},
+        tokens_in=tokens_in,
+        tokens_out=tokens_out,
+        cost_usd=cost_usd,
+        error=error,
+        started_at=started_at or store.now(),
+        tool_call_id=tool_call_id,
+    )
+    span.ended_at = store.now()
+    return store.insert_span(conn, span)
+
+
+def jsonable(value):
+    try:
+        json.dumps(value, ensure_ascii=False)
+        return value
+    except (TypeError, ValueError):
+        return repr(value)
+
+
+def _bound_args(fn, args, kwargs) -> dict:
+    try:
+        bound = inspect.signature(fn).bind_partial(*args, **kwargs)
+        return {k: jsonable(v) for k, v in bound.arguments.items()}
+    except TypeError:
+        return {
+            "args": [jsonable(a) for a in args],
+            "kwargs": {k: jsonable(v) for k, v in kwargs.items()},
+        }
+
+
+def _resolve_tool_call_id(tname: str) -> str | None:
+    """The most recent model span's unresolved tool_call of this name, if any — so a dispatched
+    tool span links back to the call that requested it without the caller passing the id."""
+    try:
+        conn = get_conn()
+        spans = store.list_spans(conn, current_episode().id)
+    except Exception:
+        return None
+    model_spans = [s for s in spans if s.kind == "model"]
+    if not model_spans:
+        return None
+    used = {s.tool_call_id for s in spans if s.kind == "tool" and s.tool_call_id}
+    calls = (model_spans[-1].output or {}).get("message", {}).get("tool_calls") or []
+    for tc in calls:
+        if tc.get("name") == tname and tc.get("id") and tc.get("id") not in used:
+            return tc["id"]
+    return None
+
+
+def _record_tool(tname, bound, tool_call_id, started, *, result=None, error=None) -> None:
+    def _do():
+        call_id = tool_call_id if tool_call_id is not None else _resolve_tool_call_id(tname)
+        output = None if error else {"result": jsonable(result)}
+        add_span("tool", tname, input={"name": tname, "arguments": bound}, output=output,
+                 error=error, started_at=started, tool_call_id=call_id)
+
+    _safe("tool span", _do)  # a recording failure must never lose the tool's own result
+
+
+def tool(fn=None, *, name: str | None = None):
+    """Decorator recording a tool span (name/arguments/result/error/tool_call_id). Works sync and
+    async. A caller may pass `tool_call_id=` to link the span to a specific model tool_call;
+    otherwise it is auto-linked to the latest unresolved call of the same name."""
+
+    def decorate(f):
+        tname = name or f.__name__
+
+        if inspect.iscoroutinefunction(f):
+            @wraps(f)
+            async def awrapper(*args, **kwargs):
+                tcid = kwargs.pop("tool_call_id", None)
+                started = store.now()
+                bound = _bound_args(f, args, kwargs)
+                try:
+                    result = await f(*args, **kwargs)
+                except Exception as exc:
+                    _record_tool(tname, bound, tcid, started, error=repr(exc))
+                    raise
+                _record_tool(tname, bound, tcid, started, result=result)
+                return result
+
+            return awrapper
+
+        @wraps(f)
+        def wrapper(*args, **kwargs):
+            tcid = kwargs.pop("tool_call_id", None)
+            started = store.now()
+            bound = _bound_args(f, args, kwargs)
+            try:
+                result = f(*args, **kwargs)
+            except Exception as exc:
+                _record_tool(tname, bound, tcid, started, error=repr(exc))
+                raise
+            _record_tool(tname, bound, tcid, started, result=result)
+            return result
+
+        return wrapper
+
+    return decorate(fn) if fn else decorate

touchstone/capture/litellm.py

@@ -0,0 +1,95 @@
+"""litellm CustomLogger callback that records an llm span. litellm is imported lazily and never
+required; the logger works even when litellm is absent (base class falls back to object)."""
+
+from __future__ import annotations
+
+import functools
+from datetime import datetime
+
+from . import spans
+from .patch_openai import _extract  # litellm ModelResponse is OpenAI-shaped
+
+
+def _base():
+    try:
+        from litellm.integrations.custom_logger import CustomLogger
+
+        return CustomLogger
+    except Exception:
+        return object
+
+
+def _iso(t) -> str:
+    if isinstance(t, datetime):
+        return t.isoformat()
+    from .. import store
+
+    return store.now()
+
+
+class TouchstoneLogger(_base()):
+    def _record(self, kwargs, response_obj, start_time, error):
+        model = kwargs.get("model")
+        messages = kwargs.get("messages") or []
+        opt = kwargs.get("optional_params") or {}
+        tools = opt.get("tools") or kwargs.get("tools")
+        result = _extract(response_obj) if response_obj is not None else spans._ERROR
+        err = error or (repr(kwargs["exception"]) if kwargs.get("exception") else None)
+        try:
+            spans.record("litellm", model, messages, tools, {}, result, err, _iso(start_time))
+        except Exception as exc:  # capture must never break the caller
+            spans._log.warning("touchstone capture failed: %r", exc)
+
+    def log_success_event(self, kwargs, response_obj, start_time, end_time):
+        self._record(kwargs, response_obj, start_time, None)
+
+    def log_failure_event(self, kwargs, response_obj, start_time, end_time):
+        self._record(kwargs, response_obj, start_time, "failure")
+
+    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        self._record(kwargs, response_obj, start_time, None)
+
+    async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
+        self._record(kwargs, response_obj, start_time, "failure")
+
+
+def _override(orig):
+    """Wrap a litellm entrypoint so TOUCHSTONE_MODEL rewrites `model` before the call runs."""
+    @functools.wraps(orig)
+    def wrapper(*args, **kwargs):
+        spans.override_model(kwargs)
+        return orig(*args, **kwargs)
+    wrapper._touchstone = True
+    return wrapper
+
+
+def _override_async(orig):
+    @functools.wraps(orig)
+    async def wrapper(*args, **kwargs):
+        spans.override_model(kwargs)
+        return await orig(*args, **kwargs)
+    wrapper._touchstone = True
+    return wrapper
+
+
+def _patch_model_override(litellm) -> None:
+    """Rewrite the model on `litellm.completion`/`acompletion` (attribute callers only; `from
+    litellm import completion` binds before this runs, so those callers keep the recorded model)."""
+    fn = getattr(litellm, "completion", None)
+    if fn is not None and not getattr(fn, "_touchstone", False):
+        litellm.completion = _override(fn)
+    afn = getattr(litellm, "acompletion", None)
+    if afn is not None and not getattr(afn, "_touchstone", False):
+        litellm.acompletion = _override_async(afn)
+
+
+def install() -> bool:
+    """Register the logger with litellm and rewire the model setting. False if litellm is absent."""
+    try:
+        import litellm
+    except Exception:
+        return False
+    logger = TouchstoneLogger()
+    litellm.callbacks = list(getattr(litellm, "callbacks", []) or []) + [logger]
+    _patch_model_override(litellm)
+    return True

touchstone/capture/openinference.py

@@ -0,0 +1,169 @@
+"""Optional ingest of OpenInference-attributed OpenTelemetry spans into the same tables.
+
+A team already instrumented with OpenInference (Arize) auto-instrumentors can point Touchstone at
+their traces with `touchstone.trace(otel=True)`: this registers an in-process OTel `SpanExporter`
+that normalizes each span's flattened `openinference.*` attributes into canonical messages and
+writes model / tool spans under one episode per trace. No collector process is involved. Requires
+the `touchstone[otel]` extra (opentelemetry-sdk); the module imports without it for testing against
+hand-built ReadableSpan-like objects.
+"""
+
+from __future__ import annotations
+
+from datetime import UTC, datetime
+
+from .. import store
+from ..messages import _args_str, canonical
+from . import context
+
+
+def _base_exporter():
+    try:
+        from opentelemetry.sdk.trace.export import SpanExporter
+
+        return SpanExporter
+    except Exception:
+        return object
+
+
+def _iso(ns) -> str:
+    return datetime.fromtimestamp(ns / 1e9, UTC).isoformat() if ns else store.now()
+
+
+def _indexed(attrs: dict, prefix: str) -> dict[int, dict]:
+    """Group `prefix.{i}.<tail>` attributes by index i."""
+    out: dict[int, dict] = {}
+    for key, val in attrs.items():
+        if not key.startswith(prefix + "."):
+            continue
+        i, _, tail = key[len(prefix) + 1:].partition(".")
+        if i.isdigit():
+            out.setdefault(int(i), {})[tail] = val
+    return out
+
+
+def _tool_calls(fields: dict) -> list[dict]:
+    by_j = _indexed(fields, "message.tool_calls")
+    calls = []
+    for j in sorted(by_j):
+        f = by_j[j]
+        calls.append({"id": f.get("tool_call.id"),
+                      "name": f.get("tool_call.function.name"),
+                      "arguments": _args_str(f.get("tool_call.function.arguments"))})
+    return calls
+
+
+def _message(fields: dict) -> dict:
+    msg = {"role": fields.get("message.role", "user"),
+           "content": fields.get("message.content", "") or ""}
+    calls = _tool_calls(fields)
+    if calls:
+        msg["tool_calls"] = calls
+    if fields.get("message.tool_call_id"):
+        msg["tool_call_id"] = fields["message.tool_call_id"]
+    return msg
+
+
+def _messages(attrs: dict, prefix: str) -> list[dict]:
+    by_i = _indexed(attrs, prefix)
+    return canonical([_message(by_i[i]) for i in sorted(by_i)])
+
+
+class OpenInferenceSpanExporter(_base_exporter()):
+    """Writes OpenInference OTel spans into Touchstone's SQLite, one episode per trace."""
+
+    def __init__(self):
+        self._episodes: dict[int, str] = {}
+        self._span_ids: dict[int, str] = {}
+
+    def export(self, spans):
+        for span in spans:
+            try:
+                self._ingest(span)
+            except Exception:  # ingest must never break the app's tracer pipeline
+                pass
+        return _success()
+
+    def shutdown(self):
+        pass
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        return True
+
+    def _episode_for(self, conn, span) -> str:
+        trace_id = span.context.trace_id
+        ep_id = self._episodes.get(trace_id)
+        if ep_id is None:
+            name = span.name if span.parent is None else "openinference"
+            ep_id = store.insert_episode(
+                conn, store.Episode(name=name, source="openinference")).id
+            self._episodes[trace_id] = ep_id
+        return ep_id
+
+    def _parent_id(self, span) -> str | None:
+        return self._span_ids.get(span.parent.span_id) if span.parent else None
+
+    def _ingest(self, span) -> None:
+        attrs = dict(span.attributes or {})
+        kind = attrs.get("openinference.span.kind")
+        conn = context.get_conn()
+        ep_id = self._episode_for(conn, span)
+        if kind == "LLM":
+            record = self._llm(ep_id, span, attrs)
+        elif kind == "TOOL":
+            record = self._tool(ep_id, span, attrs)
+        else:
+            return  # AGENT / CHAIN / RETRIEVER only seed the episode
+        self._span_ids[span.context.span_id] = store.insert_span(conn, record).id
+
+    def _llm(self, ep_id, span, attrs) -> store.Span:
+        outputs = _messages(attrs, "llm.output_messages")
+        reply = outputs[-1] if outputs else {"role": "assistant", "content": ""}
+        usage = {}
+        if attrs.get("llm.token_count.prompt_details.cache_read") is not None:
+            usage["cached_tokens"] = attrs["llm.token_count.prompt_details.cache_read"]
+        output = {"message": reply}
+        if usage:
+            output["usage"] = usage
+        return store.Span(
+            episode_id=ep_id, kind="model", name=span.name, parent_id=self._parent_id(span),
+            model=attrs.get("llm.model_name"),
+            input={"messages": _messages(attrs, "llm.input_messages"), "tools": [], "params": {}},
+            output=output,
+            tokens_in=attrs.get("llm.token_count.prompt"),
+            tokens_out=attrs.get("llm.token_count.completion"),
+            started_at=_iso(span.start_time), ended_at=_iso(span.end_time))
+
+    def _tool(self, ep_id, span, attrs) -> store.Span:
+        name = attrs.get("tool.name") or span.name
+        return store.Span(
+            episode_id=ep_id, kind="tool", name=name, parent_id=self._parent_id(span),
+            input={"name": name, "arguments": attrs.get("tool.parameters")},
+            output={"result": attrs.get("output.value")},
+            tool_call_id=attrs.get("tool_call.id"),
+            started_at=_iso(span.start_time), ended_at=_iso(span.end_time))
+
+
+def _success():
+    try:
+        from opentelemetry.sdk.trace.export import SpanExportResult
+
+        return SpanExportResult.SUCCESS
+    except Exception:
+        return None
+
+
+def register() -> bool:
+    """Register the exporter on the global TracerProvider. False if opentelemetry-sdk is absent."""
+    try:
+        from opentelemetry import trace as ot
+        from opentelemetry.sdk.trace import TracerProvider
+        from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+    except Exception:
+        return False
+    provider = ot.get_tracer_provider()
+    if not isinstance(provider, TracerProvider):
+        provider = TracerProvider()
+        ot.set_tracer_provider(provider)
+    provider.add_span_processor(SimpleSpanProcessor(OpenInferenceSpanExporter()))
+    return True