veris-ai 1.0.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

veris_ai/__init__.py

@@ -1,7 +1,41 @@
 """Veris AI Python SDK."""
 
+from typing import Any
+
 __version__ = "0.1.0"
 
+# Import lightweight modules that only use base dependencies
+from .jaeger_interface import JaegerClient, SearchQuery
 from .tool_mock import veris
 
-__all__ = ["veris"]
+# Lazy import for modules with heavy dependencies
+_instrument = None
+
+
+def instrument(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
+    """Lazy loader for the instrument function from braintrust_tracing.
+
+    This function requires the 'instrument' extra dependencies:
+    pip install veris-ai[instrument]
+    """
+    global _instrument  # noqa: PLW0603
+    if _instrument is None:
+        try:
+            from .braintrust_tracing import instrument as _instrument_impl  # noqa: PLC0415
+
+            _instrument = _instrument_impl
+        except ImportError as e:
+            error_msg = (
+                "The 'instrument' function requires additional dependencies. "
+                "Please install them with: pip install veris-ai[instrument]"
+            )
+            raise ImportError(error_msg) from e
+    return _instrument(*args, **kwargs)
+
+
+__all__ = [
+    "veris",
+    "JaegerClient",
+    "SearchQuery",
+    "instrument",
+]

veris_ai/braintrust_tracing.py

@@ -0,0 +1,282 @@
+"""Non-invasive Braintrust + Jaeger (OTEL) instrumentation helper for the `openai-agents` SDK.
+
+Typical usage
+-------------
+>>> from our_sdk import braintrust_tracing
+>>> braintrust_tracing.instrument(service_name="openai-agent")
+
+After calling :func:`instrument`, any later call to
+``agents.set_trace_processors([...])`` will be transparently patched so that:
+  • the list always contains a BraintrustTracingProcessor (for Braintrust UI)
+  • *and* an OpenTelemetry bridge processor that mirrors every span to the
+    global OTEL tracer provider (Jaeger by default).
+
+Goal: deliver full OTEL compatibility while keeping the official Braintrust
+SDK integration unchanged – **no code modifications required** besides the
+single `instrument()` call.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+from typing import Any, cast
+
+import wrapt  # type: ignore[import-untyped, import-not-found]
+from braintrust.wrappers.openai import (
+    BraintrustTracingProcessor,  # type: ignore[import-untyped, import-not-found]
+)
+from opentelemetry import context as otel_context  # type: ignore[import-untyped, import-not-found]
+from opentelemetry import trace  # type: ignore[import-untyped, import-not-found]
+from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
+    OTLPSpanExporter,  # type: ignore[import-untyped, import-not-found]
+)
+from opentelemetry.sdk.resources import (  # type: ignore[import-untyped, import-not-found]
+    SERVICE_NAME,
+    Resource,
+)
+from opentelemetry.sdk.trace import TracerProvider  # type: ignore[import-untyped, import-not-found]
+from opentelemetry.sdk.trace.export import (
+    BatchSpanProcessor,  # type: ignore[import-untyped, import-not-found]
+)
+from opentelemetry.trace import SpanKind  # type: ignore[import-untyped, import-not-found]
+
+from veris_ai.tool_mock import _session_id_context
+
+# ---------------------------------------------------------------------------
+#  Optional import of *agents* – we fail lazily at runtime if missing.
+# ---------------------------------------------------------------------------
+try:
+    import agents  # type: ignore[import-untyped]  # noqa: TC002
+    from agents import TracingProcessor  # type: ignore[import-untyped]
+
+    try:
+        from agents.tracing import get_trace_provider  # type: ignore[import-untyped]
+    except ImportError:
+        # Fallback for newer versions that have GLOBAL_TRACE_PROVIDER instead
+        from agents.tracing import (  # type: ignore[import-untyped, attr-defined, import-not-found]
+            GLOBAL_TRACE_PROVIDER,  # type: ignore[import-untyped, attr-defined, import-not-found]
+        )
+
+        get_trace_provider = lambda: GLOBAL_TRACE_PROVIDER  # type: ignore[no-any-return]  # noqa: E731
+except ModuleNotFoundError as exc:  # pragma: no cover
+    _IMPORT_ERR: ModuleNotFoundError | None = exc
+    TracingProcessor = object  # type: ignore[assignment, misc]
+    get_trace_provider = None  # type: ignore[assignment]
+else:
+    _IMPORT_ERR = None
+
+__all__ = ["instrument"]
+
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+#  Internal helper – OTEL bridge processor
+# ---------------------------------------------------------------------------
+
+
+class AgentsOTELBridgeProcessor(TracingProcessor):  # type: ignore[misc]
+    """Mirrors every Agents span into a dedicated OTEL tracer provider."""
+
+    def __init__(
+        self,
+        braintrust_processor: BraintrustTracingProcessor,
+        *,
+        service_name: str,  # noqa: ARG002
+        tracer_provider: trace.TracerProvider,
+    ) -> None:  # noqa: D401,E501
+        self._braintrust = braintrust_processor
+        self._tracer = tracer_provider.get_tracer(__name__)
+        self._otel_spans: dict[str, trace.Span] = {}
+        self._provider = tracer_provider
+
+    # ----------------------------- utils ---------------------------------
+    @staticmethod
+    def _flatten(prefix: str, obj: Any, out: dict[str, Any]) -> None:  # noqa: PLR0911, ANN401
+        """Flatten complex objects into OTEL-compatible primitives."""
+        if isinstance(obj, dict):
+            for k, v in obj.items():
+                AgentsOTELBridgeProcessor._flatten(f"{prefix}.{k}" if prefix else str(k), v, out)
+        elif isinstance(obj, str | int | float | bool) or obj is None:
+            out[prefix] = obj
+        elif isinstance(obj, list | tuple):
+            try:
+                if all(isinstance(i, str | int | float | bool) or i is None for i in obj):
+                    out[prefix] = list(obj)
+                else:
+                    out[prefix] = json.dumps(obj, default=str)
+            except Exception:  # pragma: no cover – defensive
+                out[prefix] = json.dumps(obj, default=str)
+        else:
+            out[prefix] = str(obj)
+
+    def _log_data_attributes(self, span_obj: agents.tracing.Span) -> dict[str, Any]:  # type: ignore[name-defined]
+        data = self._braintrust._log_data(span_obj)  # pyright: ignore[reportPrivateUsage]  # noqa: SLF001
+        flat: dict[str, Any] = {}
+        self._flatten("bt", data, flat)
+
+        # Add session_id if available
+        session_id = _session_id_context.get()
+        if session_id:
+            flat["veris.session_id"] = session_id
+
+        return {k: v for k, v in flat.items() if v is not None}
+
+    # --------------------- Agents lifecycle hooks ------------------------
+    def on_trace_start(self, trace_obj: Any) -> None:  # noqa: ANN401, D102
+        # Get session_id at trace start
+        session_id = _session_id_context.get()
+        attributes = {"veris.session_id": session_id} if session_id else {}
+
+        otel_span = self._tracer.start_span(
+            name=trace_obj.name or "agent-trace",
+            kind=SpanKind.INTERNAL,
+            attributes=attributes,
+        )
+        self._otel_spans[trace_obj.trace_id] = otel_span
+        logger.info(f"VERIS AI BraintrustTracingProcessor: on_trace_start: {trace_obj.trace_id}")
+
+    def on_trace_end(self, trace_obj: Any) -> None:  # noqa: ANN401, D102
+        span = self._otel_spans.pop(trace_obj.trace_id, None)
+        if span:
+            span.end()
+
+    def on_span_start(self, span: Any) -> None:  # noqa: ANN401, D102
+        parent_otel = (
+            self._otel_spans.get(span.parent_id)
+            if span.parent_id
+            else self._otel_spans.get(span.trace_id)
+        )
+        parent_ctx = (
+            trace.set_span_in_context(parent_otel) if parent_otel else otel_context.get_current()
+        )
+
+        # Get session_id at span start
+        session_id = _session_id_context.get()
+        attributes = {"veris.session_id": session_id} if session_id else {}
+
+        child = self._tracer.start_span(
+            name=span.span_data.__class__.__name__,
+            context=parent_ctx,
+            kind=SpanKind.INTERNAL,
+            attributes=attributes,
+        )
+        self._otel_spans[span.span_id] = child
+        logger.info(f"VERIS AI BraintrustTracingProcessor: on_span_start: {span.span_id}")
+
+    def on_span_end(self, span: Any) -> None:  # noqa: ANN401, D102
+        child = self._otel_spans.pop(span.span_id, None)
+        logger.info(f"VERIS AI BraintrustTracingProcessor: on_span_end: {span.span_id}")
+        if child:
+            for k, v in self._log_data_attributes(span).items():
+                try:
+                    child.set_attribute(k, v)
+                except Exception:  # pragma: no cover – bad value type  # noqa: S112
+                    continue
+            child.end()
+
+    # --------------------- house-keeping ---------------------------------
+    def shutdown(self) -> None:  # noqa: D401
+        provider = cast("TracerProvider", self._provider)
+        provider.shutdown()  # type: ignore[attr-defined]
+
+    def force_flush(self) -> None:  # noqa: D401
+        provider = cast("TracerProvider", self._provider)
+        provider.force_flush()  # type: ignore[attr-defined]
+
+
+# ---------------------------------------------------------------------------
+#  Public entry point
+# ---------------------------------------------------------------------------
+
+_PATCHED: bool = False  # ensure idempotent patching
+
+
+def instrument(
+    *,
+    service_name: str | None = None,
+    otlp_endpoint: str | None = None,
+) -> None:
+    """Bootstrap Braintrust + OTEL instrumentation and patch Agents SDK.
+
+    Invoke once at any point before `set_trace_processors` is called.
+    """
+    global _PATCHED  # noqa: PLW0603
+    if _PATCHED:
+        return  # already done
+
+    if _IMPORT_ERR is not None or get_trace_provider is None:  # pragma: no cover
+        error_msg = "The `agents` package is required but not installed"
+        raise RuntimeError(error_msg) from _IMPORT_ERR
+
+    # ------------------ 0. Validate inputs -----------------------------
+    # Resolve service name ─ explicit argument → env var → error
+    if not service_name or not str(service_name).strip():
+        service_name = os.getenv("VERIS_SERVICE_NAME")
+
+    if not service_name or not str(service_name).strip():
+        error_msg = (
+            "`service_name` must be provided either as an argument or via the "
+            "VERIS_SERVICE_NAME environment variable"
+        )
+        raise ValueError(error_msg)
+
+    # Resolve OTLP endpoint ─ explicit argument → env var → error
+    if not otlp_endpoint or not str(otlp_endpoint).strip():
+        otlp_endpoint = os.getenv("VERIS_OTLP_ENDPOINT")
+
+    if not otlp_endpoint or not str(otlp_endpoint).strip():
+        error_msg = (
+            "`otlp_endpoint` must be provided either as an argument or via the "
+            "VERIS_OTLP_ENDPOINT environment variable"
+        )
+        raise ValueError(error_msg)
+
+    logger.info(f"service_name: {service_name}")
+    logger.info(f"otlp_endpoint: {otlp_endpoint}")
+
+    # ------------------ 1. Configure OTEL provider ---------------------
+    # We create our own provider instance and do NOT set it globally.
+    # This avoids conflicts with any other OTEL setup in the application.
+    otel_provider = TracerProvider(resource=Resource.create({SERVICE_NAME: service_name}))
+    otel_provider.add_span_processor(
+        BatchSpanProcessor(OTLPSpanExporter(endpoint=otlp_endpoint, insecure=True)),
+    )
+
+    # ------------------ 2. Define wrapper for patching -------------------
+    def _wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any:  # noqa: ANN401, ARG001
+        """This function wraps `TraceProvider.set_processors`."""
+        processors = args[0] if args else []
+
+        # Find the user's Braintrust processor to pass to our bridge.
+        bt_processor = next(
+            (p for p in processors if isinstance(p, BraintrustTracingProcessor)), None
+        )
+
+        # If no Braintrust processor is present, our bridge is useless.
+        # Also, if a bridge is already there, don't add another one.
+        has_bridge = any(isinstance(p, AgentsOTELBridgeProcessor) for p in processors)
+        if not bt_processor or has_bridge:
+            return wrapped(*args, **kwargs)
+
+        # Create the bridge and add it to the list of processors.
+        bridge = AgentsOTELBridgeProcessor(
+            bt_processor,
+            service_name=service_name,
+            tracer_provider=otel_provider,
+        )
+        new_processors = list(processors) + [bridge]
+
+        # Call the original function with the augmented list.
+        new_args = (new_processors,) + args[1:]
+        logger.info(f"VERIS AI BraintrustTracingProcessor: {new_args}")
+        return wrapped(*new_args, **kwargs)
+
+    # ------------------ 3. Patch the provider instance -------------------
+    # This is more robust than patching the function, as it's independent
+    # of how the user imports `set_trace_processors`.
+    provider_instance = get_trace_provider()
+    wrapt.wrap_function_wrapper(provider_instance, "set_processors", _wrapper)
+
+    _PATCHED = True

veris_ai/jaeger_interface/README.md

@@ -0,0 +1,109 @@
+# Jaeger Interface
+
+This sub-package ships a **thin synchronous wrapper** around the
+[Jaeger Query Service](https://www.jaegertracing.io/docs/) HTTP API so
+that you can **search for and retrieve traces** directly from Python
+with minimal boilerplate.
+
+> The client relies on `requests` (already included in the SDK’s
+> dependencies) and uses *pydantic* for full type-safety.
+
+---
+
+## Installation
+
+`veris-ai` already lists both `requests` and `pydantic` as hard
+requirements, so **no additional dependencies are required**.
+
+```bash
+pip install veris-ai
+```
+
+---
+
+## Quick-start
+
+```python
+from veris_ai.jaeger_interface import JaegerClient, SearchQuery
+import json
+from veris_ai.jaeger_interface.models import Trace
+# Replace with the URL of your Jaeger Query Service instance
+client = JaegerClient("http://localhost:16686")
+
+# --- 1. Search traces --------------------------------------------------
+resp = client.search(
+    SearchQuery(
+        service="veris-agent",
+        limit=3,
+        operation="CustomSpanData",
+        tags={"bt.metadata.session_id":"oRL1_IhMP3s7T7mYrixCW"}
+    )
+)
+
+# save to json
+with open("resp.json", "w") as f:
+    f.write(resp.model_dump_json(indent=2))
+
+# Guard clause
+if not resp or not resp.data:
+    print("No data found")
+    exit(1)
+
+# Print trace ids
+for trace in resp.data:
+    if isinstance(trace, Trace):
+        print("TRACE ID:", trace.traceID, len(trace.spans), "spans")
+
+# --- 2. Retrieve a specific trace -------------------------------------
+if isinstance(resp.data, list):
+    trace_id = resp.data[0].traceID
+else:
+    trace_id = resp.data.traceID
+
+detailed = client.get_trace(trace_id)
+# save detailed to json
+with open("detailed.json", "w") as f:
+    f.write(detailed.model_dump_json(indent=2))
+```
+
+---
+
+## API Reference
+
+### `JaegerClient`
+
+| Method | Description |
+| -------- | ----------- |
+| `search(query: SearchQuery) -> SearchResponse` | Search for traces matching the given parameters (wrapper around `/api/traces`). |
+| `get_trace(trace_id: str) -> GetTraceResponse` | Fetch a single trace by ID  (wrapper around `/api/traces/{id}`). |
+
+### `SearchQuery`
+
+`SearchQuery` is a *pydantic* model for building the query-string sent to
+Jaeger’s `/api/traces` endpoint.
+
+Parameter logic:
+
+* **service** – mandatory, sets the top-level service filter.
+* **operation** – optional; a trace is kept only if *any* of its spans has
+  this exact `operationName`.
+* **tags** – dict of key/value filters; *all* pairs must match on a single
+  span (logical AND). Example: `{"error": "true"}`.
+* **limit** – applied after all filters and trims the result list.
+
+Any other fields are forwarded untouched, so you can experiment with new
+Jaeger parameters without waiting for an SDK update.
+
+---
+
+## Compatibility
+
+The implementation targets **Jaeger v1.x** REST endpoints. For clusters
+backed by **OpenSearch** storage the same endpoints apply. Should you
+need API v3 support feel free to open an issue or contribution—thanks!
+
+---
+
+## Licence
+
+This package is released under the **MIT license**.

veris_ai/jaeger_interface/__init__.py

@@ -0,0 +1,26 @@
+"""Jaeger interface for searching and retrieving traces.
+
+This sub-package provides a thin synchronous wrapper around the Jaeger
+Query Service HTTP API.
+
+Typical usage example::
+
+    from veris_ai.jaeger_interface import JaegerClient, SearchQuery
+
+    client = JaegerClient("http://localhost:16686")
+
+    traces = client.search(
+        SearchQuery(service="veris-agent", limit=20)
+    )
+
+    trace = client.get_trace(traces.data[0].traceID)
+
+The implementation uses *requests* under the hood and all public functions
+are fully typed using *pydantic* models so that IDEs can provide proper
+autocomplete and type checking.
+"""
+
+from .client import JaegerClient
+from .models import SearchQuery
+
+__all__ = ["JaegerClient", "SearchQuery"]

veris_ai/jaeger_interface/client.py

@@ -0,0 +1,133 @@
+"""Synchronous Jaeger Query Service client built on **requests**.
+
+This implementation keeps dependencies minimal while providing fully-typed
+*pydantic* models for both **request** and **response** bodies.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Self
+
+import requests
+
+from .models import GetTraceResponse, SearchQuery, SearchResponse
+
+if TYPE_CHECKING:
+    from types import TracebackType
+
+__all__ = ["JaegerClient"]
+
+
+class JaegerClient:  # noqa: D101
+    def __init__(
+        self,
+        base_url: str,
+        *,
+        timeout: float | None = 10.0,
+        session: requests.Session | None = None,
+        headers: dict[str, str] | None = None,
+    ) -> None:
+        """Create a new *JaegerClient* instance.
+
+        Args:
+            base_url: Base URL of the Jaeger Query Service (e.g. ``http://localhost:16686``).
+            timeout: Request timeout in **seconds** (applied to every call).
+            session: Optional pre-configured :class:`requests.Session` to reuse.
+            headers: Optional default headers to send with every request.
+        """
+        # Normalise to avoid trailing slash duplicates
+        self._base_url = base_url.rstrip("/")
+        self._timeout = timeout
+        self._external_session = session  # If provided we won't close it
+        self._headers = headers or {}
+
+    # ---------------------------------------------------------------------
+    # Internal helpers
+    # ---------------------------------------------------------------------
+
+    def _make_session(self) -> tuple[requests.Session, bool]:  # noqa: D401
+        """Return a *(session, should_close)* tuple.
+
+        If an external session was supplied we **must not** close it after the
+        request, hence the boolean flag letting callers know whether they are
+        responsible for closing the session.
+        """
+        if self._external_session is not None:
+            return self._external_session, False
+
+        # Reuse the session opened via the context manager if available
+        if hasattr(self, "_session_ctx"):
+            return self._session_ctx, False
+
+        session = requests.Session()
+        session.headers.update(self._headers)
+        return session, True
+
+    # ---------------------------------------------------------------------
+    # Public API
+    # ---------------------------------------------------------------------
+
+    def search(self, query: SearchQuery) -> SearchResponse:  # noqa: D401
+        """Search traces using the *v1* ``/api/traces`` endpoint.
+
+        Args:
+            query: :class:`~veris_ai.jaeger_interface.models.SearchQuery` instance.
+
+        Returns:
+            Parsed :class:`~veris_ai.jaeger_interface.models.SearchResponse` model.
+        """
+        params = query.to_params()
+        session, should_close = self._make_session()
+        try:
+            url = f"{self._base_url}/api/traces"
+            response = session.get(url, params=params, timeout=self._timeout)
+            response.raise_for_status()
+            data = response.json()
+        finally:
+            if should_close:
+                session.close()
+        return SearchResponse.model_validate(data)  # type: ignore[arg-type]
+
+    def get_trace(self, trace_id: str) -> GetTraceResponse:  # noqa: D401
+        """Retrieve a single trace by *trace_id*.
+
+        Args:
+            trace_id: The Jaeger trace identifier.
+
+        Returns:
+            Parsed :class:`~veris_ai.jaeger_interface.models.GetTraceResponse` model.
+        """
+        if not trace_id:
+            error_msg = "trace_id must be non-empty"
+            raise ValueError(error_msg)
+
+        session, should_close = self._make_session()
+        try:
+            url = f"{self._base_url}/api/traces/{trace_id}"
+            response = session.get(url, timeout=self._timeout)
+            response.raise_for_status()
+            data = response.json()
+        finally:
+            if should_close:
+                session.close()
+        return GetTraceResponse.model_validate(data)  # type: ignore[arg-type]
+
+    # ------------------------------------------------------------------
+    # Context-manager helpers (optional but convenient)
+    # ------------------------------------------------------------------
+
+    def __enter__(self) -> Self:
+        """Enter the context manager."""
+        self._session_ctx, self._should_close_ctx = self._make_session()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        """Exit the context manager."""
+        # Only close if we created the session
+        if getattr(self, "_should_close_ctx", False):
+            self._session_ctx.close()