agentos-python 0.1.0__py3-none-any.whl

agentos/tracing.py

@@ -0,0 +1,168 @@
+"""Trace context management — ID generation, context propagation, spans.
+
+Uses ``contextvars`` for async-safe, thread-safe context propagation.
+Supports nested traces and spans via context managers.
+"""
+
+from __future__ import annotations
+
+import contextvars
+import os
+import uuid
+from collections.abc import Generator
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any
+
+# ---------------------------------------------------------------------------
+# ID generators
+# ---------------------------------------------------------------------------
+
+
+def generate_trace_id() -> str:
+    """Generate a W3C-compatible trace ID (32 lowercase hex chars)."""
+    return uuid.uuid4().hex
+
+
+def generate_span_id() -> str:
+    """Generate a W3C-compatible span ID (16 lowercase hex chars)."""
+    return os.urandom(8).hex()
+
+
+def generate_event_id() -> str:
+    """Generate a UUID string for event_id (v7 if available, else v4)."""
+    try:
+        return str(uuid.uuid7())
+    except AttributeError:
+        return str(uuid.uuid4())
+
+
+def utcnow_iso() -> str:
+    """Return current UTC time as ISO 8601 string with timezone."""
+    return datetime.now(timezone.utc).isoformat()
+
+
+# ---------------------------------------------------------------------------
+# Trace context
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class TraceContext:
+    """Holds the current trace/span identity and metadata.
+
+    Passed through ``contextvars`` so nested ``@observe`` / ``with span()``
+    calls automatically inherit the parent context.
+    """
+
+    trace_id: str
+    span_id: str
+    agent_id: str
+    parent_span_id: str | None = None
+    task_run_id: str | None = None
+    user_id: str | None = None
+    session_id: str | None = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    # Mutable observation state (updated via trace_context helper)
+    _observation_updates: dict[str, Any] = field(default_factory=dict, repr=False)
+    _trace_updates: dict[str, Any] = field(default_factory=dict, repr=False)
+    _scores: list[dict[str, Any]] = field(default_factory=list, repr=False)
+
+
+# The global context variable holding the current trace context.
+_current_context: contextvars.ContextVar[TraceContext | None] = contextvars.ContextVar(
+    "agentos_trace_context", default=None
+)
+
+
+def get_current_context() -> TraceContext | None:
+    """Return the current trace context, or None if not inside a trace."""
+    return _current_context.get()
+
+
+def _set_context(ctx: TraceContext | None) -> contextvars.Token[TraceContext | None]:
+    """Set the current trace context. Returns a token for restoring."""
+    return _current_context.set(ctx)
+
+
+# ---------------------------------------------------------------------------
+# Context managers
+# ---------------------------------------------------------------------------
+
+
+@contextmanager
+def trace(
+    agent_id: str,
+    *,
+    trace_id: str | None = None,
+    task_run_id: str | None = None,
+    user_id: str | None = None,
+    session_id: str | None = None,
+) -> Generator[TraceContext, None, None]:
+    """Start a new trace. Creates a root span.
+
+    Usage::
+
+        with trace(agent_id="my-agent") as ctx:
+            # ctx.trace_id, ctx.span_id available
+            do_work()
+    """
+    ctx = TraceContext(
+        trace_id=trace_id or generate_trace_id(),
+        span_id=generate_span_id(),
+        agent_id=agent_id,
+        task_run_id=task_run_id,
+        user_id=user_id,
+        session_id=session_id,
+    )
+    token = _set_context(ctx)
+    try:
+        yield ctx
+    finally:
+        _current_context.reset(token)
+
+
+@contextmanager
+def span(
+    *,
+    agent_id: str | None = None,
+) -> Generator[TraceContext, None, None]:
+    """Start a child span within the current trace.
+
+    Inherits trace_id, agent_id, user_id, session_id, task_run_id from parent.
+    Sets parent_span_id to the parent's span_id.
+
+    Usage::
+
+        with trace(agent_id="bot") as t:
+            with span() as s:
+                # s.parent_span_id == t.span_id
+                do_step()
+    """
+    parent = _current_context.get()
+    if parent is None:
+        # No parent trace — create a standalone span with a new trace
+        effective_agent_id = agent_id or "unknown"
+        ctx = TraceContext(
+            trace_id=generate_trace_id(),
+            span_id=generate_span_id(),
+            agent_id=effective_agent_id,
+        )
+    else:
+        ctx = TraceContext(
+            trace_id=parent.trace_id,
+            span_id=generate_span_id(),
+            agent_id=agent_id or parent.agent_id,
+            parent_span_id=parent.span_id,
+            task_run_id=parent.task_run_id,
+            user_id=parent.user_id,
+            session_id=parent.session_id,
+        )
+
+    token = _set_context(ctx)
+    try:
+        yield ctx
+    finally:
+        _current_context.reset(token)

agentos/transport.py

@@ -0,0 +1,285 @@
+"""HTTP transport with batch queue and retry logic.
+
+All events flow through the transport layer:
+1. ``enqueue()`` adds to an in-memory queue
+2. Background thread flushes periodically (sync) or asyncio task (async)
+3. ``_send_batch()`` posts to ``/v1/events/batch`` with retry + backoff
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+import time
+from typing import Any
+
+import httpx
+
+from agentos.config import AgentOSConfig
+
+logger = logging.getLogger("agentos.transport")
+
+SDK_NAME = "agentos-python"
+
+
+class Transport:
+    """Synchronous transport — uses httpx.Client and a background flush thread."""
+
+    def __init__(self, config: AgentOSConfig, version: str = "0.1.0") -> None:
+        self._config = config
+        self._version = version
+        self._queue: list[dict[str, Any]] = []
+        self._lock = threading.Lock()
+        self._closed = False
+
+        self._client = httpx.Client(
+            base_url=config.base_url,
+            headers={
+                "X-API-Key": config.api_key,
+                "Content-Type": "application/json",
+                "User-Agent": f"{SDK_NAME}/{version}",
+            },
+            timeout=httpx.Timeout(30.0, connect=10.0),
+        )
+
+        self._flush_thread: threading.Thread | None = None
+        if config.flush_interval > 0:
+            self._flush_thread = threading.Thread(
+                target=self._periodic_flush, daemon=True, name="agentos-flush"
+            )
+            self._flush_thread.start()
+
+    def enqueue(self, event: dict[str, Any]) -> None:
+        """Add an event to the queue. Auto-flushes when batch_size is reached."""
+        if self._closed:
+            return
+
+        with self._lock:
+            if len(self._queue) >= self._config.max_queue_size:
+                if self._config.debug:
+                    logger.warning("Queue full (%d), dropping oldest event", len(self._queue))
+                self._queue.pop(0)
+            self._queue.append(event)
+            should_flush = len(self._queue) >= self._config.batch_size
+
+        if should_flush:
+            self.flush()
+
+    def flush(self) -> None:
+        """Send all queued events immediately."""
+        with self._lock:
+            if not self._queue:
+                return
+            batch = self._queue[:]
+            self._queue.clear()
+
+        # Send in chunks of 500 (server limit)
+        for i in range(0, len(batch), 500):
+            chunk = batch[i : i + 500]
+            try:
+                self._send_batch(chunk)
+            except Exception:
+                logger.exception("Failed to send batch of %d events", len(chunk))
+
+    def _send_batch(self, events: list[dict[str, Any]]) -> None:
+        """POST events to /v1/events/batch with retry and exponential backoff."""
+        payload = {"events": events}
+        backoff = 1.0
+
+        for attempt in range(self._config.max_retries + 1):
+            try:
+                response = self._client.post("/v1/events/batch", json=payload)
+
+                if response.status_code == 200:
+                    if self._config.debug:
+                        logger.debug("Sent batch of %d events", len(events))
+                    return
+
+                if response.status_code == 429:
+                    retry_after = response.headers.get("Retry-After")
+                    wait = float(retry_after) if retry_after else backoff
+                    if self._config.debug:
+                        logger.debug("Rate limited, waiting %.1fs", wait)
+                    time.sleep(wait)
+                    backoff = min(backoff * 2, 30.0)
+                    continue
+
+                if response.status_code >= 500:
+                    if attempt < self._config.max_retries:
+                        time.sleep(backoff)
+                        backoff = min(backoff * 2, 30.0)
+                        continue
+                    logger.error(
+                        "Server error %d after %d retries",
+                        response.status_code,
+                        self._config.max_retries,
+                    )
+                    return
+
+                # 4xx (not 429) — don't retry
+                if self._config.debug:
+                    logger.warning("Client error %d: %s", response.status_code, response.text[:200])
+                return
+
+            except httpx.TransportError as exc:
+                if attempt < self._config.max_retries:
+                    if self._config.debug:
+                        logger.debug("Transport error: %s, retrying in %.1fs", exc, backoff)
+                    time.sleep(backoff)
+                    backoff = min(backoff * 2, 30.0)
+                    continue
+                logger.error("Transport error after %d retries: %s", self._config.max_retries, exc)
+
+    def _periodic_flush(self) -> None:
+        """Background thread that flushes the queue on an interval."""
+        while not self._closed:
+            time.sleep(self._config.flush_interval)
+            if self._closed:
+                break
+            try:
+                self.flush()
+            except Exception:
+                logger.exception("Error in periodic flush")
+
+    def shutdown(self) -> None:
+        """Flush remaining events and close the HTTP client."""
+        self._closed = True
+        try:
+            self.flush()
+        except Exception:
+            logger.exception("Error flushing during shutdown")
+        finally:
+            self._client.close()
+
+    @property
+    def queue_size(self) -> int:
+        """Number of events currently queued."""
+        with self._lock:
+            return len(self._queue)
+
+
+class AsyncTransport:
+    """Asynchronous transport — uses httpx.AsyncClient and asyncio for flush."""
+
+    def __init__(self, config: AgentOSConfig, version: str = "0.1.0") -> None:
+        self._config = config
+        self._version = version
+        self._queue: list[dict[str, Any]] = []
+        self._closed = False
+        self._flush_task: Any = None  # asyncio.Task, typed as Any for import-free
+
+        self._client = httpx.AsyncClient(
+            base_url=config.base_url,
+            headers={
+                "X-API-Key": config.api_key,
+                "Content-Type": "application/json",
+                "User-Agent": f"{SDK_NAME}/{version}",
+            },
+            timeout=httpx.Timeout(30.0, connect=10.0),
+        )
+
+    async def start(self) -> None:
+        """Start the periodic flush task."""
+        import asyncio
+
+        if self._config.flush_interval > 0 and self._flush_task is None:
+            self._flush_task = asyncio.create_task(self._periodic_flush())
+
+    async def enqueue(self, event: dict[str, Any]) -> None:
+        """Add an event to the queue."""
+        if self._closed:
+            return
+
+        if len(self._queue) >= self._config.max_queue_size:
+            if self._config.debug:
+                logger.warning("Queue full (%d), dropping oldest event", len(self._queue))
+            self._queue.pop(0)
+        self._queue.append(event)
+
+        if len(self._queue) >= self._config.batch_size:
+            await self.flush()
+
+    async def flush(self) -> None:
+        """Send all queued events immediately."""
+        if not self._queue:
+            return
+        batch = self._queue[:]
+        self._queue.clear()
+
+        for i in range(0, len(batch), 500):
+            chunk = batch[i : i + 500]
+            try:
+                await self._send_batch(chunk)
+            except Exception:
+                logger.exception("Failed to send batch of %d events", len(chunk))
+
+    async def _send_batch(self, events: list[dict[str, Any]]) -> None:
+        """POST events with retry and exponential backoff."""
+        import asyncio
+
+        payload = {"events": events}
+        backoff = 1.0
+
+        for attempt in range(self._config.max_retries + 1):
+            try:
+                response = await self._client.post("/v1/events/batch", json=payload)
+
+                if response.status_code == 200:
+                    if self._config.debug:
+                        logger.debug("Sent batch of %d events", len(events))
+                    return
+
+                if response.status_code == 429:
+                    retry_after = response.headers.get("Retry-After")
+                    wait = float(retry_after) if retry_after else backoff
+                    await asyncio.sleep(wait)
+                    backoff = min(backoff * 2, 30.0)
+                    continue
+
+                if response.status_code >= 500:
+                    if attempt < self._config.max_retries:
+                        await asyncio.sleep(backoff)
+                        backoff = min(backoff * 2, 30.0)
+                        continue
+                    logger.error("Server error %d after retries", response.status_code)
+                    return
+
+                if self._config.debug:
+                    logger.warning("Client error %d: %s", response.status_code, response.text[:200])
+                return
+
+            except httpx.TransportError as exc:
+                if attempt < self._config.max_retries:
+                    await asyncio.sleep(backoff)
+                    backoff = min(backoff * 2, 30.0)
+                    continue
+                logger.error("Transport error after retries: %s", exc)
+
+    async def _periodic_flush(self) -> None:
+        """Periodic flush task."""
+        import asyncio
+
+        while not self._closed:
+            await asyncio.sleep(self._config.flush_interval)
+            if self._closed:
+                break
+            try:
+                await self.flush()
+            except Exception:
+                logger.exception("Error in periodic flush")
+
+    async def shutdown(self) -> None:
+        """Flush remaining events and close the client."""
+        self._closed = True
+        if self._flush_task is not None:
+            self._flush_task.cancel()
+        try:
+            await self.flush()
+        except Exception:
+            logger.exception("Error flushing during shutdown")
+        finally:
+            await self._client.aclose()
+
+    @property
+    def queue_size(self) -> int:
+        return len(self._queue)

agentos/types.py

@@ -0,0 +1,101 @@
+"""Type definitions for the Agent OS SDK.
+
+Mirrors the server's event schema at shared/schema/events.py.
+Uses string enums and TypedDicts for lightweight, dependency-free types.
+"""
+
+from __future__ import annotations
+
+import enum
+from typing import Any
+
+# ---------------------------------------------------------------------------
+# Enums — match server values exactly
+# ---------------------------------------------------------------------------
+
+
+class EventType(str, enum.Enum):
+    """All supported event types in the Agent OS event stream."""
+
+    LLM_CALL = "agent.llm_call"
+    TOOL_CALL = "agent.tool_call"
+    AGENT_HANDOFF = "agent.handoff"
+    RETRIEVAL_QUERY = "agent.retrieval_query"
+    EVAL = "agent.eval"
+    SECURITY_ALERT = "agent.security_alert"
+    FLAG_CHECK = "agent.flag_check"
+    BUSINESS_EVENT = "agent.business_event"
+    ROUTING_DECISION = "agent.routing_decision"
+    ROUTING_OUTCOME = "agent.routing_outcome"
+
+
+class FinishReason(str, enum.Enum):
+    STOP = "stop"
+    LENGTH = "length"
+    TOOL_CALLS = "tool_calls"
+    CONTENT_FILTER = "content_filter"
+
+
+class ToolStatus(str, enum.Enum):
+    SUCCESS = "success"
+    ERROR = "error"
+    TIMEOUT = "timeout"
+
+
+class SecurityAlertType(str, enum.Enum):
+    PROMPT_INJECTION = "prompt_injection"
+    PII_DETECTED = "pii_detected"
+    JAILBREAK = "jailbreak"
+    TOXIC_CONTENT = "toxic_content"
+    DATA_EXFILTRATION = "data_exfiltration"
+    POLICY_VIOLATION = "policy_violation"
+
+
+class Severity(str, enum.Enum):
+    CRITICAL = "critical"
+    HIGH = "high"
+    MEDIUM = "medium"
+    LOW = "low"
+    INFO = "info"
+
+
+class SecurityAction(str, enum.Enum):
+    BLOCKED = "blocked"
+    FLAGGED = "flagged"
+    ALLOWED = "allowed"
+    REDACTED = "redacted"
+
+
+class EvalEvaluator(str, enum.Enum):
+    HUMAN = "human"
+    LLM = "llm"
+    HEURISTIC = "heuristic"
+    CODE = "code"
+
+
+# ---------------------------------------------------------------------------
+# TypedDict-style types (using plain dicts for Python 3.9 compat)
+# ---------------------------------------------------------------------------
+
+# These are documentation types — the SDK uses plain dicts internally
+# to avoid requiring TypedDict backport on 3.9/3.10.
+
+LLMMessage = dict[str, Any]
+"""{"role": str, "content": str | None, "tool_calls": list | None}"""
+
+ErrorInfo = dict[str, str]
+"""{"type": str, "message": str}"""
+
+CostInfo = dict[str, float | str]
+"""{"input_cost": float, "output_cost": float, "total_cost": float, "currency": str}"""
+
+RetrievalDocument = dict[str, Any]
+"""{"id": str, "score": float | None, "content": str | None, "metadata": dict | None}"""
+
+EventPayload = dict[str, Any]
+"""The full event dict sent to the ingestion API (EventInput wire format).
+
+Required keys: event_type, trace_id, span_id, agent_id
+Optional keys: event_id, timestamp, parent_span_id, task_run_id, user_id,
+    session_id, environment, sdk_name, sdk_version, properties
+"""