LiteLLM demo: trace a Python tool call

This LiteLLM demo sends a question to a model, executes one local Python function, and sends the function's result back for a final answer. Staso records the model calls and the executed function under one agent run, so you can check what actually happened between the two responses.

The order is fictional and lives in memory. Running the script makes real OpenAI requests and sends trace data to your Staso workspace. It does not connect to an order system or change customer data. The final answer is generated by your model call; there is no prerecorded model response in this example.

What the example demonstrates

LiteLLM's function-calling documentation describes the request, dispatch, and continuation steps. A returned tool_calls entry is a proposal. Your application still has to choose whether to execute the named function and return its result.

This uses direct, non-streaming litellm.completion() through openai/gpt-4o-mini, the provider route tested with Staso. Staso's LiteLLM extra selects the supported dependency range. LiteLLM Proxy, streaming, and its newer agent harness are outside this example's supported configuration.

Install and select a workspace

Use Python 3.11 or newer in a virtual environment:

python3 -m venv .venv
source .venv/bin/activate
python3 -m pip install "staso[litellm]"

Set OPENAI_API_KEY, STASO_API_KEY, and STASO_WORKSPACE_SLUG in your environment. The Staso key must belong to the workspace you will inspect and have traces:write. Evaluating Guard policies also needs guard:evaluate. The OpenAI key pays for the two model requests on the successful path.

Use the sample order below while testing capture. Request and output content can become trace data; the capture controls let you choose what to retain before using real customer inputs.

Run the complete function-calling demo

Save this as litellm_demo.py. The application accepts only lookup_order with one string argument. It rejects an unexpected proposal before dispatch, and the function only recognizes the sample order.

import json

import litellm
import staso as st

st.init(agent_name="litellm-demo", environment="development")

MODEL = "openai/gpt-4o-mini"
TOOLS = [{
    "type": "function",
    "function": {
        "name": "lookup_order",
        "description": "Look up the fictional demo order's status.",
        "parameters": {
            "type": "object",
            "properties": {"order_id": {"type": "string"}},
            "required": ["order_id"],
            "additionalProperties": False,
        },
    },
}]


@st.tool
def lookup_order(order_id: str) -> dict[str, str]:
    if order_id != "demo-1001":
        raise ValueError("Only the fictional order demo-1001 is available")
    return {"order_id": order_id, "status": "shipped"}


@st.agent
def answer(question: str) -> str:
    messages = [{"role": "user", "content": question}]
    first = litellm.completion(
        model=MODEL,
        messages=messages,
        tools=TOOLS,
        tool_choice={"type": "function", "function": {"name": "lookup_order"}},
        parallel_tool_calls=False,
    )
    proposal = first.choices[0].message
    calls = proposal.tool_calls or []
    if len(calls) != 1 or calls[0].function.name != "lookup_order":
        raise ValueError("Expected one lookup_order proposal")
    call = calls[0]
    arguments = json.loads(call.function.arguments)
    if (
        not isinstance(arguments, dict)
        or set(arguments) != {"order_id"}
        or not isinstance(arguments["order_id"], str)
    ):
        raise ValueError("Expected one string order_id argument")

    print("Proposed:", call.function.name)
    result = lookup_order(**arguments)
    print("Executed:", json.dumps(result))
    messages.append(proposal.model_dump(exclude_none=True))
    messages.append({
        "role": "tool",
        "tool_call_id": call.id,
        "content": json.dumps(result),
    })
    final = litellm.completion(model=MODEL, messages=messages)
    return final.choices[0].message.content or ""


try:
    print(answer("What is the status of order demo-1001?"))
except st.GuardBlocked as blocked:
    print("Stopped by Guard before dispatch:", blocked.reason)
finally:
    delivery = st.shutdown()
    if not delivery.ok:
        print("Trace delivery incomplete:", delivery.last_failure)

Run python3 litellm_demo.py. On a successful allowed run, the script prints the proposed function name, the executed lookup result, and the model's answer. The order result is always the local sample; the wording of the final answer can vary. A provider error, unexpected proposal, or policy block can end the run before the second model request.

Verify the execution in Staso

Open Traces > Traces in the selected workspace and find the new litellm-demo run. Inspect the answer root, both LLM calls, and the lookup_order execution. The tool span is the evidence that the Python function ran. A proposed call in the first LLM response alone does not establish that.

Token usage depends on the fields returned by LiteLLM. Cost also requires an exact known provider and model price. Missing evidence stays Unknown or Partial; it is not proof that the request cost zero. Check the captured result rather than expecting a particular token count or latency from this guide.

st.shutdown() attempts to drain pending events and reports delivery problems. Successful delivery is not proof that the run is queryable yet. Confirm the trace in the intended workspace before treating setup as complete.

Where a Guard policy changes the flow

With a matching enabled policy, Staso's supported LiteLLM adapter checks the proposed call before returning it to this dispatch code. Audit mode records a would-block finding and allows the action. Enforce mode can raise st.GuardBlocked before lookup_order executes, where the current plan permits enforcement. Installation alone does not create a matching policy.

The example's name and argument checks validate the local dispatch contract. They do not decide whether an action matches the user's broader intent. Configure a policy for agent litellm-demo, environment development, and tool lookup_order when testing that decision. Review the Guard setup and failure behavior; fail-open is the default when evaluation is unavailable.

@st.tool records execution. Calling that function directly from another path requires an explicit st.guard(...) check when enforcement is needed. The tool-call guardrails walkthrough demonstrates that separate boundary.

Troubleshoot the missing step

For the asynchronous setup and the full integration limits, continue with the LiteLLM SDK guide.