velune-cli 0.9.0__py3-none-any.whl

velune/telemetry/spans.py

@@ -0,0 +1,277 @@
+"""Span context system for correlation IDs and trace propagation.
+
+Provides:
+- Unique run_id (UUID) for each orchestration run
+- Span IDs for agent turns and tool executions
+- Parent-child span relationships
+- Automatic context binding via structlog
+"""
+
+from __future__ import annotations
+
+import contextvars
+import uuid
+from collections.abc import AsyncIterator, Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import Any
+
+import structlog
+
+logger = structlog.get_logger()
+
+# Context variables for span propagation
+_run_id_var: contextvars.ContextVar[str | None] = contextvars.ContextVar("run_id", default=None)
+_span_id_var: contextvars.ContextVar[str | None] = contextvars.ContextVar("span_id", default=None)
+_parent_span_id_var: contextvars.ContextVar[str | None] = contextvars.ContextVar(
+    "parent_span_id", default=None
+)
+
+
+@dataclass
+class SpanContext:
+    """Represents a span context for tracing."""
+
+    run_id: str
+    span_id: str
+    parent_span_id: str | None = None
+    operation_name: str = ""
+    start_time: float = field(default_factory=lambda: __import__("time").time())
+    end_time: float | None = None
+    duration_ms: float | None = None
+    status: str = "running"  # running, completed, failed
+    error: str | None = None
+    attributes: dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for logging."""
+        return {
+            "run_id": self.run_id,
+            "span_id": self.span_id,
+            "parent_span_id": self.parent_span_id,
+            "operation_name": self.operation_name,
+            "status": self.status,
+            "duration_ms": self.duration_ms,
+            "error": self.error,
+        }
+
+
+def create_run_id() -> str:
+    """Create a unique run ID for an orchestration run."""
+    return str(uuid.uuid4())[:8]
+
+
+def create_span_id() -> str:
+    """Create a unique span ID for an operation."""
+    return str(uuid.uuid4())[:8]
+
+
+def get_current_run_id() -> str | None:
+    """Get the current run ID from context."""
+    return _run_id_var.get()
+
+
+def get_current_span_id() -> str | None:
+    """Get the current span ID from context."""
+    return _span_id_var.get()
+
+
+def get_current_parent_span_id() -> str | None:
+    """Get the current parent span ID from context."""
+    return _parent_span_id_var.get()
+
+
+@contextmanager
+def span(
+    operation_name: str,
+    parent_run_id: str | None = None,
+    parent_span_id: str | None = None,
+    **attributes: Any,
+) -> Iterator[SpanContext]:
+    """Context manager for creating a new span.
+
+    Automatically binds span context to structlog for duration of operation.
+    Handles timing and status tracking.
+
+    Args:
+        operation_name: Name of the operation (e.g., "coder_execution")
+        parent_run_id: Run ID for this span (creates new if None)
+        parent_span_id: Parent span ID (creates new span if None)
+        **attributes: Additional attributes to track
+
+    Example:
+        with span("coder_execution", parent_run_id=run_id) as span_ctx:
+            logger.info("Starting coder", operation=span_ctx.operation_name)
+            result = await coder.run()
+            # Span duration automatically tracked
+    """
+    import time
+
+    # Create or use provided IDs
+    run_id = parent_run_id or get_current_run_id() or create_run_id()
+    span_id = create_span_id()
+    current_parent = parent_span_id or get_current_span_id()
+
+    # Create span context
+    span_ctx = SpanContext(
+        run_id=run_id,
+        span_id=span_id,
+        parent_span_id=current_parent,
+        operation_name=operation_name,
+        attributes=attributes,
+    )
+
+    # Save old context
+    old_run_id = _run_id_var.get()
+    old_span_id = _span_id_var.get()
+    old_parent_span_id = _parent_span_id_var.get()
+
+    try:
+        # Bind to context variables
+        _run_id_var.set(run_id)
+        _span_id_var.set(span_id)
+        _parent_span_id_var.set(current_parent)
+
+        # Bind to structlog
+        structlog.contextvars.bind_contextvars(
+            run_id=run_id,
+            span_id=span_id,
+            operation_name=operation_name,
+            **attributes,
+        )
+
+        logger.debug(
+            "Span started",
+            span_context=span_ctx.to_dict(),
+        )
+
+        yield span_ctx
+
+    except Exception as e:
+        span_ctx.status = "failed"
+        span_ctx.error = str(e)
+        logger.error(
+            "Span failed",
+            span_context=span_ctx.to_dict(),
+            error=str(e),
+        )
+        raise
+
+    finally:
+        # Record end time and duration
+        span_ctx.end_time = time.time()
+        span_ctx.duration_ms = (span_ctx.end_time - span_ctx.start_time) * 1000
+
+        if span_ctx.status == "running":
+            span_ctx.status = "completed"
+
+        logger.debug(
+            "Span ended",
+            span_context=span_ctx.to_dict(),
+        )
+
+        # Unbind from structlog
+        unbind_keys = ["operation_name"] + list(attributes.keys())
+        structlog.contextvars.unbind_contextvars(*unbind_keys)
+
+        # Restore old context
+        if old_run_id is not None:
+            _run_id_var.set(old_run_id)
+        else:
+            _run_id_var.set(None)
+
+        if old_span_id is not None:
+            _span_id_var.set(old_span_id)
+        else:
+            _span_id_var.set(None)
+
+        if old_parent_span_id is not None:
+            _parent_span_id_var.set(old_parent_span_id)
+        else:
+            _parent_span_id_var.set(None)
+
+
+async def async_span(
+    operation_name: str,
+    parent_run_id: str | None = None,
+    parent_span_id: str | None = None,
+    **attributes: Any,
+) -> AsyncIterator[SpanContext]:
+    """Async context manager for creating a new span.
+
+    Same as span() but for async operations.
+
+    Example:
+        async with async_span("council_run", parent_run_id=run_id) as span_ctx:
+            result = await orchestrator.run()
+    """
+    import time
+
+    run_id = parent_run_id or get_current_run_id() or create_run_id()
+    span_id = create_span_id()
+    current_parent = parent_span_id or get_current_span_id()
+
+    span_ctx = SpanContext(
+        run_id=run_id,
+        span_id=span_id,
+        parent_span_id=current_parent,
+        operation_name=operation_name,
+        attributes=attributes,
+    )
+
+    old_run_id = _run_id_var.get()
+    old_span_id = _span_id_var.get()
+    old_parent_span_id = _parent_span_id_var.get()
+
+    try:
+        _run_id_var.set(run_id)
+        _span_id_var.set(span_id)
+        _parent_span_id_var.set(current_parent)
+
+        structlog.contextvars.bind_contextvars(
+            run_id=run_id,
+            span_id=span_id,
+            operation_name=operation_name,
+            **attributes,
+        )
+
+        logger.debug("Span started", span_context=span_ctx.to_dict())
+
+        yield span_ctx
+
+    except Exception as e:
+        span_ctx.status = "failed"
+        span_ctx.error = str(e)
+        logger.error(
+            "Span failed",
+            span_context=span_ctx.to_dict(),
+            error=str(e),
+        )
+        raise
+
+    finally:
+        span_ctx.end_time = time.time()
+        span_ctx.duration_ms = (span_ctx.end_time - span_ctx.start_time) * 1000
+
+        if span_ctx.status == "running":
+            span_ctx.status = "completed"
+
+        logger.debug("Span ended", span_context=span_ctx.to_dict())
+
+        unbind_keys = ["operation_name"] + list(attributes.keys())
+        structlog.contextvars.unbind_contextvars(*unbind_keys)
+
+        if old_run_id is not None:
+            _run_id_var.set(old_run_id)
+        else:
+            _run_id_var.set(None)
+
+        if old_span_id is not None:
+            _span_id_var.set(old_span_id)
+        else:
+            _span_id_var.set(None)
+
+        if old_parent_span_id is not None:
+            _parent_span_id_var.set(old_parent_span_id)
+        else:
+            _parent_span_id_var.set(None)

velune/telemetry/token_tracker.py

@@ -0,0 +1,140 @@
+"""Token usage tracking and cost estimation for all inference calls."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime
+
+# ---------------------------------------------------------------------------
+# Per-provider cost table (USD per 1 000 tokens, blended input+output rate).
+# Zero means free tier.
+# ---------------------------------------------------------------------------
+PROVIDER_COSTS: dict[str, dict[str, float]] = {
+    "anthropic": {
+        "claude-opus-4-5": 0.015,
+        "claude-sonnet-4-5": 0.003,
+        "claude-haiku-4-5": 0.00025,
+    },
+    "openai": {
+        "gpt-4o": 0.005,
+        "gpt-4o-mini": 0.00015,
+        "o1": 0.015,
+        "o1-mini": 0.003,
+    },
+    "xai": {
+        "grok-2": 0.002,
+        "grok-2-mini": 0.0002,
+    },
+    "google": {
+        "gemini-2.0-flash": 0.000075,
+        "gemini-1.5-pro": 0.00125,
+        "gemini-1.5-flash": 0.000075,
+        "gemini-2.0-flash-thinking-exp": 0.0,
+    },
+    "groq": {
+        # All free tier
+        "llama-3.3-70b-versatile": 0.0,
+        "llama-3.1-8b-instant": 0.0,
+        "mixtral-8x7b-32768": 0.0,
+        "gemma2-9b-it": 0.0,
+        "llama-3.2-11b-vision-preview": 0.0,
+    },
+    "together": {
+        "meta-llama/Llama-3.3-70B-Instruct-Turbo": 0.00088,
+        "meta-llama/Llama-3.2-11B-Vision-Instruct-Turbo": 0.00018,
+        "Qwen/Qwen2.5-Coder-32B-Instruct": 0.0008,
+        "deepseek-ai/DeepSeek-R1": 0.003,
+        "mistralai/Mistral-7B-Instruct-v0.3": 0.0002,
+    },
+    "fireworks": {
+        "accounts/fireworks/models/llama-v3p3-70b-instruct": 0.0009,
+        "accounts/fireworks/models/deepseek-r1": 0.003,
+        "accounts/fireworks/models/qwen2p5-coder-32b-instruct": 0.0009,
+        "accounts/fireworks/models/mixtral-8x22b-instruct": 0.0009,
+    },
+    "openrouter": {},  # Dynamic — set from model metadata
+    "ollama": {},  # Always free (local)
+    "lmstudio": {},  # Always free (local)
+    "llamacpp": {},  # Always free (local)
+}
+
+
+@dataclass
+class TokenUsage:
+    model_id: str
+    provider_id: str
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+    cost_usd: float
+    timestamp: datetime = field(default_factory=datetime.now)
+
+    @classmethod
+    def from_response(
+        cls,
+        provider_id: str,
+        model_id: str,
+        prompt_tokens: int,
+        completion_tokens: int,
+    ) -> TokenUsage:
+        total = prompt_tokens + completion_tokens
+        cost = cls._calculate_cost(provider_id, model_id, total)
+        return cls(
+            model_id=model_id,
+            provider_id=provider_id,
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            total_tokens=total,
+            cost_usd=cost,
+        )
+
+    @staticmethod
+    def _calculate_cost(provider_id: str, model_id: str, tokens: int) -> float:
+        rate = PROVIDER_COSTS.get(provider_id, {}).get(model_id, 0.0)
+        return (tokens / 1000) * rate
+
+
+@dataclass
+class SessionUsage:
+    usages: list[TokenUsage] = field(default_factory=list)
+
+    def add(self, usage: TokenUsage) -> None:
+        self.usages.append(usage)
+
+    @property
+    def total_tokens(self) -> int:
+        return sum(u.total_tokens for u in self.usages)
+
+    @property
+    def total_cost(self) -> float:
+        return sum(u.cost_usd for u in self.usages)
+
+    @property
+    def prompt_tokens(self) -> int:
+        return sum(u.prompt_tokens for u in self.usages)
+
+    @property
+    def completion_tokens(self) -> int:
+        return sum(u.completion_tokens for u in self.usages)
+
+    def by_provider(self) -> dict[str, int]:
+        result: dict[str, int] = {}
+        for u in self.usages:
+            result[u.provider_id] = result.get(u.provider_id, 0) + u.total_tokens
+        return result
+
+    def summary_line(self) -> str:
+        parts = [f"{self.total_tokens:,} tokens"]
+        if self.total_cost > 0:
+            parts.append(f"~${self.total_cost:.4f}")
+        else:
+            parts.append("free")
+        parts.append("session total")
+        return " · ".join(parts)
+
+
+# ---------------------------------------------------------------------------
+# Process-level session tracker — populated by inference call sites
+# ---------------------------------------------------------------------------
+
+current_session: SessionUsage = SessionUsage()