Workflows¶

import os
import sys

import dotenv
from loguru import logger

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

DEFAULT_LLM_MODEL = "openai/gpt-5.6-luna"
assert os.getenv("OPENAI_API_KEY"), "Set OPENAI_API_KEY to run the LLM nodes."

logger.remove()
_ = logger.add(sys.stderr, level="WARNING")

from kavalai.agent_service import AgentService
from kavalai.db import db_manager
from kavalai.workflow import SqliteTaskLogger, WorkflowEngine

# Create the runtime tables in an in-memory SQLite database.
await db_manager.init_sqlite()


def new_agent_service() -> AgentService:
    """An AgentService over the in-memory SQLite database."""
    return AgentService(db_manager.get_sqlite_sessionmaker())

greeter_yaml = """
name: Village greeter
description: Greets a visitor to Green Village.
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 the warm, concise greeter of Green Village (104 residents, one pub,
      one very old grain tower).
      Reply to the visitor 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,
    default_llm_model=DEFAULT_LLM_MODEL,
    agent_service=new_agent_service(),
)
state = await engine.run({"user_message": "Hi, I'm visiting from Tallinn!"})

import textwrap

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

print("status       :", state.status)
print("invocation_id:", state.invocation_id)
print("token_usage  :", state.token_usage)
print()
print("serialised   :", state.to_json()[:160], "…")

from uuid import UUID

history = await engine.agent_service.get_chat_history(
    UUID(state.session_id)
)
for message in history:
    print(textwrap.fill(f"{message.role:>9}: {message.content}", 80,
                        subsequent_indent=" " * 11))

from kavalai import crud
from kavalai.db import Run

async with engine.agent_service.session_maker() as session:
    run = await crud.get_one(session, Run, UUID(state.run_id))
    print("input_data :", run.input_data)
    print("output_data:", run.output_data)
    print("context    :", sorted(run.context))

from pydantic import BaseModel

from kavalai import pythontool


class PondReading(BaseModel):
    depth_m: float
    water: str


@pythontool
def measure_pond(name: str) -> PondReading:
    """Measure a Green Village pond. A real tool would read a sensor."""
    return PondReading(depth_m=1.2, water="murky but calm")


pond_yaml = """
name: Pond reporter
description: >-
  Reads a pond gauge with a function node, then phrases it with an LLM.
data_types:
  input:
    type: object
    properties:
      pond: {type: string}
  reading:
    type: object
    properties:
      depth_m: {type: number}
      water: {type: string}
  output:
    type: object
    properties:
      agent_response: {type: string}
nodes:
  - {name: start, type: start, next: measure}
  - name: measure
    type: function
    tool: python://measure_pond
    inputs:
      name: {type: context, value: input.pond}
    output: reading
    next: phrase
  - name: phrase
    type: llm
    prompt: >-
      Turn this pond reading into one cheerful sentence for the
      village noticeboard.
    inputs:
      reading: {type: context, value: reading}
    output: output
    next: end
  - {name: end, type: end, output: output}
"""

engine = WorkflowEngine.from_yaml(
    pond_yaml,
    default_llm_model=DEFAULT_LLM_MODEL,
    agent_service=new_agent_service())
engine.kernel.register_python_tool("measure_pond", measure_pond)

state = await engine.run({"pond": "Lake Miller"})
print("tool output in context:", state.data["reading"])
print(textwrap.fill("final: " + state.output_data["agent_response"], 80))

council_yaml = """
name: Council desk
description: Classifies a message to the village council and routes it.
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 villager's message as exactly one of: repair, permit, 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:
      repair: repair_reply
      permit: permit_reply
    default: general_reply
  - name: repair_reply
    type: llm
    prompt: >-
      Acknowledge the repair request and name the next step.
      One short paragraph.
    inputs: {input: {type: context, value: input}}
    output: output
    next: end
  - name: permit_reply
    type: llm
    prompt: >-
      Explain briefly how to apply for this permit in a village of
      104 people.
    inputs: {input: {type: context, value: input}}
    output: output
    next: end
  - name: general_reply
    type: llm
    prompt: "Answer the villager's question helpfully and briefly."
    inputs: {input: {type: context, value: input}}
    output: output
    next: end
  - {name: end, type: end, output: output}
"""

for message in [
    "The church bell has been stuck since Tuesday.",
    "May I put a beehive in my front garden?",
]:
    engine = WorkflowEngine.from_yaml(
        council_yaml,
        default_llm_model=DEFAULT_LLM_MODEL,
        agent_service=new_agent_service(),
    )
    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"][:80], "…\n")

from IPython.display import SVG, display

from kavalai.workflow import render_workflow_svg

display(SVG(render_workflow_svg(engine.graph)))

from kavalai import evaluate_bool, evaluate_expression

ctx = {
    "bell": {"broken_days": 3, "status": "stuck"},
    "ponds": [{"depth_m": 1.2}],
}

print(evaluate_expression(
    "bell.broken_days >= 3 and bell.status == 'stuck'", ctx
))
print(evaluate_expression("ponds[0].depth_m > 0.8", ctx))
print(evaluate_bool("bell.status in ['stuck', 'silent']", ctx))
print(evaluate_expression("missing.value", ctx))  # unknown path -> None

import datetime


@pythontool
def current_time() -> str:
    """The current UTC time, as a weekday and clock time."""
    return datetime.datetime.now(datetime.timezone.utc).strftime("%A %H:%M UTC")


agent_yaml = """
name: Time-aware greeter
description: An agent that checks the time before greeting the visitor.
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 visitor to Green Village appropriately for the current time of
      day. Check the time with a tool first; mention it.
    inputs: {input: {type: context, value: input}}
    output: output
    allowed_tools: ["python://current_time"]
    max_steps: 4
    next: end
  - {name: end, type: end, output: output}
"""

engine = WorkflowEngine.from_yaml(
    agent_yaml,
    default_llm_model=DEFAULT_LLM_MODEL,
    agent_service=new_agent_service())
engine.kernel.register_python_tool("current_time", current_time)

state = await engine.run({"user_message": "Say hello!"})
print(textwrap.fill(state.output_data["agent_response"], 80))

engine = WorkflowEngine.from_yaml(
    council_yaml,
    default_llm_model=DEFAULT_LLM_MODEL,
    agent_service=new_agent_service(),
)

report = {"user_message": "The pub sign fell down."}
async for event in engine.run_stream(report):
    detail = ""
    if event.type == "workflow_completed":
        detail = f"  tokens={event.token_usage['total_tokens']}"
    print(f"{event.type:<20} {event.name}{detail}")

task_logger = SqliteTaskLogger()
engine = WorkflowEngine.from_yaml(
    council_yaml,
    default_llm_model=DEFAULT_LLM_MODEL,
    agent_service=new_agent_service(), task_logger=task_logger,
)
state = await engine.run({"user_message": "How do I book the village hall?"})
await task_logger.flush()  # await the fire-and-forget writes

print("path       :", " → ".join(state.trace))
print("token usage:", state.token_usage)

for task in await task_logger.get_tasks():
    print(f"  {task['name']:<14} {task['node_type']:<8} "
          f"{task['duration_seconds']:.2f}s")

calls = await task_logger.get_model_calls()
tokens = sum(call["total_tokens"] or 0 for call in calls)
print(f"\nmodel calls: {len(calls)} · tokens: {tokens}")

from kavalai.testing import ScriptedLlmClient

scripted = ScriptedLlmClient(
    lambda messages, response_model: {
        name: ("repair" if name == "intent" else "Scripted reply.")
        for name in response_model.model_fields
    }
)

engine = WorkflowEngine.from_yaml(
    council_yaml,
    default_llm_model=DEFAULT_LLM_MODEL,
    agent_service=new_agent_service(),
    client_factory=scripted,
)
state = await engine.run({"user_message": "anything at all"})

print("deterministic path:", " → ".join(state.trace))
print("reply             :", state.output_data["agent_response"])
print("model calls       :", len(scripted.calls))

from kavalai import WorkflowBuilder

engine = (
    WorkflowBuilder("Council desk (code)",
                    description="The same graph, built in Python.")
    .data_type("input", {"user_message": str})
    .data_type("classification", {"intent": str})
    .data_type("output", {"agent_response": str})
    .start("classify")
    .llm(
        "classify",
        prompt=(
            "Classify the villager's message as exactly one of: "
            "repair, permit, other. "
            "Respond with that single lowercase word in the `intent` field."
        ),
        inputs={"input": "input"},          # a bare string == a context path
        output="classification",
        next="route",
    )
    .switch(
        "route",
        expr="classification.intent",
        cases={"repair": "repair_reply", "permit": "permit_reply"},
        default="general_reply",
    )
    .llm("repair_reply",
         prompt="Acknowledge the repair request and name the next step.",
         inputs={"input": "input"}, output="output", next="end")
    .llm("permit_reply", prompt="Explain briefly how to apply for this permit.",
         inputs={"input": "input"}, output="output", next="end")
    .llm("general_reply",
         prompt="Answer the villager's question helpfully and briefly.",
         inputs={"input": "input"}, output="output", next="end")
    .end()
    .build_engine(
        default_llm_model=DEFAULT_LLM_MODEL,
        agent_service=new_agent_service(),
    )
)

state = await engine.run({"user_message": "The grain tower steps are loose."})
print("path :", " → ".join(state.trace))
print("reply:", state.output_data["agent_response"][:90], "…")