LiveKit Agents Deep Search

lk agent init my-agent --template agent-starter-python
cd my-agent
uv add httpx
echo 'SPECSTREAM_API_KEY=YOUR_SPECSTREAM_API_KEY' >> .env.local

import asyncio
import os

import httpx
from livekit.agents import RunContext, function_tool
from livekit.agents.llm import ToolError

specstream = httpx.AsyncClient(base_url="https://api.specstream.ai", timeout=60)
quotes = {}  # quote_id -> quote: the amount to approve comes from SpecStream, never from the model
watchers = set()  # background tasks that tell the caller when a table is ready
DONE = {"completed", "partial", "failed", "canceled", "expired"}


async def call(method: str, path: str, **kwargs) -> dict:
    headers = {"Authorization": f"Bearer {os.environ['SPECSTREAM_API_KEY']}"}
    r = await specstream.request(method, path, headers=headers, **kwargs)
    if r.is_error:
        try:
            message = r.json()["error"]["message"]
        except (ValueError, KeyError, TypeError):
            message = r.text[:300]
        raise ToolError(f"SpecStream answered {r.status_code}: {message}")
    return r.json()["data"]


@function_tool
async def price_deep_search(context: RunContext, request: str, rows: int) -> dict:
    """Price a table of rows from the live web: companies, places, people or other records
    that meet the caller's conditions, with the facts to bring back. Spends nothing.
    Read back how SpecStream understood the request and the maximum price, then ask the caller to approve it.

    Args:
        request: What each row is, the conditions it must meet and the facts to bring back, in the caller's words.
        rows: How many rows the caller wants, 1 to 100.
    """
    reading = await call("POST", "/v1/parse", json={"q": request})
    if reading["status"] != "runnable":
        return {
            "status": reading["status"],
            "read_back": reading.get("read_back"),
            "next": "Ask the caller for the missing detail, then price again.",
        }
    rows = max(1, min(int(rows), 100))  # the row limit holds even if the model asks for more
    quote = await call(
        "POST",
        "/v1/capacity/quotes",
        json={"ast_hash": reading["ast_hash"], "operation": "find", "requested_units": rows},
    )
    quotes[quote["quote_id"]] = quote
    return {
        "read_back": reading["read_back"],
        "rows": rows,
        "maximum_usd": quote["maximum_cents"] / 100,
        "quote_id": quote["quote_id"],
    }


@function_tool
async def start_deep_search(context: RunContext, quote_id: str) -> str:
    """Start a priced Deep Search. Call it only after the caller said yes to the maximum price you read out.
    The rows arrive later: tell the caller the search started and keep talking.

    Args:
        quote_id: The quote_id that price_deep_search returned.
    """
    quote = quotes.pop(quote_id, None)
    if quote is None:
        raise ToolError("Unknown or used quote. Price the request again.")
    run = await call(
        "POST",
        "/v1/capacity/executions",
        json={
            "quote_id": quote["quote_id"],
            "quote_hash": quote["quote_hash"],
            "confirm_max_cents": quote["maximum_cents"],
            "idempotency_key": quote["quote_id"],
        },
    )
    task = asyncio.create_task(tell_when_done(context.session, run["find_id"]))
    watchers.add(task)
    task.add_done_callback(watchers.discard)
    return f"The search started. Its find_id is {run['find_id']}."


@function_tool
async def read_deep_search(context: RunContext, find_id: str) -> str:
    """Read the first rows of a Deep Search started in this call, as a table with a Source column.

    Args:
        find_id: The find_id that start_deep_search returned.
    """
    find = await call("GET", f"/v1/finds/{find_id}")
    page = await call("GET", f"/v1/finds/{find_id}/rows", params={"format": "markdown", "limit": 10})
    return f"Status: {find['status']}. Rows so far: {page['accepted_count']}.\n\n{page['table']}"


async def tell_when_done(session, find_id: str) -> None:
    while True:
        await asyncio.sleep(10)
        try:
            find = await call("GET", f"/v1/finds/{find_id}")
        except (ToolError, httpx.HTTPError):
            continue  # a failed status check is tried again on the next pass
        if find["status"] in DONE:
            break
    await session.generate_reply(
        instructions=(
            f"Tell the caller the table for find {find_id} is ready with "
            f"{find['accepted_count']} rows, and offer to read the first ones."
        )
    )

from specstream_deep_search import price_deep_search, read_deep_search, start_deep_search

DEEP_SEARCH_RULES = """
# Deep Search

- When the caller wants many records with the same facts on each, call price_deep_search with their request and the number of rows.
- Read back how the request was understood and the maximum price in dollars, and wait for a clear yes.
- Call start_deep_search only after that yes. Never start a search the caller did not approve.
- When the table is ready, call read_deep_search and speak the first rows in short sentences, naming the site each row was read on.
"""


class Assistant(Agent):
    def __init__(self) -> None:
        super().__init__(
            llm=inference.LLM(model="google/gemma-4-31b-it"),
            instructions=textwrap.dedent(
                """\
                ...the starter's instructions, unchanged...
                """
            )
            + DEEP_SEARCH_RULES,
            tools=[price_deep_search, start_deep_search, read_deep_search],
        )

Find 20 coffee roasters in Portland, Oregon that sell wholesale, with their website and the city they roast in.