aice-observability 0.2.0__py3-none-any.whl

aice_observability/__init__.py

@@ -0,0 +1,5 @@
+from aice_observability.builder import Observability
+
+__version__ = "0.2.0"
+
+__all__ = ["Observability", "__version__"]

aice_observability/ai/agents.py

@@ -0,0 +1,169 @@
+"""
+High-level tracing helpers for AI agent orchestration.
+
+Unlike tracing.py, which records a single provider call,
+this module records the lifecycle of an AI agent.
+
+Example hierarchy
+
+invoke_agent
+    OpenAI chat
+    execute_tool
+    retrieve_documents
+    Gemini chat
+    execute_tool
+    Anthropic chat
+
+Everything underneath automatically becomes part of the same trace.
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from collections.abc import Iterator
+from opentelemetry.trace import SpanKind
+
+from aice_observability.ai.tracing import _span
+
+
+@contextmanager
+def invoke_agent(
+    *,
+    agent: str,
+    conversation_id: str | None = None,
+    session_id: str | None = None,
+    extra_attributes: dict | None = None,
+):
+    """
+    Records one complete execution of an AI agent.
+
+    This should wrap the outermost orchestration loop.
+
+    One invoke_agent span may contain
+
+        • multiple LLM calls
+        • multiple retrievals
+        • multiple tool executions
+        • sub-agents
+
+    ``session_id`` maps to ``gen_ai.session.id`` pass it only for a genuine
+    multi-run session. A single run is not a session; use
+    ``extra_attributes={"<namespace>.run.id": run_id}`` instead so dashboards
+    grouping by session do not conflate the two.
+
+    Example
+
+        with invoke_agent(agent="ResearchAgent"):
+            ...
+    """
+
+    attributes = {
+        "gen_ai.operation.name": "invoke_agent",
+        "gen_ai.agent.name": agent,
+    }
+
+    if conversation_id:
+        attributes["gen_ai.conversation.id"] = conversation_id
+
+    if session_id:
+        attributes["gen_ai.session.id"] = session_id
+
+    if extra_attributes:
+        attributes.update(extra_attributes)
+
+    with _span(
+        f"Agent {agent}",
+        kind=SpanKind.INTERNAL,
+        attributes=attributes,
+    ) as span:
+
+        yield span
+
+
+@contextmanager
+def execute_tool(
+    *,
+    tool: str,
+    extra_attributes: dict | None = None,
+):
+    """
+    Records execution of one tool.
+
+    Example
+
+        with execute_tool(tool="Vector Search"):
+            ...
+    """
+
+    attributes = {
+        "gen_ai.operation.name": "execute_tool",
+        "gen_ai.tool.name": tool,
+    }
+    if extra_attributes:
+        attributes.update(extra_attributes)
+
+    with _span(
+        f"Tool {tool}",
+        kind=SpanKind.INTERNAL,
+        attributes=attributes,
+    ) as span:
+
+        yield span
+        
+        
+@contextmanager
+def invoke_subagent(
+    *,
+    agent: str,
+):
+    """
+    Records a child agent.
+
+    This becomes a nested span beneath the current agent.
+    """
+
+    with _span(
+        f"Subagent {agent}",
+        kind=SpanKind.INTERNAL,
+        attributes={
+            "gen_ai.operation.name": "invoke_agent",
+            "gen_ai.agent.name": agent,
+        },
+    ) as span:
+
+        yield span
+        
+        
+@contextmanager
+def execute_workflow(
+    *,
+    name: str,
+):
+    """
+    Records an end-to-end AI workflow.
+
+    Example
+
+        Upload CSV
+            ▼
+        Profile data
+            ▼
+        Retrieve metadata
+            ▼
+        Run LLM
+            ▼
+        Validate
+            ▼
+        Store output
+    """
+
+    with _span(
+        f"Workflow {name}",
+        kind=SpanKind.INTERNAL,
+        attributes={
+            "gen_ai.operation.name": "workflow",
+            "workflow.name": name,
+        },
+    ) as span:
+
+        yield span

aice_observability/ai/instrumentation/anthropic.py

@@ -0,0 +1,20 @@
+from anthropic import Anthropic
+
+from aice_observability.ai.tracing import trace_provider_call
+
+
+class AnthropicProvider:
+
+    def __init__(self, client: Anthropic):
+        self.client = client
+
+    def chat(self, **kwargs):
+
+        model = kwargs.get("model", "unknown")
+
+        return trace_provider_call(
+            provider="anthropic",
+            model=model,
+            operation="messages.create",
+            fn=lambda: self.client.messages.create(**kwargs),
+        )

aice_observability/ai/instrumentation/base.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from abc import ABC
+from abc import abstractmethod
+from typing import Any
+
+
+class LLMProvider(ABC):
+    """
+    Base class implemented by every supported LLM provider.
+
+    Implementations are intentionally very small.
+
+    Their responsibility is simply to execute the SDK call while emitting
+    OpenTelemetry spans that follow the GenAI semantic conventions.
+    """
+
+    @property
+    @abstractmethod
+    def provider(self) -> str:
+        ...
+
+    @abstractmethod
+    def chat(self, **kwargs: Any) -> Any:
+        ...

aice_observability/ai/instrumentation/gemini.py

@@ -0,0 +1,20 @@
+from google import genai
+
+from aice_observability.ai.tracing import trace_provider_call
+
+
+class GeminiProvider:
+
+    def __init__(self, client: genai.Client):
+        self.client = client
+
+    def generate(self, **kwargs):
+
+        model = kwargs.get("model", "unknown")
+
+        return trace_provider_call(
+            provider="gemini",
+            model=model,
+            operation="generate_content",
+            fn=lambda: self.client.models.generate_content(**kwargs),
+        )

aice_observability/ai/instrumentation/langchain.py

@@ -0,0 +1,16 @@
+from aice_observability.ai.tracing import trace_provider_call
+
+
+class LangChainProvider:
+
+    def __init__(self, chain):
+        self.chain = chain
+
+    def invoke(self, inputs):
+
+        return trace_provider_call(
+            provider="langchain",
+            model=self.chain.__class__.__name__,
+            operation="invoke",
+            fn=lambda: self.chain.invoke(inputs),
+        )

aice_observability/ai/instrumentation/openai.py

@@ -0,0 +1,19 @@
+from openai import OpenAI
+from aice_observability.ai.tracing import trace_provider_call
+
+
+class OpenAIProvider:
+
+    def __init__(self, client: OpenAI):
+        self.client = client
+
+    def chat(self, **kwargs):
+
+        model = kwargs.get("model", "unknown")
+
+        return trace_provider_call(
+            provider="openai",
+            model=model,
+            operation="chat",
+            fn=lambda: self.client.responses.create(**kwargs),
+        )

aice_observability/ai/instrumentation/registry.py

@@ -0,0 +1,53 @@
+"""
+Registry for AI instrumentors.
+
+Observability.start() should not know about every provider individually.
+
+Instead it delegates to this registry.
+"""
+
+from __future__ import annotations
+
+import logging
+
+from aice_observability.core.config import ObservabilityConfig
+
+from .base import AIInstrumentor
+
+logger = logging.getLogger(__name__)
+
+
+class InstrumentationRegistry:
+
+    def __init__(self) -> None:
+        self._instrumentors: list[AIInstrumentor] = []
+
+    def register(
+        self,
+        instrumentor: AIInstrumentor,
+    ) -> None:
+        self._instrumentors.append(instrumentor)
+
+    def instrument_all(
+        self,
+        config: ObservabilityConfig,
+    ) -> None:
+
+        for instrumentor in self._instrumentors:
+
+            try:
+
+                success = instrumentor.instrument(config)
+
+                if success:
+                    logger.info(
+                        "Instrumented %s",
+                        instrumentor.name,
+                    )
+
+            except Exception:
+
+                logger.exception(
+                    "Failed to instrument %s",
+                    instrumentor.name,
+                )

aice_observability/ai/instrumentation/vertexai.py

@@ -0,0 +1,18 @@
+from vertexai.generative_models import GenerativeModel
+
+from aice_observability.ai.tracing import trace_provider_call
+
+
+class VertexAIProvider:
+
+    def __init__(self, model: GenerativeModel):
+        self.model = model
+
+    def generate(self, prompt: str):
+
+        return trace_provider_call(
+            provider="vertexai",
+            model=self.model._model_name,
+            operation="generate_content",
+            fn=lambda: self.model.generate_content(prompt),
+        )

aice_observability/ai/metrics.py

@@ -0,0 +1,216 @@
+"""
+OpenTelemetry metrics for AI applications.
+
+Unlike tracing.py, which records one span per request, this module
+produces aggregated metrics suitable for Prometheus.
+
+All metrics intentionally use LOW-CARDINALITY labels.
+
+Do NOT attach
+
+    user IDs
+    workspace IDs
+    prompts
+    document IDs
+    conversation IDs
+
+Those belong in traces instead.
+"""
+
+from __future__ import annotations
+
+from opentelemetry import metrics
+
+_meter = metrics.get_meter("observability.ai")
+
+
+# COUNTERS
+llm_requests = _meter.create_counter(
+    "gen_ai.client.requests",
+    description="Total LLM requests",
+)
+
+llm_failures = _meter.create_counter(
+    "gen_ai.client.failures",
+    description="Failed LLM requests",
+)
+
+tool_calls = _meter.create_counter(
+    "gen_ai.tool.calls",
+)
+
+tool_failures = _meter.create_counter(
+    "gen_ai.tool.failures",
+)
+
+retrieval_requests = _meter.create_counter(
+    "gen_ai.retrieval.requests",
+)
+
+embedding_requests = _meter.create_counter(
+    "gen_ai.embedding.requests",
+)
+
+#: Budget governance events (per-run warning / per-run exceeded / tenant
+#: window exhausted). A single `kind` label keeps cardinality at three;
+#: tenants stay off metrics per the cardinality rule.
+budget_events = _meter.create_counter(
+    "gen_ai.budget.events",
+)
+
+
+# HISTOGRAMS
+llm_latency = _meter.create_histogram(
+    "gen_ai.client.duration",
+    unit="s",
+)
+
+tool_latency = _meter.create_histogram(
+    "gen_ai.tool.duration",
+    unit="s",
+)
+
+retrieval_latency = _meter.create_histogram(
+    "gen_ai.retrieval.duration",
+    unit="s",
+)
+
+input_tokens = _meter.create_histogram(
+    "gen_ai.usage.input_tokens",
+)
+
+output_tokens = _meter.create_histogram(
+    "gen_ai.usage.output_tokens",
+)
+
+total_tokens = _meter.create_histogram(
+    "gen_ai.usage.total_tokens",
+)
+
+cost_usd = _meter.create_histogram(
+    "gen_ai.usage.cost",
+    unit="USD",
+)
+
+retrieved_documents = _meter.create_histogram(
+    "gen_ai.retrieval.documents",
+)
+
+
+def record_llm_request(*, provider: str, model: str):
+    llm_requests.add(
+        1,
+        {
+            "provider": provider,
+            "model": model,
+        },
+    )
+    
+    
+def record_llm_failure(*, provider: str, model: str):
+    llm_failures.add(
+        1,
+        {
+            "provider": provider,
+            "model": model,
+        },
+    )
+    
+    
+def record_latency(*, provider: str, model: str, duration: float):
+    llm_latency.record(
+        duration,
+        {
+            "provider": provider,
+            "model": model,
+        },
+    )
+    
+    
+def record_token_usage(*, provider: str, model: str, input: int,output: int):
+
+    input_tokens.record(
+        input,
+        {
+            "provider": provider,
+            "model": model,
+        },
+    )
+
+    output_tokens.record(
+        output,
+        {
+            "provider": provider,
+            "model": model,
+        },
+    )
+
+    total_tokens.record(
+        input + output,
+        {
+            "provider": provider,
+            "model": model,
+        },
+    )
+    
+    
+def record_cost(*, provider: str, model: str, usd: float):
+    cost_usd.record(
+        usd,
+        {
+            "provider": provider,
+            "model": model,
+        },
+    )
+    
+    
+def record_tool_call(tool: str):
+    tool_calls.add(
+        1,
+        {
+            "tool": tool,
+        },
+    )
+    
+    
+def record_tool_latency(*, tool: str, duration: float):
+    tool_latency.record(
+        duration,
+        {
+            "tool": tool,
+        },
+    )
+    
+    
+def record_retrieval(*, vector_store: str, document_count: int, duration: float):
+
+    retrieval_requests.add(
+        1,
+        {
+            "store": vector_store,
+        },
+    )
+
+    retrieval_latency.record(
+        duration,
+        {
+            "store": vector_store,
+        },
+    )
+
+    retrieved_documents.record(
+        document_count,
+        {
+            "store": vector_store,
+        },
+    )
+    
+    
+def record_budget_event(kind: str):
+    """
+    One budget governance event. ``kind`` is fixed-cardinality by design:
+    warning, run_exceeded, tenant_exhausted anything else is dropped by the
+    caller, not by an allowlist here, so a new kind surfaces loudly in a
+    dashboard rather than silently never being counted.
+    """
+    budget_events.add(1, {"kind": kind})