agentos-python 0.1.0__py3-none-any.whl

agentos/__init__.py

@@ -0,0 +1,58 @@
+"""Agent OS Python SDK — AI agent observability platform.
+
+Usage::
+
+    from agentos import AgentOS
+
+    client = AgentOS(api_key="aos_...")
+    client.llm_call("my-agent", model="gpt-4o", system="openai")
+    client.flush()
+"""
+
+from agentos.client import AgentOS, AsyncAgentOS, get_client
+from agentos.config import AgentOSConfig
+from agentos.decorators import observe, trace_agent, trace_context, trace_tool
+from agentos.integrations.anthropic import wrap_anthropic
+from agentos.integrations.openai import wrap_openai
+from agentos.tracing import TraceContext, span, trace
+from agentos.types import (
+    ErrorInfo,
+    EvalEvaluator,
+    EventPayload,
+    EventType,
+    FinishReason,
+    LLMMessage,
+    SecurityAction,
+    SecurityAlertType,
+    Severity,
+    ToolStatus,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "__version__",
+    "AgentOS",
+    "AgentOSConfig",
+    "AsyncAgentOS",
+    "ErrorInfo",
+    "EvalEvaluator",
+    "EventPayload",
+    "EventType",
+    "FinishReason",
+    "LLMMessage",
+    "SecurityAction",
+    "SecurityAlertType",
+    "Severity",
+    "ToolStatus",
+    "TraceContext",
+    "get_client",
+    "observe",
+    "span",
+    "trace",
+    "trace_agent",
+    "trace_context",
+    "trace_tool",
+    "wrap_anthropic",
+    "wrap_openai",
+]

agentos/client.py

@@ -0,0 +1,273 @@
+"""AgentOS client — the main public API for the SDK.
+
+Usage::
+
+    from agentos import AgentOS
+
+    client = AgentOS(api_key="aos_...")
+    client.llm_call("my-agent", model="gpt-4o", system="openai", input_tokens=100)
+    client.flush()
+"""
+
+from __future__ import annotations
+
+import atexit
+import logging
+import os
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import Any
+
+from agentos import events as _events
+from agentos.config import AgentOSConfig
+from agentos.tracing import TraceContext
+from agentos.tracing import span as _span
+from agentos.tracing import trace as _trace
+from agentos.transport import AsyncTransport, Transport
+
+logger = logging.getLogger("agentos")
+
+SDK_NAME = "agentos-python"
+SDK_VERSION = "0.1.0"
+
+
+class AgentOS:
+    """Synchronous Agent OS client.
+
+    All capture methods are non-blocking — events are queued and sent
+    in the background. Call ``flush()`` to ensure delivery.
+    """
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        base_url: str = "https://api.agentos.dev",
+        **kwargs: Any,
+    ) -> None:
+        resolved_key = api_key or os.environ.get("AGENTOS_API_KEY", "")
+        self._config = AgentOSConfig(api_key=resolved_key, base_url=base_url, **kwargs)
+        self._transport: Transport | None = None
+
+        if self._config.enabled and self._config.api_key:
+            self._transport = Transport(self._config, version=SDK_VERSION)
+            atexit.register(self.shutdown)
+
+    def _enqueue(self, event: dict[str, Any]) -> None:
+        """Stamp SDK metadata and enqueue for delivery."""
+        if self._transport is None:
+            return
+        event["sdk_name"] = SDK_NAME
+        event["sdk_version"] = SDK_VERSION
+        if self._config.environment:
+            event.setdefault("environment", self._config.environment)
+        try:
+            self._transport.enqueue(event)
+        except Exception:
+            logger.exception("Failed to enqueue event")
+
+    # --- Convenience methods (one per event type) ---
+
+    def capture(self, event: dict[str, Any]) -> None:
+        """Enqueue a raw event dict."""
+        self._enqueue(event)
+
+    def llm_call(self, agent_id: str, *, model: str, system: str, **kwargs: Any) -> None:
+        """Capture an LLM call event."""
+        if not self._config.capture_content:
+            kwargs.pop("input", None)
+            kwargs.pop("output", None)
+        self._enqueue(_events.llm_call(agent_id, model=model, system=system, **kwargs))
+
+    def tool_call(self, agent_id: str, *, tool_name: str, **kwargs: Any) -> None:
+        """Capture a tool call event."""
+        if not self._config.capture_content:
+            kwargs.pop("input", None)
+            kwargs.pop("output", None)
+        self._enqueue(_events.tool_call(agent_id, tool_name=tool_name, **kwargs))
+
+    def handoff(self, agent_id: str, *, target_agent_id: str, **kwargs: Any) -> None:
+        """Capture an agent handoff event."""
+        self._enqueue(_events.handoff(agent_id, target_agent_id=target_agent_id, **kwargs))
+
+    def retrieval_query(self, agent_id: str, *, source: str, **kwargs: Any) -> None:
+        """Capture a retrieval query event."""
+        self._enqueue(_events.retrieval_query(agent_id, source=source, **kwargs))
+
+    def eval(self, agent_id: str, *, eval_name: str, score: float, **kwargs: Any) -> None:
+        """Capture an evaluation event."""
+        self._enqueue(_events.eval_result(agent_id, eval_name=eval_name, score=score, **kwargs))
+
+    def security_alert(
+        self, agent_id: str, *, alert_type: str, severity: str, **kwargs: Any
+    ) -> None:
+        """Capture a security alert event."""
+        self._enqueue(
+            _events.security_alert(agent_id, alert_type=alert_type, severity=severity, **kwargs)
+        )
+
+    def flag_check(self, agent_id: str, *, flag_key: str, flag_value: Any, **kwargs: Any) -> None:
+        """Capture a flag check event."""
+        self._enqueue(
+            _events.flag_check(agent_id, flag_key=flag_key, flag_value=flag_value, **kwargs)
+        )
+
+    def business_event(self, agent_id: str, *, event_name: str, **kwargs: Any) -> None:
+        """Capture a business event."""
+        self._enqueue(_events.business_event(agent_id, event_name=event_name, **kwargs))
+
+    # --- Trace context ---
+
+    @contextmanager
+    def trace(self, agent_id: str, **kwargs: Any) -> Generator[TraceContext, None, None]:
+        """Start a trace. Yields a TraceContext."""
+        with _trace(agent_id, **kwargs) as ctx:
+            yield ctx
+
+    @contextmanager
+    def span(self, **kwargs: Any) -> Generator[TraceContext, None, None]:
+        """Start a child span within the current trace."""
+        with _span(**kwargs) as ctx:
+            yield ctx
+
+    # --- Lifecycle ---
+
+    def flush(self) -> None:
+        """Block until all queued events are sent."""
+        if self._transport:
+            self._transport.flush()
+
+    def shutdown(self) -> None:
+        """Flush and close the client. Called automatically at exit."""
+        if self._transport:
+            self._transport.shutdown()
+            self._transport = None
+
+    @property
+    def enabled(self) -> bool:
+        return self._config.enabled and self._transport is not None
+
+    @property
+    def queue_size(self) -> int:
+        return self._transport.queue_size if self._transport else 0
+
+
+class AsyncAgentOS:
+    """Asynchronous Agent OS client.
+
+    Usage::
+
+        client = AsyncAgentOS(api_key="aos_...")
+        await client.llm_call("my-agent", model="gpt-4o", system="openai")
+        await client.flush()
+        await client.shutdown()
+    """
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        base_url: str = "https://api.agentos.dev",
+        **kwargs: Any,
+    ) -> None:
+        resolved_key = api_key or os.environ.get("AGENTOS_API_KEY", "")
+        self._config = AgentOSConfig(api_key=resolved_key, base_url=base_url, **kwargs)
+        self._transport: AsyncTransport | None = None
+
+        if self._config.enabled and self._config.api_key:
+            self._transport = AsyncTransport(self._config, version=SDK_VERSION)
+
+    async def _enqueue(self, event: dict[str, Any]) -> None:
+        if self._transport is None:
+            return
+        event["sdk_name"] = SDK_NAME
+        event["sdk_version"] = SDK_VERSION
+        if self._config.environment:
+            event.setdefault("environment", self._config.environment)
+        try:
+            await self._transport.enqueue(event)
+        except Exception:
+            logger.exception("Failed to enqueue event")
+
+    async def capture(self, event: dict[str, Any]) -> None:
+        await self._enqueue(event)
+
+    async def llm_call(self, agent_id: str, *, model: str, system: str, **kwargs: Any) -> None:
+        if not self._config.capture_content:
+            kwargs.pop("input", None)
+            kwargs.pop("output", None)
+        await self._enqueue(_events.llm_call(agent_id, model=model, system=system, **kwargs))
+
+    async def tool_call(self, agent_id: str, *, tool_name: str, **kwargs: Any) -> None:
+        if not self._config.capture_content:
+            kwargs.pop("input", None)
+            kwargs.pop("output", None)
+        await self._enqueue(_events.tool_call(agent_id, tool_name=tool_name, **kwargs))
+
+    async def handoff(self, agent_id: str, *, target_agent_id: str, **kwargs: Any) -> None:
+        await self._enqueue(_events.handoff(agent_id, target_agent_id=target_agent_id, **kwargs))
+
+    async def retrieval_query(self, agent_id: str, *, source: str, **kwargs: Any) -> None:
+        await self._enqueue(_events.retrieval_query(agent_id, source=source, **kwargs))
+
+    async def eval(self, agent_id: str, *, eval_name: str, score: float, **kwargs: Any) -> None:
+        await self._enqueue(
+            _events.eval_result(agent_id, eval_name=eval_name, score=score, **kwargs)
+        )
+
+    async def security_alert(
+        self, agent_id: str, *, alert_type: str, severity: str, **kwargs: Any
+    ) -> None:
+        await self._enqueue(
+            _events.security_alert(agent_id, alert_type=alert_type, severity=severity, **kwargs)
+        )
+
+    async def flag_check(
+        self, agent_id: str, *, flag_key: str, flag_value: Any, **kwargs: Any
+    ) -> None:
+        await self._enqueue(
+            _events.flag_check(agent_id, flag_key=flag_key, flag_value=flag_value, **kwargs)
+        )
+
+    async def business_event(self, agent_id: str, *, event_name: str, **kwargs: Any) -> None:
+        await self._enqueue(_events.business_event(agent_id, event_name=event_name, **kwargs))
+
+    async def flush(self) -> None:
+        if self._transport:
+            await self._transport.flush()
+
+    async def shutdown(self) -> None:
+        if self._transport:
+            await self._transport.shutdown()
+            self._transport = None
+
+    @property
+    def enabled(self) -> bool:
+        return self._config.enabled and self._transport is not None
+
+    @property
+    def queue_size(self) -> int:
+        return self._transport.queue_size if self._transport else 0
+
+
+# ---------------------------------------------------------------------------
+# Global singleton (env-var auto-init)
+# ---------------------------------------------------------------------------
+
+_global_client: AgentOS | None = None
+
+
+def get_client() -> AgentOS | None:
+    """Return the global client, auto-initializing from env vars if needed.
+
+    Set ``AGENTOS_API_KEY`` to enable auto-init. Returns None if not configured.
+    """
+    global _global_client
+    if _global_client is None:
+        api_key = os.environ.get("AGENTOS_API_KEY")
+        if api_key:
+            _global_client = AgentOS(
+                api_key=api_key,
+                base_url=os.environ.get("AGENTOS_BASE_URL", "https://api.agentos.dev"),
+                enabled=os.environ.get("AGENTOS_ENABLED", "true").lower() == "true",
+                environment=os.environ.get("AGENTOS_ENVIRONMENT"),
+            )
+    return _global_client

agentos/compat/__init__.py

@@ -0,0 +1,4 @@
+"""Migration compatibility layers for other observability platforms.
+
+Drop-in shims that let users migrate with minimal code changes.
+"""

agentos/compat/braintrust.py

@@ -0,0 +1,97 @@
+"""Braintrust-compatible API shim.
+
+Provides ``init``, ``@traced``, ``wrap_openai``, and ``Eval``.
+
+Usage::
+
+    from agentos.compat.braintrust import init, traced, Eval
+
+    experiment = init("my-project", api_key="aos_...")
+
+    @traced
+    def my_pipeline(input: str) -> str: ...
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+from agentos.client import AgentOS, get_client
+from agentos.decorators import observe
+from agentos.integrations.openai import wrap_openai as _native_wrap_openai
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def init(
+    project: str,
+    *,
+    api_key: str | None = None,
+    **kwargs: Any,
+) -> AgentOS:
+    """Initialize a Braintrust-compatible experiment.
+
+    Returns an AgentOS client configured as a global singleton.
+    """
+    import agentos.client
+
+    client = AgentOS(api_key=api_key or "", **kwargs)
+    agentos.client._global_client = client
+    return client
+
+
+def traced(fn: F) -> F:
+    """Braintrust-compatible ``@traced`` decorator."""
+    return observe()(fn)
+
+
+def wrap_openai(openai_client: Any, **kwargs: Any) -> Any:
+    """Braintrust-compatible ``wrap_openai``."""
+    return _native_wrap_openai(openai_client, **kwargs)
+
+
+class Eval:
+    """Braintrust-compatible eval runner.
+
+    Usage::
+
+        Eval(
+            "project",
+            data=lambda: [{"input": "2+2?", "expected": "4"}],
+            task=lambda input: call_llm(input),
+            scores=[lambda output, expected: float(output == expected)],
+        )
+    """
+
+    def __init__(
+        self,
+        project: str,
+        *,
+        data: Callable[[], list[dict[str, Any]]],
+        task: Callable[..., Any],
+        scores: list[Callable[..., float]],
+        api_key: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        client = get_client()
+        if client is None and api_key:
+            client = AgentOS(api_key=api_key, flush_interval=0)
+
+        dataset = data()
+        for item in dataset:
+            input_val = item.get("input")
+            expected = item.get("expected")
+            output = task(input_val)
+
+            for scorer in scores:
+                score = scorer(output, expected)
+                if client:
+                    client.eval(
+                        project,
+                        eval_name=scorer.__name__,
+                        score=float(score),
+                    )
+
+        if client:
+            client.flush()

agentos/compat/helicone.py

@@ -0,0 +1,72 @@
+"""Helicone-compatible header translator and proxy config.
+
+Usage::
+
+    from agentos.compat.helicone import get_proxy_config, translate_headers
+
+    config = get_proxy_config(api_key="aos_...", agent_id="my-agent")
+    client = openai.OpenAI(base_url=config["base_url"], default_headers=config["headers"])
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def get_proxy_config(
+    *,
+    api_key: str,
+    agent_id: str = "default",
+    base_url: str = "https://proxy.agentos.dev/v1",
+) -> dict[str, Any]:
+    """Return proxy config dict matching Helicone's usage pattern.
+
+    Returns:
+        Dict with ``base_url`` and ``headers`` keys, ready to pass to OpenAI client.
+    """
+    return {
+        "base_url": base_url,
+        "headers": {
+            "X-API-Key": api_key,
+            "X-Agent-Id": agent_id,
+        },
+    }
+
+
+def translate_headers(
+    helicone_headers: dict[str, str],
+    *,
+    api_key: str | None = None,
+) -> dict[str, str]:
+    """Translate Helicone-style headers to AgentOS headers.
+
+    Mapping:
+    - ``Helicone-Auth`` → ``X-API-Key`` (extracts Bearer token)
+    - ``Helicone-Session-Id`` → ``X-Session-Id``
+    - ``Helicone-Property-agent`` → ``X-Agent-Id``
+    - ``Helicone-Property-*`` → ``X-Property-*``
+    """
+    result: dict[str, str] = {}
+
+    for key, value in helicone_headers.items():
+        lower = key.lower()
+
+        if lower == "helicone-auth":
+            # Extract Bearer token
+            token = value.replace("Bearer ", "").replace("bearer ", "").strip()
+            result["X-API-Key"] = api_key or token
+
+        elif lower == "helicone-session-id":
+            result["X-Session-Id"] = value
+
+        elif lower.startswith("helicone-property-"):
+            prop_name = key[len("Helicone-Property-") :]
+            if prop_name.lower() == "agent":
+                result["X-Agent-Id"] = value
+            else:
+                result[f"X-Property-{prop_name}"] = value
+
+    if api_key and "X-API-Key" not in result:
+        result["X-API-Key"] = api_key
+
+    return result