Launch quickstart

staso[openai]==0.5.15
openai==3.6.0

export STASO_API_KEY=ak_...
export OPENAI_API_KEY=sk-...
export STASO_AGENT_NAME=launch-quickstart-agent
export STASO_ENVIRONMENT=production
export STASO_WORKSPACE_SLUG=default
export STASO_SESSION_ID=launch-quickstart-first-trace

from __future__ import annotations

import json
import os
from typing import Final

from openai import OpenAI

import staso as st

REQUIRED_ENV: Final = (
    "STASO_API_KEY",
    "OPENAI_API_KEY",
    "STASO_AGENT_NAME",
    "STASO_ENVIRONMENT",
    "STASO_WORKSPACE_SLUG",
    "STASO_SESSION_ID",
)

TOOLS: Final = [
    {
        "type": "function",
        "function": {
            "name": "lookup_order",
            "description": "Look up the current status for a sample order.",
            "parameters": {
                "type": "object",
                "properties": {
                    "order_id": {
                        "type": "string",
                        "description": "The order ID to look up.",
                    }
                },
                "required": ["order_id"],
                "additionalProperties": False,
            },
        },
    }
]


def _load_settings() -> dict[str, str]:
    missing = [name for name in REQUIRED_ENV if not os.environ.get(name)]
    if missing:
        raise RuntimeError("Missing required environment variables: " + ", ".join(sorted(missing)))
    return {name: os.environ[name] for name in REQUIRED_ENV}


@st.tool(name="lookup_order")
def lookup_order(order_id: str) -> dict[str, str]:
    return {
        "order_id": order_id,
        "status": "paid",
        "shipping_state": "ready_to_ship",
    }


@st.agent(name="launch-quickstart-agent")
def answer_order_question(order_id: str = "order_123") -> dict[str, object]:
    client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
    try:
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {
                    "role": "user",
                    "content": f"Look up the shipping status for order {order_id}.",
                }
            ],
            tools=TOOLS,
            tool_choice={"type": "function", "function": {"name": "lookup_order"}},
        )
    except st.GuardBlocked as blocked:
        return {
            "status": "blocked",
            "tool": blocked.tool_name,
            "reason": blocked.reason,
        }

    tool_calls = response.choices[0].message.tool_calls or []
    if len(tool_calls) != 1:
        raise RuntimeError("Expected one lookup_order tool proposal from OpenAI.")

    tool_call = tool_calls[0]
    arguments = json.loads(tool_call.function.arguments)
    result = lookup_order(**arguments)
    return {"status": "observed", "tool_result": result}


def main() -> None:
    settings = _load_settings()
    session_id = settings["STASO_SESSION_ID"]
    st.init(
        api_key=settings["STASO_API_KEY"],
        agent_name=settings["STASO_AGENT_NAME"],
        environment=settings["STASO_ENVIRONMENT"],
        workspace_slug=settings["STASO_WORKSPACE_SLUG"],
    )

    try:
        result = answer_order_question()
        print(json.dumps(result, sort_keys=True))
    finally:
        delivery = st.shutdown(timeout=10)
        if not delivery.ok:
            raise RuntimeError(
                "Trace delivery incomplete: "
                f"last_failure={delivery.last_failure!r}, "
                f"failed={delivery.failed}, "
                f"pending={delivery.pending}, "
                f"rejected={delivery.rejected}, "
                f"last_rejection={delivery.last_rejection!r}"
            )

    print(
        "Staso accepted trace delivery for "
        f"session_id={session_id}. The dashboard/backend trace query may take a few "
        "seconds to return the durable trace."
    )


if __name__ == "__main__":
    main()