Metadata-Version: 2.4
Name: infralo
Version: 0.1.0.dev2
Summary: Infralo SDK
License: MIT
Keywords: agents,ai,llm,observability,tracing
Requires-Python: >=3.10
Requires-Dist: httpx<1.0,>=0.28.0
Provides-Extra: agno
Requires-Dist: agno>=2.0; extra == 'agno'
Provides-Extra: crewai
Requires-Dist: crewai>=1.0.0; extra == 'crewai'
Provides-Extra: google-adk
Requires-Dist: google-adk>=2.0.0; extra == 'google-adk'
Provides-Extra: langgraph
Requires-Dist: langchain-core>=1.0.0; extra == 'langgraph'
Requires-Dist: langgraph>=1.0.0; extra == 'langgraph'
Description-Content-Type: text/markdown

# 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

```bash
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:

```python
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

```python
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.

```python
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.**

```python
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**.

```python
@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.

```python
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.

```python
@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):

```python
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

```python
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

```python
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)

```python
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

```python
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

```python
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

```python
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

```python
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

```python
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`).
