Trace a LiteLLM agent run

Add Staso to a Python agent that calls LiteLLM directly. The customer integration instruments litellm.completion() and litellm.acompletion() after st.init(). This guide uses LiteLLM's openai/gpt-4o-mini route, the upstream path tested with Staso. LiteLLM's use inside Staso's own backend is separate from this SDK integration.

Install and set keys

pip install "staso[litellm]"
export STASO_API_KEY="<your Staso API key>"
export OPENAI_API_KEY="<your OpenAI API key>"

The Staso key sends traces and requests Guard decisions. The OpenAI key pays for the upstream model calls in this example. Use a key scoped to the workspace you intend to trace.

Run a synchronous call

Save this as litellm_sync.py and run python litellm_sync.py:

import litellm
import staso as st

st.init(agent_name="support-agent")


@st.agent
def answer(question: str) -> str:
    response = litellm.completion(
        model="openai/gpt-4o-mini",
        messages=[{"role": "user", "content": question}],
    )
    return response.choices[0].message.content or ""


try:
    print(answer("What is the capital of Japan? One word."))
finally:
    st.shutdown()

The explicit @st.agent root groups the model call under one run. st.init() patches an installed LiteLLM package; you do not need to call patch_litellm() separately.

Run an asynchronous call

The same installation supports non-streaming acompletion():

import asyncio

import litellm
import staso as st

st.init(agent_name="async-support-agent")


@st.agent
async def answer(question: str) -> str:
    response = await litellm.acompletion(
        model="openai/gpt-4o-mini",
        messages=[{"role": "user", "content": question}],
    )
    return response.choices[0].message.content or ""


try:
    print(asyncio.run(answer("What is the capital of Japan? One word.")))
finally:
    st.shutdown()

Proposed tools and Guard

LiteLLM returns model-proposed tool calls. With an enabled Guard policy scoped to this agent, environment, and tool, Staso checks each supported proposal before the response reaches your dispatch code. Audit mode records a would-block finding and returns the proposal. Enforce mode raises st.GuardBlocked on a blocked proposal. No rule runs merely because the SDK is installed; configure a policy and verify its mode and scope first.

import json

import litellm
import staso as st

st.init(agent_name="weather-agent")

TOOLS = [{
    "type": "function",
    "function": {
        "name": "get_weather",
        "description": "Get the weather for a city.",
        "parameters": {
            "type": "object",
            "properties": {"city": {"type": "string"}},
            "required": ["city"],
            "additionalProperties": False,
        },
    },
}]


@st.tool
def get_weather(city: str) -> dict:
    return {"city": city, "temperature_c": 18}


@st.agent
def answer(question: str) -> str:
    messages = [{"role": "user", "content": question}]
    try:
        first = litellm.completion(
            model="openai/gpt-4o-mini",
            messages=messages,
            tools=TOOLS,
            tool_choice={"type": "function", "function": {"name": "get_weather"}},
        )
    except st.GuardBlocked as blocked:
        return f"Blocked by Guard: {blocked.reason}"

    calls = first.choices[0].message.tool_calls or []
    if not calls:
        return first.choices[0].message.content or ""

    call = calls[0]
    result = get_weather(**json.loads(call.function.arguments))
    messages.append(first.choices[0].message.model_dump(exclude_none=True))
    messages.append({
        "role": "tool",
        "tool_call_id": call.id,
        "content": json.dumps(result),
    })
    final = litellm.completion(model="openai/gpt-4o-mini", messages=messages)
    return final.choices[0].message.content or ""


try:
    print(answer("What's the weather in Paris?"))
finally:
    st.shutdown()

@st.tool records the Python function only when your host calls it. A proposed call is not evidence of execution. A direct Python tool call that does not come from a supported model proposal needs an explicit st.guard(...) check before execution if you want Guard enforcement. See manual checks.

What appears in Staso

A completed non-streaming call creates an LLM span under the agent root. A tool proposal records its ID and function name for correlation; the host-executed @st.tool call creates a separate tool span. Guard findings appear only when an active policy evaluates the proposal. Token usage appears when LiteLLM's response includes supported usage fields. Staso estimates cost only when its registry has an exact provider and resolved model price; otherwise cost is labeled Unknown or Partial. It does not treat missing usage as zero.

Capture controls can suppress request, system prompt, tool schema, tool arguments, and output content independently. Tool IDs and names can remain as operational correlation data. See SDK setup for st.init(capture_...) options.

Current limits

This guide covers direct Python completion() and acompletion() calls, non-streaming, through the tested OpenAI route. LiteLLM Proxy traffic and other providers have not been release-proven here.
stream=True is outside this integration's supported launch path. Do not infer streaming support from LiteLLM itself.
Guard can check the model's proposed call only when the adapter has usable tool-call evidence and a matching enabled policy. Missing or suppressed arguments/schema may cause a rule to skip or abstain.
Staso does not discover arbitrary Python tool execution. Decorate the host function with @st.tool for execution evidence; use st.guard(...) before a direct call when enforcement is required.
Guard runs inline and adds latency according to policy, provider, and network conditions. Fail-open is the default when a decision is unavailable; manual checks explain fail-closed configuration.