infralo 0.1.0.dev1

Infralo SDK

Infralo Python SDK

Observability SDK for agentic AI workflows. Submits tool call spans to Infralo in the background — your user-facing code is never blocked.

Installation

pip install infralo
# or
uv add infralo

---

Core concept

You create one Trace per user request. The trace lives in a ContextVar and propagates automatically through async/await and thread contexts — no need to pass it around manually.

Every gateway LLM call needs two Infralo headers. trace.headers() builds them for you from the current context:

extra_headers=trace.headers()
# → {"x-infralo-trace-id": "...", "x-infralo-parent-span-id": "..."}
# The parent_span_id is automatically the last tool span submitted, or absent
# on the first call. You never track trace.last_tool_span_id yourself.

After the LLM responds, trace.update_from_response(response) reads back the gateway's x-infralo-span-id header so the next tool spans know their parent.

---

Quick Start

from infralo import Infralo
from openai import OpenAI

infralo = Infralo(api_key="vk_your_key")
client = OpenAI(base_url="https://your-infralo-gateway.com/v1")

---

Pattern A — Context manager (explicit, most control)

Use when you iterate over tool_calls from the LLM response or need fine control over what gets recorded.

def run_agent(user_message: str, session_id: str) -> str:
    with infralo.start_trace(session_id=session_id) as trace:

        # ── First LLM call ──────────────────────────────────────────────────
        response = client.chat.completions.create(
            model="my-deployment",
            messages=[{"role": "user", "content": user_message}],
            extra_headers=trace.headers(),   # ← always just this, nothing else
        )
        trace.update_from_response(response) # read x-infralo-span-id back

        # ── Execute tools and record each as a span ─────────────────────────
        tool_results = []
        for tool_call in response.choices[0].message.tool_calls:
            with trace.tool(
                name=tool_call.function.name,
                tool_call_id=tool_call.id,
                # parent_span_id defaults to the last one from update_from_response()
            ) as span:
                args = json.loads(tool_call.function.arguments)
                span.set_input(args)               # optional: record inputs
                result = dispatch_tool(tool_call)
                span.set_output(result)            # optional: record outputs
                span.set_metadata(version="1.0")   # optional: custom metadata
            tool_results.append(result)

        # ── Second LLM call ─────────────────────────────────────────────────
        # trace.headers() now automatically includes x-infralo-parent-span-id
        # pointing to the last tool span — no manual tracking needed.
        response2 = client.chat.completions.create(
            model="my-deployment",
            messages=[..., *tool_messages(response, tool_results)],
            extra_headers=trace.headers(),   # ← same call, updated automatically
        )
        return response2.choices[0].message.content

    # Trace exits: spans already queued; background thread flushes them.

Error handling in Pattern A

If the tool raises an exception inside with trace.tool(...) as span:, the span is automatically marked status="error" with the exception message. The exception still propagates normally — the SDK never swallows it.

with trace.tool(name="search", tool_call_id="call_1") as span:
    span.set_input({"query": "..."})
    result = search_api(query)    # ← if this raises, span is auto-marked as error
    span.set_output(result)       # ← only called if search_api succeeded
# ^ __exit__ is called regardless; exc_type != None → status="error"

You do not need a try/except to get error recording. You only need one if you want to recover and continue — in which case catch the exception, then decide whether to re-raise or return a fallback.

---

Pattern B — @infralo.tool decorator (ergonomic)

Best for clean code where timing and I/O capture can be automatic. The decorator reads the active trace from context — the same error-handling guarantee applies: if the wrapped function raises, the span is marked as error automatically.

@infralo.tool
def search_web(query: str) -> list[dict]:
    """Timing and input/output captured automatically."""
    return call_search_api(query)

@infralo.tool(name="fetch_doc", capture_input=False)   # omit sensitive args from payload
async def fetch_doc(url: str, auth_token: str) -> str:
    async with httpx.AsyncClient() as c:
        r = await c.get(url, headers={"Authorization": auth_token})
        return r.text

# Must be called inside an active trace
with infralo.start_trace(session_id="s1") as trace:
    response = client.chat.completions.create(...)
    trace.update_from_response(response)

    # Decorated tools automatically read trace context — no headers to pass
    results = search_web(query="latest AI news")
    doc = await fetch_doc(url="https://...", auth_token="secret")

    # Attach per-call metadata via _infralo_metadata kwarg
    results2 = search_web(
        query="page 2",
        _infralo_metadata={"source": "google", "page": 2},
    )

---

Pattern C — Manual span (full control)

Use when you need try/except for recovery logic, or when you can't use a context manager (e.g., tool starts in one coroutine and finishes in another).

Unlike Pattern A, you must call span.end() explicitly in all code paths.

with infralo.start_trace(session_id="s1") as trace:
    trace.update_from_response(llm_response)

    span = trace.start_tool_span(
        name="slow_db_query",
        tool_call_id="call_xyz",
    )
    try:
        result = db.query("SELECT ...")
        span.end(output={"rows": len(result)})              # status="success"
    except TimeoutError:
        span.end(status="timeout", error="Query exceeded 30s")
    except Exception as e:
        span.end(status="error", error=str(e), error_code="DB_ERROR")

[!IMPORTANT] Always call span.end() in every branch. If you forget, the span is never submitted. Pattern A's context manager handles this automatically — prefer it when possible.

---

Parallel Tool Execution

Use async with trace.parallel(): to give all tool spans inside the block the same sequence_number, correctly representing that they ran concurrently.

@infralo.tool
async def search(query: str) -> list[dict]: ...

@infralo.tool
async def fetch(url: str) -> str: ...

with infralo.start_trace() as trace:
    response = client.chat.completions.create(...)
    trace.update_from_response(response)

    # Parallel @tool calls — share sequence_number automatically
    async with trace.parallel():
        results = await asyncio.gather(
            search(query="q1"),
            fetch(url="https://..."),
        )
    # After the block: sequence counter advances past this group.
    # The next tool() call gets the next sequential number.

For parallel execution with explicit context managers (Pattern A style):

async with trace.parallel():
    async def run_search():
        with trace.tool("search", "call_1") as span:
            result = await search_api("q1")
            span.set_output(result)
            return result

    async def run_fetch():
        with trace.tool("fetch", "call_2") as span:
            result = await fetch_api("url")
            span.set_output(result)
            return result

    results = await asyncio.gather(run_search(), run_fetch())

Both approaches work the same way: async with trace.parallel(): pins a shared sequence number in a ContextVar before the tasks are created, so all tasks inside inherit it automatically.

---

Edge Cases

Header Propagation Failure

If trace.update_from_response(response) cannot find the x-infralo-span-id header (e.g., if a proxy or gateway stripped it):

1. A warning is logged via Python's standard logging library under the logger name infralo.trace.
2. The active parent span ID is cleared (trace._current_parent_span_id = None) to prevent subsequent tool spans from accidentally being linked to an older LLM span from a previous step.

Nested Traces

Starting a trace while another trace is already active on the same thread/async task is fully supported and behaves like a stack:

• Entering the inner with infralo.start_trace() block temporarily registers the inner trace as the active one.
• Any spans created inside the inner block belong to the inner trace.
• Exiting the inner block automatically restores the outer trace as the active trace.

---

Configuration

infralo = Infralo(
    api_key="vk_your_key",
    endpoint="https://your-infralo-gateway.com",  # default: https://api.infralo.com
    flush_interval=5.0,    # seconds between background flushes (default: 5.0)
    max_batch_size=50,     # flush early if queue reaches this size (default: 50)
    max_queue_size=10000,  # limit in-memory queue to prevent OOM on outages (default: 10000)
)

---

Framework Integration

Web Frameworks

FastAPI

from contextlib import asynccontextmanager
from fastapi import FastAPI, Request

@asynccontextmanager
async def lifespan(app: FastAPI):
    yield
    infralo.shutdown()   # flush remaining spans on shutdown

app = FastAPI(lifespan=lifespan)

@app.middleware("http")
async def infralo_middleware(request: Request, call_next):
    session_id = request.headers.get("x-session-id", "")
    with infralo.start_trace(session_id=session_id) as trace:
        request.state.infralo_trace = trace
        return await call_next(request)

Django (WSGI — each thread has its own isolated ContextVar)

class InfraLoMiddleware:
    def __init__(self, get_response):
        self.get_response = get_response

    def __call__(self, request):
        session_id = request.META.get("HTTP_X_SESSION_ID", "")
        with infralo.start_trace(session_id=session_id) as trace:
            request.infralo_trace = trace
            return self.get_response(request)

Plain Python script

with infralo.start_trace() as trace:
    ...
# atexit flushes remaining spans automatically — no manual flush needed.

---

Agent Framework Integration

Infralo provides first-class integrations for the most popular agentic frameworks. Install only the extras you need — the core infralo package has zero framework dependencies.

Framework	Install	Import
Agno	pip install infralo[agno]	infralo.integrations.agno
LangGraph / LangChain	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().

LLM span linking

Tool spans are always captured regardless of LLM routing. For the full LLM → tool parent link, route your model calls through the Infralo gateway. The gateway injects x-infralo-span-id into the HTTP response; each integration's model callback reads this header and calls trace.set_parent_span_id() automatically.

When the gateway header is unavailable, tool spans appear as roots in the trace tree (still useful) and a DEBUG-level log message is emitted.

---

Agno

from infralo import Infralo
from infralo.integrations.agno import make_agno_callbacks
from agno.agent import Agent
from agno.models.openai import OpenAIChat

infralo = Infralo(api_key="vk_...")
callbacks = make_agno_callbacks(infralo)

with infralo.start_trace(session_id="user-123") as trace:
    agent = Agent(
        model=OpenAIChat(
            id="my-deployment",
            base_url="http://localhost:8000/v1",  # Infralo gateway
            api_key="vk_...",
        ),
        tools=[search_web, get_weather],
        **callbacks,   # injects before/after_tool + after_model
    )
    agent.print_response("What's the weather in London?")

make_agno_callbacks() returns three hooks as a dict:

Key	Purpose
before_tool_callback	Starts a ToolSpan; records tool name and input args
after_tool_callback	Ends the span; records output and status
after_model_callback	Reads x-infralo-span-id from the HTTP response; sets parent span

---

LangGraph / LangChain

from infralo import Infralo
from infralo.integrations.langgraph import InfraloCallbackHandler
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent
from langchain_core.messages import HumanMessage

infralo = Infralo(api_key="vk_...")

llm = ChatOpenAI(
    model="my-deployment",
    base_url="http://localhost:8000/v1",
    api_key="vk_...",
)

with infralo.start_trace(session_id="user-123") as trace:
    agent = create_react_agent(llm, [search_web, get_weather])
    handler = InfraloCallbackHandler()   # one per invocation is safest

    result = await agent.ainvoke(
        {"messages": [HumanMessage(content="Search for LangGraph news")]},
        config={"callbacks": [handler]},
    )

[!NOTE] LangChain's callback abstraction does not expose raw HTTP response headers, so x-infralo-span-id cannot be extracted from on_llm_end automatically. Tool spans are always captured. Full LLM↔tool linkage is available when you inject the span ID into llm_output via a custom ChatOpenAI subclass or middleware.

---

CrewAI

from infralo import Infralo
from infralo.integrations.crewai import InfraloCrewAIListener
from crewai import Agent, Task, Crew

infralo = Infralo(api_key="vk_...")

# Register once at module level — self-registers on the CrewAI event bus
InfraloCrewAIListener()

with infralo.start_trace(session_id="user-123") as trace:
    crew = Crew(agents=[...], tasks=[...])
    crew.kickoff()

The listener handles ToolUsageStartedEvent, ToolUsageFinishedEvent, and ToolUsageErrorEvent — no changes to your crew, agents, or task definitions.

---

Google ADK

from infralo import Infralo
from infralo.integrations.google_adk import make_google_adk_callbacks
from google.adk.agents import Agent
from google.adk.models.lite_llm import LiteLlm
from google.adk.runners import Runner

infralo = Infralo(api_key="vk_...")
callbacks = make_google_adk_callbacks(infralo)

with infralo.start_trace(session_id="user-123") as trace:
    agent = Agent(
        name="my_agent",
        model=LiteLlm(
            model="openai/my-deployment",
            api_base="http://localhost:8000/v1",
        ),
        tools=[search_web, get_weather],
        **callbacks,   # injects before/after_tool + after_model
    )
    runner = Runner(agent=agent, ...)
    async for event in runner.run_async(...):
        ...

make_google_adk_callbacks() returns the same three hook keys as the Agno integration (before_tool_callback, after_tool_callback, after_model_callback).