tracecast 0.1.0__py3-none-any.whl

tracecast/__init__.py

@@ -0,0 +1,6 @@
+from .core.tracer import Tracer
+from .core.cost_calculator import calculate_cost
+from .models.trace import Trace
+from .models.span import Span, SpanType
+
+__all__ = ["Tracer", "calculate_cost", "Trace", "Span", "SpanType"]

tracecast/core/cost_calculator.py

@@ -0,0 +1,55 @@
+from typing import Optional, Dict
+
+PRICE_TABLE: Dict[str, Dict[str, float]] = {
+    "gpt-5":             {"input": 0.00125,  "output": 0.01000},
+    "gpt-5.4":           {"input": 0.00250,  "output": 0.01500},
+    "gpt-4o":            {"input": 0.00250,  "output": 0.01000},
+    "gpt-4o-mini":       {"input": 0.00015,  "output": 0.00060},
+    "gpt-4-turbo":       {"input": 0.01000,  "output": 0.03000},
+    "gpt-4.1":           {"input": 0.00200,  "output": 0.00800},
+    "gpt-4.1-mini":      {"input": 0.00040,  "output": 0.00160},
+    "o3":                {"input": 0.00200,  "output": 0.00800},
+    "o4-mini":           {"input": 0.00110,  "output": 0.00440},
+
+    "claude-opus-4-6":   {"input": 0.00500,  "output": 0.02500},
+    "claude-sonnet-4-6": {"input": 0.00300,  "output": 0.01500},
+    "claude-haiku-4-5":  {"input": 0.00100,  "output": 0.00500},
+    "claude-opus-4-5":   {"input": 0.00500,  "output": 0.02500},
+    "claude-sonnet-4-5": {"input": 0.00300,  "output": 0.01500},
+    "claude-opus-4":     {"input": 0.01500,  "output": 0.07500},
+    "claude-sonnet-4":   {"input": 0.00300,  "output": 0.01500},
+    "claude-haiku-3-5":  {"input": 0.00080,  "output": 0.00400},
+
+    "gemini-3.1-pro":    {"input": 0.00200,  "output": 0.01200},
+    "gemini-3-flash":    {"input": 0.00050,  "output": 0.00300},
+    "gemini-2.5-flash":  {"input": 0.00030,  "output": 0.00250},
+    "gemini-2.5-pro":    {"input": 0.00125,  "output": 0.01000},
+    "gemini-2.0-flash":  {"input": 0.00010,  "output": 0.00040},
+
+    "llama-3.3-70b-versatile":                       {"input": 0.00059, "output": 0.00079},
+    "meta-llama/llama-4-scout-17b-16e-instruct":     {"input": 0.00011, "output": 0.00034},
+    "meta-llama/llama-4-maverick-17b-128e-instruct": {"input": 0.00020, "output": 0.00060},
+    "llama-3.3-70b":     {"input": 0.00059, "output": 0.00079},
+    "llama-4-scout":     {"input": 0.00011, "output": 0.00034},
+    "llama-4-maverick":  {"input": 0.00020, "output": 0.00060},
+
+    "ollama/*":          {"input": 0.0,     "output": 0.0},
+}
+
+
+def calculate_cost(
+    model: str,
+    tokens_in: int,
+    tokens_out: int,
+    custom_prices: Optional[Dict] = None,
+) -> float:
+    prices = {**PRICE_TABLE, **custom_prices} if custom_prices else PRICE_TABLE
+    table = prices.get(model) or _prefix_match(model, prices)
+    if not table:
+        return 0.0
+    return (tokens_in / 1000 * table["input"]) + (tokens_out / 1000 * table["output"])
+
+
+def _prefix_match(model: str, prices: Dict) -> Optional[Dict]:
+    provider = model.split("/")[0] + "/*"
+    return prices.get(provider)

tracecast/core/logger.py

@@ -0,0 +1,103 @@
+import logging
+import re
+from typing import Optional
+
+_logger = logging.getLogger("tracecast")
+
+
+def _inline(text: str, max_chars: int = 50) -> str:
+    normalized = re.sub(r"[\r\n\t]+", " ", text)
+    normalized = re.sub(r" {2,}", " ", normalized).strip()
+    if len(normalized) > max_chars:
+        normalized = normalized[:max_chars] + "..."
+    return normalized
+
+
+class TraceCastLogger:
+    def __init__(self, prefix: Optional[str] = None):
+        self._prefix = prefix
+
+    def _fmt(self, prefix: str, msg: str) -> str:
+        return f"[{prefix}] {msg}"
+
+    def trace_start(self, trace_name: str) -> None:
+        prefix = self._prefix or trace_name
+        _logger.info(self._fmt(prefix, "Trace started"))
+
+    def trace_end(
+        self,
+        trace_name: str,
+        *,
+        total_tokens: int,
+        cost_usd: float,
+        latency_ms: Optional[int],
+        tools_used: dict,
+    ) -> None:
+        prefix = self._prefix or trace_name
+        latency_str = f"{latency_ms / 1000:.2f}s" if latency_ms is not None else "n/a"
+        tools_str = ""
+        if tools_used:
+            parts = [f"{k}×{v}" for k, v in tools_used.items()]
+            tools_str = " | tools: " + ", ".join(parts)
+        _logger.info(
+            self._fmt(
+                prefix,
+                f"Trace finished → total: {total_tokens} tokens"
+                f" | ${cost_usd:.4f}"
+                f" | {latency_str}"
+                f"{tools_str}",
+            )
+        )
+
+    def llm_start(self, trace_name: str, *, model: str) -> None:
+        prefix = self._prefix or trace_name
+        _logger.info(self._fmt(prefix, f"LLM started → {model}"))
+
+    def llm_end(
+        self,
+        trace_name: str,
+        *,
+        model: str,
+        tokens_in: int,
+        tokens_out: int,
+        cost_usd: float,
+        latency_ms: Optional[float],
+    ) -> None:
+        prefix = self._prefix or trace_name
+        latency_str = f"{latency_ms / 1000:.2f}s" if latency_ms is not None else "n/a"
+        _logger.info(
+            self._fmt(
+                prefix,
+                f"LLM end → {model}"
+                f" | tokens: {tokens_in} in / {tokens_out} out"
+                f" | ${cost_usd:.4f}"
+                f" | {latency_str}",
+            )
+        )
+
+    def llm_error(self, trace_name: str, *, model: str, error: str) -> None:
+        prefix = self._prefix or trace_name
+        _logger.warning(self._fmt(prefix, f"LLM error → {model} | ⚠ {_inline(error)}"))
+
+    def tool_start(self, trace_name: str, *, name: str, input_str: str) -> None:
+        prefix = self._prefix or trace_name
+        _logger.info(self._fmt(prefix, f"Tool call → {name} | {_inline(input_str, 80)}"))
+
+    def tool_end(
+        self, trace_name: str, *, name: str, latency_ms: Optional[float]
+    ) -> None:
+        prefix = self._prefix or trace_name
+        latency_str = f"{latency_ms / 1000:.2f}s" if latency_ms is not None else "n/a"
+        _logger.info(self._fmt(prefix, f"Tool end → {name} | {latency_str}"))
+
+    def tool_error(self, trace_name: str, *, name: str, error: str) -> None:
+        prefix = self._prefix or trace_name
+        _logger.warning(self._fmt(prefix, f"Tool error → {name} | ⚠ {_inline(error)}"))
+
+    def chain_start(self, trace_name: str, *, name: str) -> None:
+        prefix = self._prefix or trace_name
+        _logger.info(self._fmt(prefix, f"Chain → {name}"))
+
+    def chain_error(self, trace_name: str, *, name: str, error: str) -> None:
+        prefix = self._prefix or trace_name
+        _logger.warning(self._fmt(prefix, f"Chain error → {name} | ⚠ {_inline(error)}"))

tracecast/core/token_counter.py

@@ -0,0 +1,39 @@
+from typing import Any
+
+
+def extract_tokens(response: Any, provider: str) -> dict:
+    extractors = {
+        "openai":    _from_openai,
+        "anthropic": _from_anthropic,
+        "langchain": _from_langchain_response,
+    }
+    fn = extractors.get(provider, _fallback)
+    return fn(response)
+
+
+def _from_openai(r) -> dict:
+    usage = getattr(r, "usage", None) or {}
+    return {
+        "input":  getattr(usage, "prompt_tokens", 0),
+        "output": getattr(usage, "completion_tokens", 0),
+    }
+
+
+def _from_anthropic(r) -> dict:
+    usage = getattr(r, "usage", None) or {}
+    return {
+        "input":  getattr(usage, "input_tokens", 0),
+        "output": getattr(usage, "output_tokens", 0),
+    }
+
+
+def _from_langchain_response(r) -> dict:
+    usage = r.get("token_usage") or r.get("usage", {})
+    return {
+        "input":  usage.get("prompt_tokens") or usage.get("input_tokens", 0),
+        "output": usage.get("completion_tokens") or usage.get("output_tokens", 0),
+    }
+
+
+def _fallback(r) -> dict:
+    return {"input": 0, "output": 0}

tracecast/core/tracer.py

@@ -0,0 +1,106 @@
+import uuid
+from contextvars import ContextVar
+from contextlib import contextmanager, asynccontextmanager
+from datetime import datetime, timezone
+from typing import Optional, List
+
+from ..models.trace import Trace
+from ..exporters.base import BaseExporter
+
+
+_current_trace: ContextVar[Optional[Trace]] = ContextVar("_current_trace", default=None)
+
+
+class Tracer:
+
+    def __init__(
+        self,
+        exporters: Optional[List[BaseExporter]] = None,
+        logging: bool = False,
+        log_prefix: Optional[str] = None,
+    ):
+        self.exporters = exporters or []
+        self._tc_logger = None
+        if logging:
+            from .logger import TraceCastLogger
+            self._tc_logger = TraceCastLogger(prefix=log_prefix)
+
+    @contextmanager
+    def trace(self, name: str, session_id=None, user_id=None, project_id=None, metadata=None):
+        t = Trace(
+            trace_id=str(uuid.uuid4()),
+            name=name,
+            session_id=session_id,
+            user_id=user_id,
+            project_id=project_id,
+            started_at=datetime.now(timezone.utc),
+            metadata=metadata or {},
+        )
+        if self._tc_logger:
+            self._tc_logger.trace_start(name)
+        token = _current_trace.set(t)
+        try:
+            yield t
+        finally:
+            t.finished_at = datetime.now(timezone.utc)
+            t._finalize()
+            _current_trace.reset(token)
+            if self._tc_logger:
+                self._tc_logger.trace_end(
+                    name,
+                    total_tokens=t.total_tokens,
+                    cost_usd=t.cost_usd,
+                    latency_ms=t.latency_ms,
+                    tools_used=t.tools_used,
+                )
+            self._export(t)
+
+    @asynccontextmanager
+    async def atrace(self, name: str, session_id=None, user_id=None, project_id=None, metadata=None):
+        t = Trace(
+            trace_id=str(uuid.uuid4()),
+            name=name,
+            session_id=session_id,
+            user_id=user_id,
+            project_id=project_id,
+            started_at=datetime.now(timezone.utc),
+            metadata=metadata or {},
+        )
+        if self._tc_logger:
+            self._tc_logger.trace_start(name)
+        token = _current_trace.set(t)
+        try:
+            yield t
+        finally:
+            t.finished_at = datetime.now(timezone.utc)
+            t._finalize()
+            _current_trace.reset(token)
+            if self._tc_logger:
+                self._tc_logger.trace_end(
+                    name,
+                    total_tokens=t.total_tokens,
+                    cost_usd=t.cost_usd,
+                    latency_ms=t.latency_ms,
+                    tools_used=t.tools_used,
+                )
+            await self._aexport(t)
+
+    def _export(self, trace: Trace) -> None:
+        for exporter in self.exporters:
+            try:
+                exporter.export(trace)
+            except Exception as exc:
+                import warnings
+                warnings.warn(f"TraceCast: exporter {type(exporter).__name__} failed: {exc}", stacklevel=2)
+
+    async def _aexport(self, trace: Trace) -> None:
+        for exporter in self.exporters:
+            try:
+                await exporter.aexport(trace)
+            except Exception as exc:
+                import warnings
+                warnings.warn(f"TraceCast: exporter {type(exporter).__name__} failed: {exc}", stacklevel=2)
+
+    @staticmethod
+    def current() -> Optional[Trace]:
+        return _current_trace.get()

tracecast/exporters/__init__.py

@@ -0,0 +1,15 @@
+from .base import BaseExporter
+from .json_file import JsonFileExporter
+from .dict_exporter import DictExporter
+
+try:
+    from .mongo import MongoExporter
+except ImportError:
+    MongoExporter = None
+
+try:
+    from .postgres import PostgresExporter
+except ImportError:
+    PostgresExporter = None
+
+__all__ = ["BaseExporter", "JsonFileExporter", "DictExporter", "MongoExporter", "PostgresExporter"]

tracecast/exporters/base.py

@@ -0,0 +1,15 @@
+from abc import ABC, abstractmethod
+from ..models.trace import Trace
+
+
+class BaseExporter(ABC):
+    @abstractmethod
+    def export(self, trace: Trace) -> None:
+        ...
+
+    def export_batch(self, traces: list) -> None:
+        for trace in traces:
+            self.export(trace)
+
+    async def aexport(self, trace: Trace) -> None:
+        self.export(trace)

tracecast/exporters/dict_exporter.py

@@ -0,0 +1,35 @@
+from typing import Any, Callable, Dict, Iterable, List, Optional, Set
+from .base import BaseExporter
+from ..models.trace import Trace
+
+
+def _filter_dict(doc: dict, include: Optional[Set[str]], exclude: Optional[Set[str]]) -> dict:
+    if include is not None:
+        return {k: v for k, v in doc.items() if k in include}
+    if exclude is not None:
+        return {k: v for k, v in doc.items() if k not in exclude}
+    return doc
+
+
+class DictExporter(BaseExporter):
+
+    def __init__(
+        self,
+        on_trace: Optional[Callable[[Dict[str, Any]], None]] = None,
+        include_fields: Optional[Iterable[str]] = None,
+        exclude_fields: Optional[Iterable[str]] = None,
+    ):
+        self._on_trace = on_trace
+        self._include: Optional[Set[str]] = set(include_fields) if include_fields is not None else None
+        self._exclude: Optional[Set[str]] = set(exclude_fields) if exclude_fields is not None else None
+        self.traces: List[Dict[str, Any]] = []
+
+    def export(self, trace: Trace) -> None:
+        doc = _filter_dict(trace.to_dict(), self._include, self._exclude)
+        if self._on_trace is not None:
+            self._on_trace(doc)
+        else:
+            self.traces.append(doc)
+
+    def clear(self) -> None:
+        self.traces.clear()

tracecast/exporters/json_file.py

@@ -0,0 +1,37 @@
+import json
+import asyncio
+from pathlib import Path
+from typing import Iterable, Optional, Set
+from .base import BaseExporter
+from ..models.trace import Trace
+
+
+def _filter_dict(doc: dict, include: Optional[Set[str]], exclude: Optional[Set[str]]) -> dict:
+    if include is not None:
+        return {k: v for k, v in doc.items() if k in include}
+    if exclude is not None:
+        return {k: v for k, v in doc.items() if k not in exclude}
+    return doc
+
+
+class JsonFileExporter(BaseExporter):
+
+    def __init__(
+        self,
+        path: str = "./traces.jsonl",
+        include_fields: Optional[Iterable[str]] = None,
+        exclude_fields: Optional[Iterable[str]] = None,
+    ):
+        self.path = Path(path)
+        if self.path.parent != Path("."):
+            self.path.parent.mkdir(parents=True, exist_ok=True)
+        self._include: Optional[Set[str]] = set(include_fields) if include_fields is not None else None
+        self._exclude: Optional[Set[str]] = set(exclude_fields) if exclude_fields is not None else None
+
+    def export(self, trace: Trace) -> None:
+        doc = _filter_dict(trace.to_dict(), self._include, self._exclude)
+        with self.path.open("a", encoding="utf-8") as f:
+            f.write(json.dumps(doc, default=str) + "\n")
+
+    async def aexport(self, trace: Trace) -> None:
+        await asyncio.to_thread(self.export, trace)

tracecast/exporters/mongo.py

@@ -0,0 +1,34 @@
+from datetime import datetime, timezone
+from typing import Iterable, Optional, Set
+from pymongo import MongoClient
+from .base import BaseExporter
+from ..models.trace import Trace
+
+
+def _filter_doc(doc: dict, include: Optional[Set[str]], exclude: Optional[Set[str]]) -> dict:
+    if include is not None:
+        return {k: v for k, v in doc.items() if k in include}
+    if exclude is not None:
+        return {k: v for k, v in doc.items() if k not in exclude}
+    return doc
+
+
+class MongoExporter(BaseExporter):
+
+    def __init__(
+        self,
+        uri: str,
+        db: str = "tracecast",
+        collection: str = "traces",
+        include_fields: Optional[Iterable[str]] = None,
+        exclude_fields: Optional[Iterable[str]] = None,
+    ):
+        self.col = MongoClient(uri)[db][collection]
+        self._include: Optional[Set[str]] = set(include_fields) if include_fields is not None else None
+        self._exclude: Optional[Set[str]] = set(exclude_fields) if exclude_fields is not None else None
+
+    def export(self, trace: Trace) -> None:
+        doc = trace.to_dict()
+        doc["exported_at"] = datetime.now(timezone.utc)
+        doc = _filter_doc(doc, self._include, self._exclude)
+        self.col.insert_one(doc)

tracecast/exporters/postgres.py

@@ -0,0 +1,164 @@
+import json
+from datetime import datetime, timezone
+from typing import Any, Dict, Iterable, List, Optional, Set
+from .base import BaseExporter
+from ..models.trace import Trace
+
+_ALL_COLUMNS: List[str] = [
+    "trace_id",
+    "name",
+    "session_id",
+    "user_id",
+    "project_id",
+    "model",
+    "total_tokens_in",
+    "total_tokens_out",
+    "total_tokens",
+    "cost_usd",
+    "latency_ms",
+    "tools_used",
+    "spans",
+    "metadata",
+    "started_at",
+    "finished_at",
+    "exported_at",
+]
+
+_COLUMN_DEFS: Dict[str, str] = {
+    "trace_id":         "TEXT        NOT NULL",
+    "name":             "TEXT        NOT NULL",
+    "session_id":       "TEXT",
+    "user_id":          "TEXT",
+    "project_id":       "TEXT",
+    "model":            "TEXT",
+    "total_tokens_in":  "INTEGER     DEFAULT 0",
+    "total_tokens_out": "INTEGER     DEFAULT 0",
+    "total_tokens":     "INTEGER     DEFAULT 0",
+    "cost_usd":         "DOUBLE PRECISION DEFAULT 0",
+    "latency_ms":       "INTEGER",
+    "tools_used":       "JSONB       DEFAULT '{}'",
+    "spans":            "JSONB       DEFAULT '[]'",
+    "metadata":         "JSONB       DEFAULT '{}'",
+    "started_at":       "TIMESTAMPTZ NOT NULL",
+    "finished_at":      "TIMESTAMPTZ",
+    "exported_at":      "TIMESTAMPTZ NOT NULL",
+}
+
+_REQUIRED_COLUMNS: Set[str] = {"trace_id", "started_at", "exported_at"}
+
+_JSONB_COLUMNS: Set[str] = {"tools_used", "spans", "metadata"}
+
+
+def _build_create_sql(table: str, columns: List[str]) -> str:
+    col_parts = ["id BIGSERIAL PRIMARY KEY"]
+    for col in columns:
+        col_parts.append(f"    {col} {_COLUMN_DEFS[col]}")
+    if "trace_id" in columns:
+        col_parts.append(f"    CONSTRAINT {table}_trace_id_unique UNIQUE (trace_id)")
+    return f"CREATE TABLE IF NOT EXISTS {table} (\n" + ",\n".join(col_parts) + "\n);"
+
+
+def _build_insert_sql(table: str, columns: List[str]) -> str:
+    update_cols = [c for c in columns if c not in ("trace_id", "started_at")]
+    placeholders = ", ".join(f"%({c})s" for c in columns)
+    col_list = ", ".join(columns)
+    set_clause = ",\n    ".join(f"{c} = EXCLUDED.{c}" for c in update_cols)
+    return (
+        f"INSERT INTO {table} ({col_list})\n"
+        f"VALUES ({placeholders})\n"
+        f"ON CONFLICT (trace_id) DO UPDATE SET\n    {set_clause};"
+    )
+
+
+class PostgresExporter(BaseExporter):
+
+    def __init__(
+        self,
+        dsn: str,
+        table: str = "traces",
+        autocommit: bool = True,
+        include_fields: Optional[Iterable[str]] = None,
+        exclude_fields: Optional[Iterable[str]] = None,
+    ):
+        try:
+            import psycopg2
+            import psycopg2.extras
+        except ImportError as exc:
+            raise ImportError(
+                "psycopg2 is required for PostgresExporter. "
+                "Install it with: pip install psycopg2-binary"
+            ) from exc
+
+        self._psycopg2 = psycopg2
+        self._extras = psycopg2.extras
+        self._dsn = dsn
+        self._table = table
+        self._autocommit = autocommit
+        self._conn = None
+
+        if include_fields is not None:
+            chosen = set(include_fields) | _REQUIRED_COLUMNS
+            self._columns = [c for c in _ALL_COLUMNS if c in chosen]
+        elif exclude_fields is not None:
+            excluded = set(exclude_fields) - _REQUIRED_COLUMNS
+            self._columns = [c for c in _ALL_COLUMNS if c not in excluded]
+        else:
+            self._columns = list(_ALL_COLUMNS)
+
+        self._create_sql = _build_create_sql(self._table, self._columns)
+        self._insert_sql = _build_insert_sql(self._table, self._columns)
+        self._ensure_table()
+
+    def _get_conn(self):
+        if self._conn is None or self._conn.closed:
+            self._conn = self._psycopg2.connect(self._dsn)
+            self._conn.autocommit = self._autocommit
+        return self._conn
+
+    def _ensure_table(self) -> None:
+        conn = self._get_conn()
+        with conn.cursor() as cur:
+            cur.execute(self._create_sql)
+        if not self._autocommit:
+            conn.commit()
+
+    def _build_row(self, doc: dict) -> dict:
+        full_row: Dict[str, Any] = {
+            "trace_id":         doc["trace_id"],
+            "name":             doc["name"],
+            "session_id":       doc.get("session_id"),
+            "user_id":          doc.get("user_id"),
+            "project_id":       doc.get("project_id"),
+            "model":            doc.get("model"),
+            "total_tokens_in":  doc["total_tokens_in"],
+            "total_tokens_out": doc["total_tokens_out"],
+            "total_tokens":     doc["total_tokens"],
+            "cost_usd":         doc["cost_usd"],
+            "latency_ms":       doc.get("latency_ms"),
+            "tools_used":       self._extras.Json(doc["tools_used"]),
+            "spans":            self._extras.Json(doc["spans"]),
+            "metadata":         self._extras.Json(doc["metadata"]),
+            "started_at":       doc["started_at"],
+            "finished_at":      doc.get("finished_at"),
+            "exported_at":      datetime.now(timezone.utc).isoformat(),
+        }
+        return {k: v for k, v in full_row.items() if k in self._columns}
+
+    def export(self, trace: Trace) -> None:
+        doc = trace.to_dict()
+        row = self._build_row(doc)
+        conn = self._get_conn()
+        with conn.cursor() as cur:
+            cur.execute(self._insert_sql, row)
+        if not self._autocommit:
+            conn.commit()
+
+    def close(self) -> None:
+        if self._conn and not self._conn.closed:
+            self._conn.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *_):
+        self.close()