agentos-python 0.1.0__py3-none-any.whl

agentos/compat/langfuse.py

@@ -0,0 +1,317 @@
+"""Langfuse-compatible API shim.
+
+Drop-in replacement: change ``from langfuse import Langfuse`` to
+``from agentos.compat.langfuse import Langfuse``. All existing code works.
+
+Maps Langfuse primitives to AgentOS events:
+- trace() → TraceContext
+- generation() → agent.llm_call
+- span() → child span
+- score() → agent.eval
+- event() → agent.business_event
+"""
+
+from __future__ import annotations
+
+import contextlib
+from typing import Any
+
+from agentos.client import AgentOS
+from agentos.tracing import (
+    generate_span_id,
+    generate_trace_id,
+)
+
+
+class StatefulGenerationClient:
+    """Represents a Langfuse generation (LLM call)."""
+
+    def __init__(
+        self,
+        client: AgentOS,
+        trace_ctx: _TraceShim,
+        *,
+        name: str | None = None,
+        model: str | None = None,
+        model_parameters: dict[str, Any] | None = None,
+        input: Any | None = None,
+        output: Any | None = None,
+        usage: dict[str, Any] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        self._client = client
+        self._trace = trace_ctx
+        self._span_id = generate_span_id()
+        self._name = name
+        self._model = model or "unknown"
+        self._model_parameters = model_parameters or {}
+        self._input = input
+        self._output = output
+        self._usage = usage or {}
+        self._metadata = metadata or {}
+        self._ended = False
+
+    def end(
+        self,
+        *,
+        output: Any | None = None,
+        usage: dict[str, Any] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """End the generation and send the event."""
+        if output is not None:
+            self._output = output
+        if usage:
+            self._usage.update(usage)
+        if metadata:
+            self._metadata.update(metadata)
+        self._flush_event()
+        self._ended = True
+
+    def _flush_event(self) -> None:
+        """Send the llm_call event."""
+        props: dict[str, Any] = {
+            "model": self._model,
+            "system": "unknown",  # Langfuse doesn't require provider
+        }
+
+        if self._usage:
+            if "input" in self._usage:
+                props["input_tokens"] = self._usage["input"]
+            elif "promptTokens" in self._usage:
+                props["input_tokens"] = self._usage["promptTokens"]
+            if "output" in self._usage:
+                props["output_tokens"] = self._usage["output"]
+            elif "completionTokens" in self._usage:
+                props["output_tokens"] = self._usage["completionTokens"]
+            if "total" in self._usage:
+                props["total_tokens"] = self._usage["total"]
+
+        if self._model_parameters.get("temperature") is not None:
+            props["temperature"] = self._model_parameters["temperature"]
+        if self._model_parameters.get("max_tokens") is not None:
+            props["max_tokens"] = self._model_parameters["max_tokens"]
+
+        if self._input is not None:
+            if isinstance(self._input, list):
+                props["input"] = self._input
+            else:
+                props["input"] = [{"role": "user", "content": str(self._input)}]
+
+        if self._output is not None:
+            if isinstance(self._output, list):
+                props["output"] = self._output
+            else:
+                props["output"] = [{"role": "assistant", "content": str(self._output)}]
+
+        self._client.llm_call(
+            self._trace._agent_id,
+            trace_id=self._trace._trace_id,
+            span_id=self._span_id,
+            parent_span_id=self._trace._current_span_id,
+            user_id=self._trace._user_id,
+            session_id=self._trace._session_id,
+            **props,
+        )
+
+    def __del__(self) -> None:
+        if not self._ended:
+            with contextlib.suppress(Exception):
+                self._flush_event()
+
+
+class StatefulSpanClient:
+    """Represents a Langfuse span."""
+
+    def __init__(
+        self,
+        client: AgentOS,
+        trace_ctx: _TraceShim,
+        *,
+        name: str | None = None,
+        input: Any | None = None,
+        output: Any | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        self._client = client
+        self._trace = trace_ctx
+        self._span_id = generate_span_id()
+        self._parent_span_id = trace_ctx._current_span_id
+        self._name = name
+        self._input = input
+        self._output = output
+        self._metadata = metadata or {}
+
+    def end(
+        self,
+        *,
+        output: Any | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """End the span."""
+        if output is not None:
+            self._output = output
+        if metadata:
+            self._metadata.update(metadata)
+
+    def generation(self, **kwargs: Any) -> StatefulGenerationClient:
+        """Create a nested generation within this span."""
+        return StatefulGenerationClient(
+            self._client,
+            self._trace,
+            **kwargs,
+        )
+
+    def span(self, **kwargs: Any) -> StatefulSpanClient:
+        """Create a nested span."""
+        return StatefulSpanClient(self._client, self._trace, **kwargs)
+
+
+class _TraceShim:
+    """Internal shim representing a Langfuse trace."""
+
+    def __init__(
+        self,
+        client: AgentOS,
+        *,
+        id: str | None = None,
+        name: str | None = None,
+        user_id: str | None = None,
+        session_id: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        tags: list[str] | None = None,
+        input: Any | None = None,
+        output: Any | None = None,
+        **kwargs: Any,
+    ) -> None:
+        self._client = client
+        self._trace_id = id or generate_trace_id()
+        self._agent_id = name or "langfuse-trace"
+        self._user_id = user_id
+        self._session_id = session_id
+        self._metadata = metadata or {}
+        self._current_span_id = generate_span_id()  # root span
+
+    def generation(self, **kwargs: Any) -> StatefulGenerationClient:
+        """Create a generation (LLM call) on this trace."""
+        return StatefulGenerationClient(self._client, self, **kwargs)
+
+    def span(self, **kwargs: Any) -> StatefulSpanClient:
+        """Create a span on this trace."""
+        return StatefulSpanClient(self._client, self, **kwargs)
+
+    def score(
+        self,
+        *,
+        name: str,
+        value: float,
+        data_type: str | None = None,
+        comment: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Attach a score to this trace."""
+        self._client.eval(
+            self._agent_id,
+            eval_name=name,
+            score=value,
+            trace_id=self._trace_id,
+            span_id=self._current_span_id,
+            user_id=self._user_id,
+            session_id=self._session_id,
+        )
+
+    def event(
+        self,
+        *,
+        name: str,
+        input: Any | None = None,
+        output: Any | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Create a lightweight event on this trace."""
+        self._client.business_event(
+            self._agent_id,
+            event_name=name,
+            trace_id=self._trace_id,
+            span_id=self._current_span_id,
+            user_id=self._user_id,
+            session_id=self._session_id,
+            metadata=metadata,
+        )
+
+
+class Langfuse:
+    """Langfuse-compatible client. Drop-in replacement.
+
+    Usage::
+
+        from agentos.compat.langfuse import Langfuse
+
+        langfuse = Langfuse(public_key="aos_...")
+        trace = langfuse.trace(name="my-trace")
+        gen = trace.generation(model="gpt-4o", usage={"input": 100, "output": 50})
+        trace.score(name="quality", value=0.95)
+        langfuse.flush()
+    """
+
+    def __init__(
+        self,
+        public_key: str | None = None,
+        secret_key: str | None = None,  # accepted but ignored
+        host: str | None = None,
+        *,
+        release: str | None = None,
+        debug: bool = False,
+        threads: int = 1,
+        flush_at: int = 15,
+        flush_interval: float = 0.5,
+        enabled: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        api_key = public_key or ""
+        base_url = host or "https://api.agentos.dev"
+
+        self._client = AgentOS(
+            api_key=api_key,
+            base_url=base_url,
+            batch_size=flush_at,
+            flush_interval=flush_interval,
+            enabled=enabled,
+            debug=debug,
+        )
+
+    def trace(self, **kwargs: Any) -> _TraceShim:
+        """Create a new trace."""
+        return _TraceShim(self._client, **kwargs)
+
+    def score(
+        self,
+        *,
+        trace_id: str | None = None,
+        observation_id: str | None = None,
+        name: str,
+        value: float,
+        data_type: str | None = None,
+        comment: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Create a score (standalone, not attached to a trace object)."""
+        self._client.eval(
+            "langfuse-score",
+            eval_name=name,
+            score=value,
+            trace_id=trace_id or generate_trace_id(),
+        )
+
+    def flush(self) -> None:
+        """Flush all pending events."""
+        self._client.flush()
+
+    def shutdown(self) -> None:
+        """Flush and close."""
+        self._client.shutdown()

agentos/compat/langsmith.py

@@ -0,0 +1,79 @@
+"""LangSmith-compatible API shim.
+
+Provides ``@traceable`` decorator and ``Client`` with feedback API.
+
+Usage::
+
+    from agentos.compat.langsmith import traceable, Client, wrap_openai
+
+    @traceable(name="my-chain", run_type="chain")
+    def my_chain(input: str) -> str: ...
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+from agentos.decorators import observe
+from agentos.integrations.openai import wrap_openai as _native_wrap_openai
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def traceable(
+    *,
+    name: str | None = None,
+    run_type: str | None = None,
+    **kwargs: Any,
+) -> Callable[[F], F]:
+    """LangSmith-compatible ``@traceable`` decorator.
+
+    Maps ``run_type`` to AgentOS observation type:
+    - ``"llm"`` → ``@observe(as_type="generation")``
+    - ``"chain"`` / ``"tool"`` / None → ``@observe()``
+    """
+    as_type = "generation" if run_type == "llm" else None
+    return observe(name=name, as_type=as_type)
+
+
+def wrap_openai(openai_client: Any, **kwargs: Any) -> Any:
+    """LangSmith-compatible ``wrap_openai``."""
+    return _native_wrap_openai(openai_client, **kwargs)
+
+
+class Client:
+    """LangSmith-compatible Client with feedback API.
+
+    Usage::
+
+        client = Client(api_key="aos_...")
+        client.create_feedback(run_id="...", key="quality", score=0.9)
+    """
+
+    def __init__(self, api_key: str | None = None, **kwargs: Any) -> None:
+        from agentos.client import AgentOS, get_client
+
+        if api_key:
+            self._client = AgentOS(api_key=api_key, flush_interval=0, **kwargs)
+        else:
+            self._client = get_client()
+
+    def create_feedback(
+        self,
+        run_id: str,
+        key: str,
+        *,
+        score: float | None = None,
+        value: Any | None = None,
+        comment: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Create feedback (maps to agent.eval event)."""
+        if self._client is None:
+            return
+        self._client.eval(
+            "langsmith-feedback",
+            eval_name=key,
+            score=score if score is not None else (1.0 if value else 0.0),
+        )

agentos/compat/otel.py

@@ -0,0 +1,167 @@
+"""OpenTelemetry / Arize Phoenix compatibility — OTLP SpanExporter.
+
+For Phoenix users: swap the OTLP endpoint to AgentOS.
+
+Usage::
+
+    from agentos.compat.otel import register
+
+    tracer_provider = register(
+        project_name="my-project",
+        endpoint="https://ingest.agentos.dev/v1/otlp",
+        api_key="aos_...",
+    )
+
+    from openinference.instrumentation.openai import OpenAIInstrumentor
+    OpenAIInstrumentor().instrument(tracer_provider=tracer_provider)
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Sequence
+from typing import Any
+
+logger = logging.getLogger("agentos.compat.otel")
+
+# OpenInference span kind → AgentOS event type
+_SPAN_KIND_MAP = {
+    "LLM": "agent.llm_call",
+    "TOOL": "agent.tool_call",
+    "CHAIN": None,  # parent context only
+    "AGENT": None,
+    "RETRIEVER": "agent.retrieval_query",
+    "EMBEDDING": "agent.tool_call",
+    "RERANKER": "agent.tool_call",
+    "GUARDRAIL": "agent.security_alert",
+}
+
+# OpenInference attribute → AgentOS property
+_ATTR_MAP = {
+    "llm.model_name": "gen_ai.request.model",
+    "llm.token_count.prompt": "gen_ai.usage.input_tokens",
+    "llm.token_count.completion": "gen_ai.usage.output_tokens",
+    "llm.token_count.total": "gen_ai.usage.total_tokens",
+    "input.value": "input",
+    "output.value": "output",
+    "llm.invocation_parameters": None,  # handled specially
+}
+
+
+def register(
+    *,
+    project_name: str = "default",
+    endpoint: str = "https://ingest.agentos.dev/v1/otlp",
+    api_key: str | None = None,
+) -> Any:
+    """Register an OTEL TracerProvider that exports to AgentOS.
+
+    Requires ``opentelemetry-sdk`` to be installed.
+    Returns a ``TracerProvider`` compatible with OpenInference instrumentors.
+    """
+    try:
+        from opentelemetry.sdk.trace import TracerProvider
+        from opentelemetry.sdk.trace.export import BatchSpanProcessor
+    except ImportError as err:
+        raise ImportError(
+            "OpenTelemetry SDK required. Install with: pip install agentos-python[otel]"
+        ) from err
+
+    exporter = AgentOSSpanExporter(
+        endpoint=endpoint,
+        api_key=api_key or "",
+        project_name=project_name,
+    )
+
+    provider = TracerProvider()
+    provider.add_span_processor(BatchSpanProcessor(exporter))
+
+    return provider
+
+
+class AgentOSSpanExporter:
+    """Custom OTEL SpanExporter that converts spans to AgentOS events."""
+
+    def __init__(
+        self,
+        endpoint: str,
+        api_key: str,
+        project_name: str = "default",
+    ) -> None:
+        from agentos.client import AgentOS
+
+        self._client = AgentOS(
+            api_key=api_key,
+            base_url=endpoint.replace("/v1/otlp", ""),
+            flush_interval=1.0,
+        )
+        self._project_name = project_name
+
+    def export(self, spans: Sequence[Any]) -> Any:
+        """Convert OTEL spans to AgentOS events and enqueue."""
+        try:
+            from opentelemetry.sdk.trace.export import SpanExportResult
+        except ImportError:
+            return None
+
+        for otel_span in spans:
+            try:
+                self._export_span(otel_span)
+            except Exception:
+                logger.exception("Failed to export span")
+
+        return SpanExportResult.SUCCESS
+
+    def _export_span(self, otel_span: Any) -> None:
+        """Convert a single OTEL span to an AgentOS event."""
+        attrs = dict(otel_span.attributes or {})
+
+        # Determine event type from OpenInference span kind
+        span_kind = attrs.get("openinference.span.kind", "CHAIN")
+        event_type = _SPAN_KIND_MAP.get(str(span_kind).upper())
+
+        if event_type is None:
+            return  # Skip CHAIN/AGENT spans (they're just context)
+
+        # Build properties
+        props: dict[str, Any] = {}
+        for otel_key, aos_key in _ATTR_MAP.items():
+            if otel_key in attrs and aos_key:
+                props[aos_key] = attrs[otel_key]
+
+        # Map to specific event type
+        if event_type == "agent.llm_call":
+            props.setdefault("gen_ai.system", "unknown")
+            props.setdefault("gen_ai.request.model", "unknown")
+            self._client.llm_call(
+                self._project_name,
+                model=props.pop("gen_ai.request.model"),
+                system=props.pop("gen_ai.system"),
+                trace_id=format(otel_span.context.trace_id, "032x"),
+                span_id=format(otel_span.context.span_id, "016x"),
+                **{k: v for k, v in props.items() if not k.startswith("gen_ai.")},
+                **{k: v for k, v in props.items() if k.startswith("gen_ai.")},
+            )
+        elif event_type == "agent.tool_call":
+            tool_name = attrs.get("tool.name", otel_span.name or "unknown")
+            self._client.tool_call(
+                self._project_name,
+                tool_name=str(tool_name),
+                trace_id=format(otel_span.context.trace_id, "032x"),
+                span_id=format(otel_span.context.span_id, "016x"),
+            )
+        elif event_type == "agent.retrieval_query":
+            source = attrs.get("retrieval.source", "unknown")
+            self._client.retrieval_query(
+                self._project_name,
+                source=str(source),
+                trace_id=format(otel_span.context.trace_id, "032x"),
+                span_id=format(otel_span.context.span_id, "016x"),
+            )
+
+    def shutdown(self) -> None:
+        self._client.shutdown()
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        self._client.flush()
+        return True

agentos/compat/weave.py

@@ -0,0 +1,93 @@
+"""W&B Weave-compatible API shim.
+
+Provides ``init``, ``@op()``, and ``Evaluation``.
+
+Usage::
+
+    from agentos.compat import weave
+    weave.init("my-project", api_key="aos_...")
+
+    @weave.op()
+    def call_llm(prompt: 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
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def init(project: str, *, api_key: str | None = None, **kwargs: Any) -> None:
+    """Initialize Weave-compatible tracing."""
+    import agentos.client
+
+    client = AgentOS(api_key=api_key or "", **kwargs)
+    agentos.client._global_client = client
+
+
+def op(**kwargs: Any) -> Callable[[F], F]:
+    """Weave-compatible ``@op()`` decorator. Equivalent to ``@observe()``."""
+    return observe(**kwargs)
+
+
+class Evaluation:
+    """Weave-compatible evaluation runner.
+
+    Usage::
+
+        evaluation = Evaluation(
+            dataset=[{"input": "2+2?", "expected": "4"}],
+            scorers=[accuracy_scorer],
+        )
+        await evaluation.evaluate(my_pipeline)
+    """
+
+    def __init__(
+        self,
+        *,
+        dataset: list[dict[str, Any]],
+        scorers: list[Callable[..., Any]],
+    ) -> None:
+        self._dataset = dataset
+        self._scorers = scorers
+
+    async def evaluate(self, fn: Callable[..., Any]) -> list[dict[str, Any]]:
+        """Run evaluation and send score events."""
+        import inspect
+
+        client = get_client()
+        results = []
+
+        for item in self._dataset:
+            input_val = item.get("input")
+            expected = item.get("expected")
+
+            if inspect.iscoroutinefunction(fn):
+                output = await fn(input_val)
+            else:
+                output = fn(input_val)
+
+            for scorer in self._scorers:
+                if inspect.iscoroutinefunction(scorer):
+                    score = await scorer(output, expected)
+                else:
+                    score = scorer(output, expected)
+
+                if client:
+                    client.eval(
+                        "weave-eval",
+                        eval_name=scorer.__name__,
+                        score=float(score),
+                    )
+
+                results.append({"input": input_val, "output": output, "score": score})
+
+        if client:
+            client.flush()
+
+        return results

agentos/config.py

@@ -0,0 +1,34 @@
+"""Configuration for the Agent OS SDK."""
+
+from __future__ import annotations
+
+import dataclasses
+
+
+@dataclasses.dataclass
+class AgentOSConfig:
+    """SDK configuration.
+
+    Args:
+        api_key: Agent OS API key (format: ``aos_...``).
+        base_url: Ingestion API base URL.
+        batch_size: Max events per batch before auto-flush.
+        flush_interval: Seconds between periodic flushes.
+        max_retries: Max retry attempts for failed HTTP requests.
+        capture_content: If False, strips message content from LLM and tool events.
+        enabled: Kill switch — when False, all capture calls are no-ops.
+        environment: Environment tag (e.g. production, staging, development).
+        debug: If True, log SDK internals to stderr.
+        max_queue_size: Max events in the queue before dropping oldest.
+    """
+
+    api_key: str = ""
+    base_url: str = "https://api.agentos.dev"
+    batch_size: int = 50
+    flush_interval: float = 5.0
+    max_retries: int = 3
+    capture_content: bool = True
+    enabled: bool = True
+    environment: str | None = None
+    debug: bool = False
+    max_queue_size: int = 10_000