agent-observatory-cli 0.1.0__py3-none-any.whl

observatory/auto.py

@@ -0,0 +1,212 @@
+"""Launch a Python agent with model and framework instrumentation enabled."""
+from __future__ import annotations
+
+import importlib
+import importlib.metadata
+import importlib.util
+import os
+import platform
+import queue
+import runpy
+import secrets
+import socket
+import sys
+import threading
+import time
+import urllib.error
+import urllib.request
+from pathlib import Path
+from urllib.parse import urlsplit
+
+from .client import Client
+
+
+class SetupError(Exception):
+    """A problem the user can fix before their agent starts."""
+
+
+INTEGRATIONS = (
+    ("openai", "openai", "OpenAIInstrumentor"),
+    ("anthropic", "anthropic", "AnthropicInstrumentor"),
+    ("langchain_core", "langchain", "LangChainInstrumentor"),
+)
+
+
+def _receiver_endpoint(endpoint):
+    parsed = urlsplit(endpoint)
+    if (parsed.scheme not in {"http", "https"} or not parsed.hostname or
+            parsed.username or parsed.password or parsed.query or parsed.fragment):
+        raise SetupError("Use the dashboard's HTTP address for --endpoint, such as http://127.0.0.1:4318.")
+    return endpoint.rstrip("/")
+
+
+def _check_receiver(endpoint, token):
+    headers = {"Authorization": f"Bearer {token}"} if token else {}
+    try:
+        with urllib.request.urlopen(urllib.request.Request(endpoint + "/api/health", headers=headers),
+                                    timeout=3) as response:
+            if response.status != 200:
+                raise SetupError("The Observatory receiver is not ready.")
+    except urllib.error.HTTPError as exc:
+        if exc.code == 401:
+            raise SetupError("This receiver needs a token. Set OBSERVATORY_TOKEN in your agent's terminal.") from exc
+        raise SetupError(f"The receiver returned HTTP {exc.code}. Check --endpoint.") from exc
+    except (OSError, urllib.error.URLError) as exc:
+        raise SetupError(f"Cannot reach Observatory at {endpoint}. Start the dashboard first, then retry.") from exc
+
+
+class _Connection:
+    def __init__(self, endpoint, service, integrations, capture_content, token, source_id, events):
+        self.client = Client(endpoint=endpoint, service=service, token=token, timeout=2)
+        self.payload = {"source_id": source_id, "service": service,
+                        "runtime": f"Python {platform.python_version()}",
+                        "integrations": integrations, "capture_content": capture_content,
+                        "process_id": os.getpid(), "state": "running"}
+        self.stop = threading.Event()
+        self.events = events
+        self.thread = threading.Thread(target=self._heartbeat, name="observatory-connection", daemon=True)
+
+    def start(self):
+        if not self.client._post("/api/connections", self.payload):
+            raise SetupError("Could not register this agent with Observatory. Check the receiver address and token.")
+        self.thread.start()
+
+    def _heartbeat(self):
+        next_heartbeat = time.monotonic() + 10
+        while not self.stop.is_set() or not self.events.empty():
+            try:
+                event = self.events.get(timeout=0.5)
+                self.client._post("/api/progress", event)
+            except queue.Empty:
+                pass
+            if time.monotonic() >= next_heartbeat:
+                self.client._post("/api/connections", self.payload)
+                next_heartbeat = time.monotonic() + 10
+
+    def close(self):
+        self.stop.set()
+        if self.thread.is_alive():
+            self.thread.join(timeout=3)
+        self.client._post("/api/connections", {**self.payload, "state": "stopped"})
+
+
+def run_agent(target, arguments=(), *, endpoint="http://127.0.0.1:4318", service="my-agent",
+              module=False, capture_content=False):
+    """Instrument installed SDKs, then run a script or module in this interpreter."""
+    endpoint = _receiver_endpoint(endpoint)
+    service = service.strip()
+    if not service or len(service) > 128:
+        raise SetupError("Choose a service name between 1 and 128 characters.")
+    path = Path(target).resolve() if not module else None
+    if path is not None and not path.is_file():
+        raise SetupError(f"Cannot find {target}. Run this command from your agent's folder or supply its full path.")
+    try:
+        from openinference.instrumentation import TraceConfig
+        from opentelemetry import trace
+        from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+        from opentelemetry.sdk.resources import Resource
+        from opentelemetry.sdk.trace import SpanLimits, TracerProvider
+        from opentelemetry.sdk.trace.export import BatchSpanProcessor
+        from opentelemetry.sdk.trace import SpanProcessor
+        from packaging.requirements import Requirement
+    except ImportError as exc:
+        raise SetupError('Automatic capture is not installed in this Python environment. Install this checkout with: python -m pip install ".[auto]"') from exc
+
+    token = os.environ.get("OBSERVATORY_TOKEN")
+    _check_receiver(endpoint, token)
+    headers = {"Authorization": f"Bearer {token}"} if token else {}
+    source_id = secrets.token_hex(16)
+    events = queue.Queue(maxsize=256)
+    provider = TracerProvider(resource=Resource.create({"service.name": service,
+        "host.name": socket.gethostname(), "host.arch": platform.machine(),
+        "os.type": platform.system().lower(), "observatory.capture": "automatic",
+        "observatory.source_id": source_id}),
+        span_limits=SpanLimits(max_attribute_length=4000))
+    provider.add_span_processor(BatchSpanProcessor(
+        OTLPSpanExporter(endpoint=endpoint + "/v1/traces", headers=headers, timeout=3),
+        schedule_delay_millis=500, max_export_batch_size=64))
+
+    class FrameworkLinks(SpanProcessor):
+        def on_start(self, span, parent_context=None):
+            # LangChain callbacks intentionally do not attach their spans to the
+            # Python runtime context. Link SDK spans to that active chain instead.
+            if span.parent is not None or importlib.util.find_spec("langchain_core") is None:
+                return
+            try:
+                from openinference.instrumentation.langchain import get_current_span
+                parent = get_current_span()
+                if parent is not None and parent.get_span_context().is_valid:
+                    span.add_link(parent.get_span_context(), {"observatory.link.reason": "langchain.runtime_context"})
+            except Exception:
+                pass
+
+    provider.add_span_processor(FrameworkLinks())
+
+    class LiveProgress(SpanProcessor):
+        def _send(self, span, stage):
+            scope = getattr(span, "instrumentation_scope", None)
+            if scope is None or scope.name not in {"openinference.instrumentation.openai", "openinference.instrumentation.anthropic"}:
+                return
+            event = {"trace_id": f"{span.context.trace_id:032x}", "stage": stage,
+                     "message": f"{service} · {span.name}", "service": service}
+            try:
+                events.put_nowait(event)
+            except queue.Full:
+                pass
+
+        def on_start(self, span, parent_context=None):
+            self._send(span, "Model call started")
+
+        def on_end(self, span):
+            self._send(span, "Model call finished")
+
+    provider.add_span_processor(LiveProgress())
+    trace.set_tracer_provider(provider)
+    config = TraceConfig(hide_inputs=not capture_content, hide_outputs=not capture_content,
+                         hide_input_images=True, hide_embeddings_vectors=True,
+                         hide_embeddings_text=not capture_content,
+                         hide_llm_invocation_parameters=not capture_content)
+    enabled, failures, instrumentors = [], [], []
+    connection = None
+    original_argv, original_path = sys.argv[:], sys.path[:]
+    try:
+        for sdk, name, class_name in INTEGRATIONS:
+            if importlib.util.find_spec(sdk) is None:
+                continue
+            try:
+                package = importlib.import_module(f"openinference.instrumentation.{name}")
+                instrumentor = getattr(package, class_name)()
+                for dependency in instrumentor.instrumentation_dependencies():
+                    required = Requirement(dependency)
+                    installed = importlib.metadata.version(required.name)
+                    if installed not in required.specifier:
+                        raise SetupError(f"requires {required}; found {installed}. Use a matching integration version or update that SDK.")
+                instrumentor.instrument(tracer_provider=provider, config=config)
+                if not instrumentor.is_instrumented_by_opentelemetry:
+                    raise SetupError("SDK and instrumentor versions do not match")
+                instrumentors.append(instrumentor)
+                enabled.append(name)
+            except Exception as exc:
+                failures.append(f"{name}: {exc}")
+        if not enabled:
+            hint = " ".join(failures) if failures else "No supported SDK found."
+            raise SetupError("Automatic capture needs the OpenAI SDK, Anthropic SDK, or LangChain in this same Python environment. " + hint)
+        if failures:
+            print("Observatory: some integrations could not start: " + ", ".join(failures), file=sys.stderr)
+        connection = _Connection(endpoint, service, enabled, capture_content, token, source_id, events)
+        connection.start()
+        print(f"Observatory is listening to {service}: {', '.join(enabled)}.\nDashboard: {endpoint}", file=sys.stderr)
+        sys.argv = [target, *arguments]
+        if module:
+            runpy.run_module(target, run_name="__main__", alter_sys=True)
+        else:
+            sys.path.insert(0, str(path.parent))
+            runpy.run_path(str(path), run_name="__main__")
+    finally:
+        # Finished streaming calls are exported before the connection is marked stopped.
+        provider.shutdown()
+        if connection is not None:
+            connection.close()
+        for instrumentor in reversed(instrumentors):
+            instrumentor.uninstrument()
+        sys.argv, sys.path = original_argv, original_path

observatory/client.py

@@ -0,0 +1,253 @@
+"""Small, dependency-free client for instrumenting Python agent code."""
+from __future__ import annotations
+
+import asyncio
+import contextvars
+import json
+import secrets
+import time
+import urllib.request
+from typing import Any
+
+_ACTIVE_RUN: contextvars.ContextVar["Run | None"] = contextvars.ContextVar("observatory_run", default=None)
+_ACTIVE_SPAN: contextvars.ContextVar["Span | None"] = contextvars.ContextVar("observatory_span", default=None)
+_AUTO_PARENT = object()
+MAX_REQUEST_BYTES = 8 * 1024 * 1024
+
+
+def _random_hex(byte_count):
+    value = secrets.token_hex(byte_count)
+    while not any(char != "0" for char in value):
+        value = secrets.token_hex(byte_count)
+    return value
+
+
+class Client:
+    """Emit framework-neutral spans and progress to an Observatory receiver."""
+
+    def __init__(self, endpoint="http://127.0.0.1:4318", service="python-agent",
+                 token=None, timeout=3.0, strict=False, capture_exception_message=False,
+                 resource_attributes=None):
+        self.endpoint = endpoint.rstrip("/")
+        self.service = str(service)
+        self.token = token
+        self.timeout = float(timeout)
+        self.strict = bool(strict)
+        self.capture_exception_message = bool(capture_exception_message)
+        self.resource_attributes = dict(resource_attributes or {})
+        self.last_error: Exception | None = None
+
+    def trace(self, name="agent.run", *, task=None, session_id=None, attributes=None):
+        return Run(self, name, task, session_id, attributes or {})
+
+    def _post_spans(self, spans):
+        base = {"service": self.service, "spans": []}
+        base_size = len(json.dumps(base, ensure_ascii=False, separators=(",", ":")).encode("utf-8"))
+        batch = []
+        batch_size = base_size
+        successful = True
+        last_error = None
+
+        def send():
+            nonlocal successful, last_error, batch, batch_size
+            if not batch:
+                return
+            if not self._post("/api/spans", {"service": self.service, "spans": batch}):
+                successful = False
+                last_error = self.last_error
+            batch = []
+            batch_size = base_size
+
+        for span in spans:
+            try:
+                item_size = len(json.dumps(span, ensure_ascii=False, separators=(",", ":")).encode("utf-8"))
+            except Exception as exc:
+                self.last_error = exc
+                if self.strict:
+                    raise
+                successful = False
+                last_error = exc
+                continue
+            if base_size + item_size > MAX_REQUEST_BYTES:
+                exc = ValueError("a span exceeds the receiver's 8 MiB request limit")
+                self.last_error = exc
+                if self.strict:
+                    raise exc
+                successful = False
+                last_error = exc
+                continue
+            separator = 1 if batch else 0
+            if batch and batch_size + separator + item_size > MAX_REQUEST_BYTES:
+                send()
+                separator = 0
+            batch.append(span)
+            batch_size += separator + item_size
+        send()
+        if not successful:
+            self.last_error = last_error
+        return successful
+
+    def _post(self, path, payload):
+        try:
+            data = json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
+            headers = {"Content-Type": "application/json"}
+            if self.token:
+                headers["Authorization"] = f"Bearer {self.token}"
+            request = urllib.request.Request(self.endpoint + path, data=data,
+                                             headers=headers, method="POST")
+            with urllib.request.urlopen(request, timeout=self.timeout) as response:
+                if not 200 <= response.status < 300:
+                    raise OSError(f"Observatory returned HTTP {response.status}")
+            return True
+        except Exception as exc:
+            self.last_error = exc
+            if self.strict:
+                raise
+            return False
+
+
+class Run:
+    def __init__(self, client, name, task, session_id, attributes):
+        self.client = client
+        self.trace_id = _random_hex(16)
+        self.session_id = str(session_id).strip() if session_id is not None else None
+        if self.session_id is not None and (not self.session_id or len(self.session_id) > 512):
+            raise ValueError("session_id must contain 1 to 512 characters")
+        self.spans: list[Span] = []
+        root_attributes = dict(attributes)
+        if task is not None:
+            root_attributes.setdefault("agent.task", str(task))
+        if self.session_id:
+            root_attributes.setdefault("session.id", self.session_id)
+        self.root = Span(self, name, root_attributes, parent_span_id=None)
+        self.spans.append(self.root)
+        self._run_token = None
+        self._closed = False
+
+    def __enter__(self):
+        if self._closed or self._run_token is not None:
+            raise RuntimeError("run cannot be entered more than once")
+        self._run_token = _ACTIVE_RUN.set(self)
+        self.root.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        self._finish(exc_type, exc, tb)
+        self._flush()
+        return False
+
+    async def __aenter__(self):
+        return self.__enter__()
+
+    async def __aexit__(self, exc_type, exc, tb):
+        self._finish(exc_type, exc, tb)
+        await asyncio.to_thread(self._flush)
+        return False
+
+    def _finish(self, exc_type, exc, tb):
+        if self._closed:
+            return
+        self.root.__exit__(exc_type, exc, tb)
+        if self._run_token is not None:
+            _ACTIVE_RUN.reset(self._run_token)
+            self._run_token = None
+        self._closed = True
+
+    def span(self, name, *, attributes=None, parent_span_id=_AUTO_PARENT):
+        if self._closed:
+            raise RuntimeError("run is already closed")
+        if parent_span_id is _AUTO_PARENT:
+            active = _ACTIVE_SPAN.get()
+            parent_span_id = active.span_id if active is not None and active.run is self else self.root.span_id
+        span = Span(self, name, attributes or {}, parent_span_id=parent_span_id)
+        self.spans.append(span)
+        return span
+
+    def progress(self, stage, message="", percent=None):
+        if self._closed:
+            raise RuntimeError("run is already closed")
+        payload = {"trace_id": self.trace_id, "stage": str(stage), "message": str(message), "service": self.client.service}
+        if percent is not None:
+            payload["percent"] = percent
+        if self.session_id:
+            payload["session_id"] = self.session_id
+        return self.client._post("/api/progress", payload)
+
+    async def aprogress(self, stage, message="", percent=None):
+        return await asyncio.to_thread(self.progress, stage, message, percent)
+
+    def _flush(self):
+        payload = []
+        for span in self.spans:
+            payload.append({"trace_id": self.trace_id, "span_id": span.span_id,
+                            "parent_span_id": span.parent_span_id, "name": span.name,
+                            "start_unix_nano": span.start_ns, "end_unix_nano": span.end_ns,
+                            "status": span.status, "attributes": span.attributes,
+                            "events": span.events,
+                            "resource_attributes": self.client.resource_attributes})
+        return self.client._post_spans(payload)
+
+
+class Span:
+    def __init__(self, run, name, attributes, parent_span_id):
+        self.run = run
+        self.name = str(name)
+        self.span_id = _random_hex(8)
+        self.parent_span_id = parent_span_id
+        self.start_ns = time.time_ns()
+        self.end_ns = self.start_ns
+        self.status = "UNSET"
+        self.attributes: dict[str, Any] = dict(attributes)
+        self.events: list[dict[str, Any]] = []
+        self._context_token = None
+        self._finished = False
+
+    def __enter__(self):
+        if self._context_token is not None or self._finished:
+            raise RuntimeError("span cannot be entered more than once")
+        self._context_token = _ACTIVE_SPAN.set(self)
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        self.finish(exc)
+        if self._context_token is not None:
+            _ACTIVE_SPAN.reset(self._context_token)
+            self._context_token = None
+        return False
+
+    def finish(self, error=None):
+        if self._finished:
+            return self
+        if error is not None:
+            self.status = "ERROR"
+            self.attributes.setdefault("exception.type", type(error).__name__)
+            exception = {"type": type(error).__name__}
+            if self.run.client.capture_exception_message:
+                message = str(error)
+                self.attributes.setdefault("exception.message", message)
+                exception["message"] = message
+            self.add_event("exception", exception)
+        self.end_ns = max(time.time_ns(), self.start_ns)
+        self._finished = True
+        return self
+
+    async def __aenter__(self):
+        return self.__enter__()
+
+    async def __aexit__(self, exc_type, exc, tb):
+        return self.__exit__(exc_type, exc, tb)
+
+    def set_attribute(self, key, value):
+        self.attributes[str(key)] = value
+        return self
+
+    def add_event(self, name, attributes=None):
+        self.events.append({"name": str(name), "time_unix_nano": time.time_ns(),
+                            "attributes": dict(attributes or {})})
+        return self
+
+    def set_status(self, status):
+        if status not in {"UNSET", "OK", "ERROR"}:
+            raise ValueError("status must be UNSET, OK, or ERROR")
+        self.status = status
+        return self