Building agents with Workflows¶

import os
import dotenv

dotenv.load_dotenv("../.env")

os.environ["KAVALAI_DEFAULT_LLM_MODEL"] = "openai/gpt-4o-mini"
assert os.getenv("OPENAI_API_KEY"), "Set OPENAI_API_KEY to run the LLM nodes."

from kavalai.workflow import (
    WorkflowEngine,
    SqliteDataStorage,
    SqliteTaskLogger,
)

greeter_yaml = """
name: Greeter
description: A one-node agent that greets the user by name.
data_types:
  input:
    type: object
    properties:
      user_message: {type: string}
  output:
    type: object
    properties:
      agent_response: {type: string}
nodes:
  - {name: start, type: start, next: reply}
  - name: reply
    type: llm
    prompt: |
      You are a warm, concise greeter.
      Reply to the user in one friendly sentence: {{ context.input.user_message }}
    inputs:
      input: {type: context, value: input}
    output: output
    next: end
  - {name: end, type: end, output: output}
"""

engine = WorkflowEngine.from_yaml(
    greeter_yaml,
    storage=SqliteDataStorage(),
    task_logger=SqliteTaskLogger(),
)
state = await engine.run({"user_message": "Hi, I'm Timo!"})

print(state.output_data["agent_response"])
print("path:", " → ".join(state.trace))

# The entire state serialises to JSON (and back via WorkflowState.from_json).
print(state.to_json()[:240], "…\n")

# It was checkpointed to storage after each node — reload it by run id.
reloaded = await engine.storage.load_state(state.run_id)
print("reloaded status:", reloaded.status, "| nodes:", reloaded.trace, "\n")

# The same backend recorded the conversation turns.
for msg in await engine.storage.get_chat_history(state.session_id):
    print(f"{msg.role:>9}: {msg.content}")

from pydantic import BaseModel
from kavalai.functionkernel import pythontool


class Weather(BaseModel):
    summary: str
    celsius: int


@pythontool
def get_weather(city: str) -> Weather:
    # A real tool would call a weather API; we fake it for the demo.
    return Weather(summary=f"clear skies over {city}", celsius=24)


weather_yaml = """
name: Weather forecaster
description: Looks up weather with a function node, then phrases it with an LLM.
data_types:
  input:
    type: object
    properties:
      user_message: {type: string}
      city: {type: string}
  weather:
    type: object
    properties:
      summary: {type: string}
      celsius: {type: integer}
  output:
    type: object
    properties:
      agent_response: {type: string}
nodes:
  - {name: start, type: start, next: fetch}
  - name: fetch
    type: function
    tool: python://get_weather
    inputs:
      city: {type: context, value: input.city}
    output: weather
    next: phrase
  - name: phrase
    type: llm
    prompt: "Turn the weather reading into a cheerful one-line forecast."
    inputs:
      weather: {type: context, value: weather}
    output: output
    next: end
  - {name: end, type: end, output: output}
"""

engine = WorkflowEngine.from_yaml(weather_yaml, storage=SqliteDataStorage())
engine.kernel.register_python_tool("get_weather", get_weather)

state = await engine.run({"user_message": "weather?", "city": "Tallinn"})
print("tool output kept in context:", state.data["weather"])
print("final:", state.output_data["agent_response"])

router_yaml = """
name: Support router
description: Classifies a support request and routes to a specialised reply.
data_types:
  input:
    type: object
    properties:
      user_message: {type: string}
  classification:
    type: object
    properties:
      intent: {type: string}
  output:
    type: object
    properties:
      agent_response: {type: string}
nodes:
  - {name: start, type: start, next: classify}
  - name: classify
    type: llm
    prompt: |
      Classify the user's request as exactly one of: refund, technical, other.
      Respond with that single lowercase word in the `intent` field.
    inputs:
      input: {type: context, value: input}
    output: classification
    next: route
  - name: route
    type: switch
    expr: classification.intent
    cases:
      refund: refund_reply
      technical: tech_reply
    default: general_reply
  - name: refund_reply
    type: llm
    prompt: "Empathetically acknowledge the refund request and outline next steps."
    inputs: {input: {type: context, value: input}}
    output: output
    next: end
  - name: tech_reply
    type: llm
    prompt: "Offer one concrete first troubleshooting step for the technical issue."
    inputs: {input: {type: context, value: input}}
    output: output
    next: end
  - name: general_reply
    type: llm
    prompt: "Answer the user's question helpfully and briefly."
    inputs: {input: {type: context, value: input}}
    output: output
    next: end
  - {name: end, type: end, output: output}
"""

for message in ["I want my money back!", "The app crashes on launch"]:
    engine = WorkflowEngine.from_yaml(router_yaml)
    state = await engine.run({"user_message": message})
    print(repr(message))
    print("  intent:", state.data["classification"]["intent"])
    print("  path:  ", " → ".join(state.trace))
    print("  reply: ", state.output_data["agent_response"][:70], "…\n")

from kavalai.workflow import evaluate_expression, evaluate_bool

ctx = {"state": {"retries": 3, "status": "ok"}, "items": [{"score": 0.9}]}
print(evaluate_expression("state.retries >= 3 and state.status == 'ok'", ctx))
print(evaluate_expression("items[0].score > 0.8", ctx))
print(evaluate_bool("state.status in ['ok', 'done']", ctx))
print(evaluate_expression("missing.value", ctx))  # unknown path -> None

import datetime


@pythontool
def current_time() -> str:
    return datetime.datetime.now(datetime.timezone.utc).strftime("%A %H:%M UTC")


agent_yaml = """
name: Time-aware assistant
description: An agent that checks the time before greeting the user.
data_types:
  input:
    type: object
    properties:
      user_message: {type: string}
  output:
    type: object
    properties:
      agent_response: {type: string}
nodes:
  - {name: start, type: start, next: assistant}
  - name: assistant
    type: agent
    prompt: "Greet the user appropriately for the current time of day. Use a tool to check the time."
    inputs: {input: {type: context, value: input}}
    output: output
    max_steps: 4
    next: end
  - {name: end, type: end, output: output}
"""

engine = WorkflowEngine.from_yaml(agent_yaml)
engine.kernel.register_python_tool("current_time", current_time)
state = await engine.run({"user_message": "Say hello!"})
print(state.output_data["agent_response"])

storage = SqliteDataStorage()
tasklog = SqliteTaskLogger()
engine = WorkflowEngine.from_yaml(router_yaml, storage=storage, task_logger=tasklog)
state = await engine.run({"user_message": "How do I reset my password?"})
await tasklog.flush()  # await the fire-and-forget writes

conn = await tasklog._connect()
print("Nodes executed:")
async with conn.execute(
    "SELECT name, node_type, duration_seconds FROM tasks ORDER BY rowid"
) as cur:
    for name, ntype, dur in await cur.fetchall():
        print(f"  {name:<14} {ntype:<8} {dur:.2f}s")

async with conn.execute("SELECT model, total_tokens FROM model_call_stats") as cur:
    rows = await cur.fetchall()
print(f"\nModel calls: {len(rows)} · tokens: {sum(r[1] or 0 for r in rows)}")

from kavalai.llm_clients.base_client import BaseLlmClient


class StubClient(BaseLlmClient):
    """Returns canned structured output — no network, fully deterministic."""

    def __init__(self, *args, **kwargs):
        super().__init__()

    async def chat_completions(self, *, chat_history, response_model=None):
        values = {
            name: ("refund" if name == "intent" else "Stubbed reply.")
            for name in response_model.model_fields
        }
        return response_model(**values)


engine = WorkflowEngine.from_yaml(
    router_yaml, client_factory=lambda *a, **k: StubClient()
)
state = await engine.run({"user_message": "anything at all"})
print("deterministic path:", " → ".join(state.trace))  # always routes via 'refund'

from kavalai.workflow import WorkflowBuilder

engine = (
    WorkflowBuilder("Support router (code)", description="Built with the builder.")
    .data_type("input", {"user_message": str})
    .data_type("classification", {"intent": str})
    .data_type("output", {"agent_response": str})
    .start("classify")
    .llm(
        "classify",
        prompt=(
            "Classify the request as exactly one of: refund, technical, other. "
            "Respond with that single lowercase word in the `intent` field."
        ),
        inputs={"input": "input"},          # a string input == a context path
        output="classification",
        next="route",
    )
    .switch(
        "route",
        expr="classification.intent",
        cases={"refund": "refund_reply", "technical": "tech_reply"},
        default="general_reply",
    )
    .llm("refund_reply", prompt="Acknowledge the refund and outline next steps.",
         inputs={"input": "input"}, output="output", next="end")
    .llm("tech_reply", prompt="Offer one concrete first troubleshooting step.",
         inputs={"input": "input"}, output="output", next="end")
    .llm("general_reply", prompt="Answer the question helpfully and briefly.",
         inputs={"input": "input"}, output="output", next="end")
    .end()
    .build_engine()                          # -> a ready WorkflowEngine
)

state = await engine.run({"user_message": "My app keeps crashing on launch"})
print("path:         ", " → ".join(state.trace))
print("reply:        ", state.output_data["agent_response"][:70], "…")
print("invocation_id:", state.invocation_id)
print("token_usage:  ", state.token_usage)