agentos-python 0.1.0__py3-none-any.whl

agentos/events.py

@@ -0,0 +1,449 @@
+"""Event builder functions — one per event type.
+
+Each function returns a dict matching the ``EventInput`` wire format
+with dot-notation property keys. All builders:
+
+- Auto-generate ``event_id``, ``timestamp``, ``span_id`` if not provided
+- Pull ``trace_id`` from the current context if not provided
+- Accept ``**extra`` for custom properties
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from agentos.tracing import (
+    generate_event_id,
+    generate_span_id,
+    generate_trace_id,
+    get_current_context,
+    utcnow_iso,
+)
+from agentos.types import EventType
+
+
+def _base_envelope(
+    event_type: EventType,
+    agent_id: str,
+    *,
+    trace_id: str | None = None,
+    span_id: str | None = None,
+    parent_span_id: str | None = None,
+    task_run_id: str | None = None,
+    user_id: str | None = None,
+    session_id: str | None = None,
+    environment: str | None = None,
+) -> dict[str, Any]:
+    """Build the common event envelope, filling from context if available."""
+    ctx = get_current_context()
+
+    envelope: dict[str, Any] = {
+        "event_id": generate_event_id(),
+        "event_type": event_type.value,
+        "timestamp": utcnow_iso(),
+        "trace_id": trace_id or (ctx.trace_id if ctx else generate_trace_id()),
+        "span_id": span_id or (ctx.span_id if ctx else generate_span_id()),
+        "agent_id": agent_id,
+    }
+
+    # Optional fields — prefer explicit args, then context, then omit
+    pid = parent_span_id or (ctx.parent_span_id if ctx else None)
+    if pid:
+        envelope["parent_span_id"] = pid
+
+    trid = task_run_id or (ctx.task_run_id if ctx else None)
+    if trid:
+        envelope["task_run_id"] = trid
+
+    uid = user_id or (ctx.user_id if ctx else None)
+    if uid:
+        envelope["user_id"] = uid
+
+    sid = session_id or (ctx.session_id if ctx else None)
+    if sid:
+        envelope["session_id"] = sid
+
+    if environment:
+        envelope["environment"] = environment
+
+    return envelope
+
+
+def _set_if_not_none(d: dict[str, Any], key: str, value: Any) -> None:
+    """Set key in dict only if value is not None."""
+    if value is not None:
+        d[key] = value
+
+
+# ---------------------------------------------------------------------------
+# Event builders
+# ---------------------------------------------------------------------------
+
+
+def llm_call(
+    agent_id: str,
+    *,
+    model: str,
+    system: str,
+    trace_id: str | None = None,
+    span_id: str | None = None,
+    parent_span_id: str | None = None,
+    task_run_id: str | None = None,
+    user_id: str | None = None,
+    session_id: str | None = None,
+    environment: str | None = None,
+    response_model: str | None = None,
+    temperature: float | None = None,
+    top_p: float | None = None,
+    max_tokens: int | None = None,
+    stop_sequences: list[str] | None = None,
+    response_id: str | None = None,
+    finish_reason: str | None = None,
+    input_tokens: int | None = None,
+    output_tokens: int | None = None,
+    total_tokens: int | None = None,
+    input: list[dict[str, Any]] | None = None,
+    output: list[dict[str, Any]] | None = None,
+    duration_ms: float | None = None,
+    error: dict[str, str] | None = None,
+    stream: bool | None = None,
+    **extra: Any,
+) -> dict[str, Any]:
+    """Build an ``agent.llm_call`` event."""
+    event = _base_envelope(
+        EventType.LLM_CALL,
+        agent_id,
+        trace_id=trace_id,
+        span_id=span_id,
+        parent_span_id=parent_span_id,
+        task_run_id=task_run_id,
+        user_id=user_id,
+        session_id=session_id,
+        environment=environment,
+    )
+
+    props: dict[str, Any] = {
+        "gen_ai.system": system,
+        "gen_ai.request.model": model,
+    }
+
+    _set_if_not_none(props, "gen_ai.response.model", response_model)
+    _set_if_not_none(props, "gen_ai.request.temperature", temperature)
+    _set_if_not_none(props, "gen_ai.request.top_p", top_p)
+    _set_if_not_none(props, "gen_ai.request.max_tokens", max_tokens)
+    _set_if_not_none(props, "gen_ai.request.stop_sequences", stop_sequences)
+    _set_if_not_none(props, "gen_ai.response.id", response_id)
+    _set_if_not_none(props, "gen_ai.response.finish_reason", finish_reason)
+    _set_if_not_none(props, "gen_ai.usage.input_tokens", input_tokens)
+    _set_if_not_none(props, "gen_ai.usage.output_tokens", output_tokens)
+    _set_if_not_none(props, "gen_ai.usage.total_tokens", total_tokens)
+    _set_if_not_none(props, "input", input)
+    _set_if_not_none(props, "output", output)
+    _set_if_not_none(props, "duration_ms", duration_ms)
+    _set_if_not_none(props, "error", error)
+    _set_if_not_none(props, "stream", stream)
+
+    props.update(extra)
+    event["properties"] = props
+    return event
+
+
+def tool_call(
+    agent_id: str,
+    *,
+    tool_name: str,
+    trace_id: str | None = None,
+    span_id: str | None = None,
+    parent_span_id: str | None = None,
+    task_run_id: str | None = None,
+    user_id: str | None = None,
+    session_id: str | None = None,
+    environment: str | None = None,
+    description: str | None = None,
+    input: Any | None = None,
+    output: Any | None = None,
+    status: str | None = None,
+    duration_ms: float | None = None,
+    error: dict[str, str] | None = None,
+    **extra: Any,
+) -> dict[str, Any]:
+    """Build an ``agent.tool_call`` event."""
+    event = _base_envelope(
+        EventType.TOOL_CALL,
+        agent_id,
+        trace_id=trace_id,
+        span_id=span_id,
+        parent_span_id=parent_span_id,
+        task_run_id=task_run_id,
+        user_id=user_id,
+        session_id=session_id,
+        environment=environment,
+    )
+
+    props: dict[str, Any] = {"tool.name": tool_name}
+    _set_if_not_none(props, "tool.description", description)
+    _set_if_not_none(props, "tool.input", input)
+    _set_if_not_none(props, "tool.output", output)
+    _set_if_not_none(props, "tool.status", status)
+    _set_if_not_none(props, "duration_ms", duration_ms)
+    _set_if_not_none(props, "error", error)
+
+    props.update(extra)
+    event["properties"] = props
+    return event
+
+
+def handoff(
+    agent_id: str,
+    *,
+    target_agent_id: str,
+    trace_id: str | None = None,
+    span_id: str | None = None,
+    parent_span_id: str | None = None,
+    task_run_id: str | None = None,
+    user_id: str | None = None,
+    session_id: str | None = None,
+    environment: str | None = None,
+    reason: str | None = None,
+    context: dict[str, Any] | None = None,
+    **extra: Any,
+) -> dict[str, Any]:
+    """Build an ``agent.handoff`` event."""
+    event = _base_envelope(
+        EventType.AGENT_HANDOFF,
+        agent_id,
+        trace_id=trace_id,
+        span_id=span_id,
+        parent_span_id=parent_span_id,
+        task_run_id=task_run_id,
+        user_id=user_id,
+        session_id=session_id,
+        environment=environment,
+    )
+
+    props: dict[str, Any] = {"target_agent_id": target_agent_id}
+    _set_if_not_none(props, "reason", reason)
+    _set_if_not_none(props, "context", context)
+
+    props.update(extra)
+    event["properties"] = props
+    return event
+
+
+def retrieval_query(
+    agent_id: str,
+    *,
+    source: str,
+    trace_id: str | None = None,
+    span_id: str | None = None,
+    parent_span_id: str | None = None,
+    task_run_id: str | None = None,
+    user_id: str | None = None,
+    session_id: str | None = None,
+    environment: str | None = None,
+    query: str | None = None,
+    top_k: int | None = None,
+    results_count: int | None = None,
+    documents: list[dict[str, Any]] | None = None,
+    duration_ms: float | None = None,
+    **extra: Any,
+) -> dict[str, Any]:
+    """Build an ``agent.retrieval_query`` event."""
+    event = _base_envelope(
+        EventType.RETRIEVAL_QUERY,
+        agent_id,
+        trace_id=trace_id,
+        span_id=span_id,
+        parent_span_id=parent_span_id,
+        task_run_id=task_run_id,
+        user_id=user_id,
+        session_id=session_id,
+        environment=environment,
+    )
+
+    props: dict[str, Any] = {"retrieval.source": source}
+    _set_if_not_none(props, "retrieval.query", query)
+    _set_if_not_none(props, "retrieval.top_k", top_k)
+    _set_if_not_none(props, "retrieval.results_count", results_count)
+    _set_if_not_none(props, "retrieval.documents", documents)
+    _set_if_not_none(props, "duration_ms", duration_ms)
+
+    props.update(extra)
+    event["properties"] = props
+    return event
+
+
+def eval_result(
+    agent_id: str,
+    *,
+    eval_name: str,
+    score: float,
+    trace_id: str | None = None,
+    span_id: str | None = None,
+    parent_span_id: str | None = None,
+    task_run_id: str | None = None,
+    user_id: str | None = None,
+    session_id: str | None = None,
+    environment: str | None = None,
+    label: str | None = None,
+    reason: str | None = None,
+    evaluator: str | None = None,
+    ref_span_id: str | None = None,
+    ref_task_run_id: str | None = None,
+    **extra: Any,
+) -> dict[str, Any]:
+    """Build an ``agent.eval`` event."""
+    event = _base_envelope(
+        EventType.EVAL,
+        agent_id,
+        trace_id=trace_id,
+        span_id=span_id,
+        parent_span_id=parent_span_id,
+        task_run_id=task_run_id,
+        user_id=user_id,
+        session_id=session_id,
+        environment=environment,
+    )
+
+    props: dict[str, Any] = {
+        "eval.name": eval_name,
+        "eval.score": score,
+    }
+    _set_if_not_none(props, "eval.label", label)
+    _set_if_not_none(props, "eval.reason", reason)
+    _set_if_not_none(props, "eval.evaluator", evaluator)
+    _set_if_not_none(props, "eval.ref_span_id", ref_span_id)
+    _set_if_not_none(props, "eval.ref_task_run_id", ref_task_run_id)
+
+    props.update(extra)
+    event["properties"] = props
+    return event
+
+
+def security_alert(
+    agent_id: str,
+    *,
+    alert_type: str,
+    severity: str,
+    trace_id: str | None = None,
+    span_id: str | None = None,
+    parent_span_id: str | None = None,
+    task_run_id: str | None = None,
+    user_id: str | None = None,
+    session_id: str | None = None,
+    environment: str | None = None,
+    description: str | None = None,
+    scanner: str | None = None,
+    ref_span_id: str | None = None,
+    action_taken: str | None = None,
+    details: dict[str, Any] | None = None,
+    **extra: Any,
+) -> dict[str, Any]:
+    """Build an ``agent.security_alert`` event."""
+    event = _base_envelope(
+        EventType.SECURITY_ALERT,
+        agent_id,
+        trace_id=trace_id,
+        span_id=span_id,
+        parent_span_id=parent_span_id,
+        task_run_id=task_run_id,
+        user_id=user_id,
+        session_id=session_id,
+        environment=environment,
+    )
+
+    props: dict[str, Any] = {
+        "security.alert_type": alert_type,
+        "security.severity": severity,
+    }
+    _set_if_not_none(props, "security.description", description)
+    _set_if_not_none(props, "security.scanner", scanner)
+    _set_if_not_none(props, "security.ref_span_id", ref_span_id)
+    _set_if_not_none(props, "security.action_taken", action_taken)
+    _set_if_not_none(props, "security.details", details)
+
+    props.update(extra)
+    event["properties"] = props
+    return event
+
+
+def flag_check(
+    agent_id: str,
+    *,
+    flag_key: str,
+    flag_value: Any,
+    trace_id: str | None = None,
+    span_id: str | None = None,
+    parent_span_id: str | None = None,
+    task_run_id: str | None = None,
+    user_id: str | None = None,
+    session_id: str | None = None,
+    environment: str | None = None,
+    variant: str | None = None,
+    experiment_id: str | None = None,
+    reason: str | None = None,
+    **extra: Any,
+) -> dict[str, Any]:
+    """Build an ``agent.flag_check`` event."""
+    event = _base_envelope(
+        EventType.FLAG_CHECK,
+        agent_id,
+        trace_id=trace_id,
+        span_id=span_id,
+        parent_span_id=parent_span_id,
+        task_run_id=task_run_id,
+        user_id=user_id,
+        session_id=session_id,
+        environment=environment,
+    )
+
+    props: dict[str, Any] = {
+        "flag.key": flag_key,
+        "flag.value": flag_value,
+    }
+    _set_if_not_none(props, "flag.variant", variant)
+    _set_if_not_none(props, "flag.experiment_id", experiment_id)
+    _set_if_not_none(props, "flag.reason", reason)
+
+    props.update(extra)
+    event["properties"] = props
+    return event
+
+
+def business_event(
+    agent_id: str,
+    *,
+    event_name: str,
+    trace_id: str | None = None,
+    span_id: str | None = None,
+    parent_span_id: str | None = None,
+    task_run_id: str | None = None,
+    user_id: str | None = None,
+    session_id: str | None = None,
+    environment: str | None = None,
+    revenue: float | None = None,
+    currency: str | None = None,
+    metadata: dict[str, Any] | None = None,
+    **extra: Any,
+) -> dict[str, Any]:
+    """Build an ``agent.business_event`` event."""
+    event = _base_envelope(
+        EventType.BUSINESS_EVENT,
+        agent_id,
+        trace_id=trace_id,
+        span_id=span_id,
+        parent_span_id=parent_span_id,
+        task_run_id=task_run_id,
+        user_id=user_id,
+        session_id=session_id,
+        environment=environment,
+    )
+
+    props: dict[str, Any] = {"business.event_name": event_name}
+    _set_if_not_none(props, "business.revenue", revenue)
+    _set_if_not_none(props, "business.currency", currency)
+    _set_if_not_none(props, "business.metadata", metadata)
+
+    props.update(extra)
+    event["properties"] = props
+    return event

agentos/integrations/__init__.py

@@ -0,0 +1,15 @@
+"""Framework integrations for automatic LLM call tracing.
+
+Usage::
+
+    from agentos import wrap_openai
+    import openai
+
+    client = openai.OpenAI()
+    client = wrap_openai(client, api_key="aos_...", agent_id="my-agent")
+"""
+
+from agentos.integrations.anthropic import wrap_anthropic
+from agentos.integrations.openai import wrap_openai
+
+__all__ = ["wrap_openai", "wrap_anthropic"]

agentos/integrations/anthropic.py

@@ -0,0 +1,130 @@
+"""Anthropic SDK integration — wraps client to auto-capture LLM calls.
+
+Usage::
+
+    from agentos.integrations.anthropic import wrap_anthropic
+    import anthropic
+
+    client = anthropic.Anthropic()
+    client = wrap_anthropic(client, api_key="aos_...", agent_id="my-agent")
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from typing import Any
+
+logger = logging.getLogger("agentos.integrations.anthropic")
+
+
+def wrap_anthropic(
+    anthropic_client: Any,
+    *,
+    api_key: str | None = None,
+    agent_id: str = "default",
+    base_url: str = "https://api.agentos.dev",
+    capture_content: bool = True,
+    **kwargs: Any,
+) -> Any:
+    """Wrap an Anthropic client to auto-capture LLM calls.
+
+    Returns the same client with patched ``messages.create`` method.
+    """
+    from agentos.client import AgentOS, get_client
+
+    aos_client = None
+    if api_key:
+        aos_client = AgentOS(api_key=api_key, base_url=base_url, flush_interval=0, **kwargs)
+    else:
+        aos_client = get_client()
+
+    if aos_client is None:
+        logger.warning("No AgentOS client configured — wrap_anthropic is a no-op")
+        return anthropic_client
+
+    original_create = anthropic_client.messages.create
+
+    def patched_create(*args: Any, **call_kwargs: Any) -> Any:
+        model_arg = call_kwargs.get("model", "unknown")
+        messages = call_kwargs.get("messages")
+        stream = call_kwargs.get("stream", False)
+
+        start = time.monotonic()
+        try:
+            response = original_create(*args, **call_kwargs)
+
+            if stream:
+                return response
+
+            duration_ms = (time.monotonic() - start) * 1000
+
+            data: dict[str, Any] = {
+                "model": getattr(response, "model", model_arg),
+                "system": "anthropic",
+                "duration_ms": duration_ms,
+            }
+
+            # Usage
+            usage = getattr(response, "usage", None)
+            if usage:
+                data["input_tokens"] = getattr(usage, "input_tokens", None)
+                data["output_tokens"] = getattr(usage, "output_tokens", None)
+
+            # Finish reason
+            stop_reason = getattr(response, "stop_reason", None)
+            if stop_reason:
+                reason_map = {"end_turn": "stop", "max_tokens": "length", "tool_use": "tool_calls"}
+                data["finish_reason"] = reason_map.get(stop_reason, stop_reason)
+
+            # Response ID
+            resp_id = getattr(response, "id", None)
+            if resp_id:
+                data["response_id"] = resp_id
+
+            # Content
+            if capture_content:
+                if messages:
+                    data["input"] = [
+                        {"role": m.get("role", ""), "content": m.get("content")}
+                        for m in messages
+                        if isinstance(m, dict)
+                    ]
+                content_blocks = getattr(response, "content", [])
+                if content_blocks:
+                    text_parts = []
+                    for block in content_blocks:
+                        if getattr(block, "type", None) == "text":
+                            text_parts.append(getattr(block, "text", ""))
+                    if text_parts:
+                        data["output"] = [{"role": "assistant", "content": "\n".join(text_parts)}]
+
+            # Temperature from call args
+            if "temperature" in call_kwargs:
+                data["temperature"] = call_kwargs["temperature"]
+            if "max_tokens" in call_kwargs:
+                data["max_tokens"] = call_kwargs["max_tokens"]
+
+            try:
+                aos_client.llm_call(agent_id, **data)
+            except Exception:
+                logger.exception("Failed to capture Anthropic call")
+
+            return response
+
+        except Exception as exc:
+            duration_ms = (time.monotonic() - start) * 1000
+            try:
+                aos_client.llm_call(
+                    agent_id,
+                    model=model_arg,
+                    system="anthropic",
+                    duration_ms=duration_ms,
+                    error={"type": type(exc).__name__, "message": str(exc)},
+                )
+            except Exception:
+                logger.exception("Failed to capture Anthropic error")
+            raise
+
+    anthropic_client.messages.create = patched_create
+    return anthropic_client