infralo 0.1.0.dev1__py3-none-any.whl

infralo/__init__.py

@@ -0,0 +1,9 @@
+"""
+infralo — Observability SDK for agentic AI workflows.
+"""
+
+from infralo.client import Infralo
+from infralo.trace import get_active_trace
+
+__all__ = ["Infralo", "get_active_trace"]
+__version__ = "0.1.0"

infralo/_models.py

@@ -0,0 +1,84 @@
+"""
+infralo — Internal span data model.
+
+ToolSpanData mirrors the ToolCallSpanCreate API schema and is serialised
+to JSON before submission to the /api/v1/traces/spans/tool-call endpoint.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any
+
+_PREVIEW_MAX_CHARS = 500
+
+
+def _to_preview(value: Any) -> tuple[str, bool]:
+    """Return (preview_str, is_truncated) for any value."""
+    text = value if isinstance(value, str) else str(value)
+    if len(text) > _PREVIEW_MAX_CHARS:
+        return text[:_PREVIEW_MAX_CHARS], True
+    return text, False
+
+
+def _isoformat(dt: datetime) -> str:
+    if dt.tzinfo is None:
+        dt = dt.replace(tzinfo=timezone.utc)
+    return dt.isoformat()
+
+
+@dataclass
+class ToolSpanData:
+    """
+    Internal representation of one tool call span.
+    Serialises to the ToolCallSpanCreate API body shape.
+    """
+
+    trace_id: str
+    span_id: str
+    parent_span_id: str
+    sequence_number: int
+    tool_name: str
+    tool_call_id: str
+    start_time: datetime
+    end_time: datetime
+
+    session_id: str = ""
+    status: str = "success"  # "success" | "error" | "timeout"
+    error_code: str | None = None
+    error_message: str | None = None
+
+    input_preview: str = ""
+    is_input_truncated: bool = False
+    output_preview: str = ""
+    is_output_truncated: bool = False
+
+    input_payload: dict | None = None
+    output_payload: Any = None
+
+    metadata: dict = field(default_factory=dict)
+
+    def to_dict(self) -> dict:
+        """Serialise to the ToolCallSpanCreate JSON shape."""
+        return {
+            "trace_id": self.trace_id,
+            "span_id": self.span_id,
+            "parent_span_id": self.parent_span_id,
+            "sequence_number": self.sequence_number,
+            "tool_name": self.tool_name,
+            "tool_call_id": self.tool_call_id,
+            "start_time": _isoformat(self.start_time),
+            "end_time": _isoformat(self.end_time),
+            "session_id": self.session_id,
+            "status": self.status,
+            "error_code": self.error_code,
+            "error_message": self.error_message,
+            "input_preview": self.input_preview,
+            "is_input_truncated": self.is_input_truncated,
+            "output_preview": self.output_preview,
+            "is_output_truncated": self.is_output_truncated,
+            "input_payload": self.input_payload,
+            "output_payload": self.output_payload,
+            "metadata": self.metadata,
+        }

infralo/client.py

@@ -0,0 +1,251 @@
+"""
+infralo — Infralo client: main entry point.
+
+    from infralo import Infralo
+
+    infralo = Infralo(api_key="vk_your_key")
+
+    with infralo.start_trace(session_id="user-123") as trace:
+        ...
+"""
+
+from __future__ import annotations
+
+import asyncio
+import atexit
+import functools
+import inspect
+import uuid
+from typing import Any, Callable, TypeVar
+
+from infralo.exporter import SpanExporter
+from infralo.trace import Trace, _active_trace
+
+F = TypeVar("F", bound=Callable)
+
+
+class Infralo:
+    """
+    Main Infralo SDK client.
+
+    Create one instance per application (typically at module level) and reuse it.
+    The background exporter starts immediately and flushes remaining spans on exit.
+
+    Args:
+        api_key: Your Infralo Virtual API Key (starts with vk_).
+        endpoint: Base URL of your Infralo gateway. Default: https://api.infralo.com
+        flush_interval: Seconds between background flush cycles. Default: 5.0
+        max_batch_size: Max spans per HTTP request. Default: 50
+    """
+
+    def __init__(
+        self,
+        api_key: str,
+        endpoint: str = "https://api.infralo.com",
+        flush_interval: float = 5.0,
+        max_batch_size: int = 50,
+        max_queue_size: int = 10000,
+    ) -> None:
+        self._exporter = SpanExporter(
+            api_key=api_key,
+            endpoint=endpoint,
+            flush_interval=flush_interval,
+            max_batch_size=max_batch_size,
+            max_queue_size=max_queue_size,
+        )
+        # Guaranteed final flush when the process exits (scripts, workers, servers)
+        atexit.register(self._exporter.flush)
+
+    # ── Trace factory ────────────────────────────────────────────────────────
+
+    def start_trace(
+        self,
+        trace_id: str | None = None,
+        session_id: str = "",
+    ) -> Trace:
+        """
+        Create a new Trace for one agentic chain run.
+
+        Args:
+            trace_id: Shared across the entire chain. Auto-generated UUID if omitted.
+                      Pass the same trace_id on every LLM gateway call via the
+                      x-infralo-trace-id header.
+            session_id: Optional user/session identifier for grouping traces
+                        (e.g., a user_id, conversation_id).
+
+        Usage:
+            with infralo.start_trace(session_id="user-123") as trace:
+                response = client.chat.completions.create(
+                    model="my-deployment",
+                    messages=[...],
+                    extra_headers={"x-infralo-trace-id": trace.trace_id},
+                )
+                trace.update_from_response(response)
+                ...
+        """
+        return Trace(
+            trace_id=trace_id or str(uuid.uuid4()),
+            session_id=session_id,
+            exporter=self._exporter,
+        )
+
+    # ── @tool decorator ──────────────────────────────────────────────────────
+
+    def tool(
+        self,
+        fn: F | None = None,
+        *,
+        name: str | None = None,
+        capture_input: bool = True,
+        capture_output: bool = True,
+    ) -> F | Callable[[F], F]:
+        """
+        Decorator that automatically wraps a sync or async function as a tool span.
+
+        Reads the active Trace from contextvars — must be called inside a
+        `with infralo.start_trace(...):` block. Works transparently with both
+        sync and async functions.
+
+        Args:
+            name: Override the span name (defaults to the function's __name__).
+            capture_input: Auto-capture all function kwargs as input_payload
+                           and input_preview. Default True.
+                           Set False to avoid storing sensitive arguments (e.g., passwords).
+            capture_output: Auto-capture the return value as output_payload
+                            and output_preview. Default True.
+                            Set False to avoid storing sensitive outputs.
+
+        Metadata:
+            Pass _infralo_metadata={...} as a keyword argument to the decorated
+            function to attach custom metadata to the span:
+                result = my_tool(query="test", _infralo_metadata={"version": "2"})
+
+        Usage:
+            @infralo.tool
+            def search_web(query: str) -> list[dict]:
+                return search_api(query)
+
+            @infralo.tool(name="custom_name", capture_input=False)
+            async def search_web(query: str) -> list[dict]:
+                return await async_search_api(query)
+
+        Parallel usage (inside `async with trace.parallel():`):
+            async with trace.parallel():
+                results = await asyncio.gather(
+                    search_web("q1"),   # @infralo.tool — shares seq number
+                    fetch_doc("url"),   # @infralo.tool — shares seq number
+                )
+        """
+
+        def decorator(func: F) -> F:
+            tool_name = name or func.__name__
+
+            def _build_input(args: tuple, kwargs: dict) -> dict | None:
+                if not capture_input:
+                    return None
+                try:
+                    sig = inspect.signature(func)
+                    bound = sig.bind(*args, **kwargs)
+                    bound.apply_defaults()
+                    # Exclude private infralo kwargs
+                    return {
+                        k: v for k, v in bound.arguments.items() if not k.startswith("_infralo_")
+                    }
+                except Exception:
+                    return None
+
+            def _pop_infralo_kwargs(kwargs: dict) -> dict:
+                """Extract and remove _infralo_* kwargs before calling the real function."""
+                return {k: kwargs.pop(k) for k in list(kwargs) if k.startswith("_infralo_")}
+
+            @functools.wraps(func)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                trace = _active_trace.get()
+                if trace is None:
+                    raise RuntimeError(
+                        f"@infralo.tool '{tool_name}' was called outside of an active trace. "
+                        "Wrap your code with `with infralo.start_trace(): ...`."
+                    )
+
+                infralo_kwargs = _pop_infralo_kwargs(kwargs)
+                metadata: dict = infralo_kwargs.get("_infralo_metadata", {})
+                tool_call_id: str = infralo_kwargs.get("_infralo_tool_call_id", str(uuid.uuid4()))
+                seq_override: int | None = infralo_kwargs.get("_infralo_sequence_number")
+                parent_override: str | None = infralo_kwargs.get("_infralo_parent_span_id")
+
+                input_data = _build_input(args, dict(kwargs))
+
+                with trace.tool(
+                    name=tool_name,
+                    tool_call_id=tool_call_id,
+                    sequence_number=seq_override,
+                    parent_span_id=parent_override,
+                ) as span:
+                    if input_data is not None:
+                        span.set_input(input_data, capture_payload=capture_input)
+                    if metadata:
+                        span.set_metadata(**metadata)
+                    result = func(*args, **kwargs)
+                    if capture_output:
+                        span.set_output(result, capture_payload=capture_output)
+                return result
+
+            @functools.wraps(func)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                trace = _active_trace.get()
+                if trace is None:
+                    raise RuntimeError(
+                        f"@infralo.tool '{tool_name}' was called outside of an active trace. "
+                        "Wrap your code with `with infralo.start_trace(): ...`."
+                    )
+
+                infralo_kwargs = _pop_infralo_kwargs(kwargs)
+                metadata: dict = infralo_kwargs.get("_infralo_metadata", {})
+                tool_call_id: str = infralo_kwargs.get("_infralo_tool_call_id", str(uuid.uuid4()))
+                seq_override: int | None = infralo_kwargs.get("_infralo_sequence_number")
+                parent_override: str | None = infralo_kwargs.get("_infralo_parent_span_id")
+
+                input_data = _build_input(args, dict(kwargs))
+
+                with trace.tool(
+                    name=tool_name,
+                    tool_call_id=tool_call_id,
+                    sequence_number=seq_override,
+                    parent_span_id=parent_override,
+                ) as span:
+                    if input_data is not None:
+                        span.set_input(input_data, capture_payload=capture_input)
+                    if metadata:
+                        span.set_metadata(**metadata)
+                    result = await func(*args, **kwargs)
+                    if capture_output:
+                        span.set_output(result, capture_payload=capture_output)
+                return result
+
+            # Transparently support both sync and async functions
+            if asyncio.iscoroutinefunction(func):
+                return async_wrapper  # type: ignore[return-value]
+            return sync_wrapper  # type: ignore[return-value]
+
+        # Support both @infralo.tool and @infralo.tool(name="...", ...)
+        if fn is not None:
+            return decorator(fn)
+        return decorator
+
+    # ── Utilities ────────────────────────────────────────────────────────────
+
+    def flush(self, timeout: float = 10.0) -> None:
+        """
+        Manually flush all pending spans. Useful in tests or before shutdown.
+
+        Calling this is NOT required in production — spans are flushed automatically
+        by the background thread and on process exit.
+        """
+        self._exporter.flush(timeout)
+
+    def shutdown(self) -> None:
+        """
+        Shutdown the background exporter gracefully. Called automatically on exit.
+        Safe to call manually if you need to ensure spans are sent before teardown.
+        """
+        self._exporter.shutdown()

infralo/exporter.py

@@ -0,0 +1,154 @@
+"""
+infralo — Background HTTP span exporter.
+
+Spans are pushed to a thread-safe queue and flushed to the Infralo API
+by a single daemon thread. The caller never waits for HTTP — enqueue()
+returns immediately.
+
+Flush schedule:
+  - Every `flush_interval` seconds (default 5 s).
+  - Early if the queue reaches `max_batch_size` (default 50).
+  - Final flush on process exit via atexit (called by Infralo.__init__).
+"""
+
+from __future__ import annotations
+
+import logging
+import queue
+import threading
+from typing import TYPE_CHECKING
+
+import httpx
+
+if TYPE_CHECKING:
+    from infralo._models import ToolSpanData
+
+logger = logging.getLogger("infralo.exporter")
+
+
+class SpanExporter:
+    """
+    Background HTTP exporter for tool call spans.
+
+    Not intended for direct use — created and managed by the Infralo client.
+    """
+
+    def __init__(
+        self,
+        api_key: str,
+        endpoint: str,
+        flush_interval: float = 5.0,
+        max_batch_size: int = 50,
+        max_queue_size: int = 10000,
+    ) -> None:
+        self.api_key = api_key
+        clean_endpoint = endpoint.rstrip("/")
+        for suffix in ("/api/v1", "/v1"):
+            if clean_endpoint.endswith(suffix):
+                clean_endpoint = clean_endpoint[: -len(suffix)]
+        self.endpoint = clean_endpoint.rstrip("/")
+        self.flush_interval = flush_interval
+        self.max_batch_size = max_batch_size
+        self.max_queue_size = max_queue_size
+
+        self._queue: queue.Queue[ToolSpanData] = queue.Queue(maxsize=max_queue_size)
+        self._flush_event = threading.Event()
+        self._stop_event = threading.Event()
+        self._flush_lock = threading.Lock()
+
+        # Retry backoff state
+        self._current_interval = flush_interval
+        self._max_backoff = 60.0
+
+        self._thread = threading.Thread(
+            target=self._flush_loop,
+            name="infralo-exporter",
+            daemon=True,  # dies with the process — atexit handles final flush
+        )
+        self._thread.start()
+
+    # ── Public API ──────────────────────────────────────────────────────────
+
+    def enqueue(self, span: "ToolSpanData") -> None:
+        """
+        Push a span onto the export queue. Always returns immediately.
+        Never raises — failures are logged and the span is re-queued.
+        """
+        try:
+            self._queue.put_nowait(span)
+        except queue.Full:
+            logger.warning("infralo: export queue full — dropping span %s", span.span_id)
+            return
+
+        if self._queue.qsize() >= self.max_batch_size:
+            self._flush_event.set()  # wake the background thread early
+
+    def flush(self, timeout: float = 10.0) -> None:
+        """
+        Synchronously drain all pending spans from the queue.
+        Called automatically on process exit (atexit). Safe to call manually in tests.
+        """
+        self._do_flush()
+
+    def shutdown(self) -> None:
+        """Stop the background thread and do a final flush. Idempotent."""
+        self._stop_event.set()
+        self._flush_event.set()
+        self._thread.join(timeout=5.0)
+        self._do_flush()
+
+    # ── Background thread ───────────────────────────────────────────────────
+
+    def _flush_loop(self) -> None:
+        """Daemon thread: flush on interval or early-wake event with backoff on failure."""
+        while not self._stop_event.is_set():
+            self._flush_event.wait(timeout=self._current_interval)
+            self._flush_event.clear()
+            self._do_flush()
+        # Final drain after stop signal
+        self._do_flush()
+
+    def _do_flush(self) -> None:
+        """Drain up to max_batch_size spans and POST them to the API."""
+        with self._flush_lock:
+            batch: list[ToolSpanData] = []
+            try:
+                while len(batch) < self.max_batch_size:
+                    batch.append(self._queue.get_nowait())
+            except queue.Empty:
+                pass
+
+            if not batch:
+                return
+
+            try:
+
+                resp = httpx.post(
+                    f"{self.endpoint}/api/v1/traces/spans/tool-call",
+                    headers={"Authorization": f"Bearer {self.api_key}"},
+                    json=[s.to_dict() for s in batch],
+                    timeout=10.0,
+                )
+                resp.raise_for_status()
+                logger.debug("infralo: flushed %d spans", len(batch))
+
+                # Reset backoff on success
+                self._current_interval = self.flush_interval
+
+            except Exception as exc:
+                logger.warning(
+                    "infralo: failed to flush %d span(s): %s — re-enqueueing for retry",
+                    len(batch),
+                    exc,
+                )
+                # Apply exponential backoff
+                self._current_interval = min(self._current_interval * 2.0, self._max_backoff)
+
+                # Best-effort re-enqueue for next flush cycle
+                for span in batch:
+                    try:
+                        self._queue.put_nowait(span)
+                    except queue.Full:
+                        logger.warning(
+                            "infralo: queue full during retry — dropping span %s", span.span_id
+                        )

infralo/integrations/__init__.py

@@ -0,0 +1,23 @@
+"""
+infralo — Agent framework integrations.
+
+Each integration is a separate optional module.  Install only the extras you
+need; the core ``infralo`` package intentionally has no framework dependencies.
+
++-----------------+-------------------------------------+---------------------------+
+| Framework       | Install extra                       | Import path               |
++-----------------+-------------------------------------+---------------------------+
+| Agno            | pip install infralo[agno]           | infralo.integrations.agno |
+| LangGraph/Chain | pip install infralo[langgraph]      | infralo.integrations.langgraph |
+| CrewAI          | pip install infralo[crewai]         | infralo.integrations.crewai |
+| Google ADK      | pip install infralo[google-adk]     | infralo.integrations.google_adk |
++-----------------+-------------------------------------+---------------------------+
+
+All integrations follow the same three-step pattern:
+
+    1. Create the Infralo client once at module level.
+    2. Create / register the framework-specific adapter.
+    3. Wrap every agentic run with ``infralo.start_trace()``.
+
+See the individual module docstrings and ``examples/`` for complete usage.
+"""