import asyncio
import json
import time
from typing import Any, AsyncGenerator, Callable, Mapping, Optional, Sequence, Type
import os
from loguru import logger
from pydantic import BaseModel, Field, ConfigDict
from kavalai.run_context import RunContext
from kavalai.utils import to_plain
from kavalai.functionkernel import FunctionKernel, _is_tool_allowed
from kavalai.llm_clients.base_client import BaseLlmClient, ChatHistory, ChatMessage
from kavalai.llm_clients.common import safe_parse_json
from kavalai.llm_clients.streamer import StreamContent
from jinja2 import Template
[docs]
def get_step_output_type(ResponseModel: Type[BaseModel]):
class StepOutput(BaseModel):
"""Data structure that helps passing around information between consecutive agent runs."""
instructions: str = Field(
description="Briefly describe the goal of this step (what you intend to achieve with these tool calls).",
)
tool_calls: list[ToolCall] = Field(
default_factory=list,
description="Add tool call requests here, their output will be available via `call_id` key for next steps.",
)
output: Optional[ResponseModel] = None
return StepOutput
[docs]
class StepStreamDemuxer:
"""Routes one agent step's raw LLM stream into named sub-streams.
Consumes the step's :class:`StreamContent` chunks (the step is always
streamed in raw-delta mode) and returns the chunks to emit, gated by the
stream flags:
- the step's raw model output under ``step<N>`` (``stream_partials``),
- the ``instructions`` field under ``instructions``, with a per-step
``complete`` (``stream_instructions``),
- the ``output`` field under ``response`` as full safe-parsed JSON
(``stream_output``),
- pass-through of auxiliary provider streams (e.g. Gemini thoughts) and
of ``restart`` chunks, which also reset the accumulated state.
One instance handles exactly one step; the raw text accumulates in
``buffer`` and the completed step's JSON is exposed as ``step_json``.
"""
def __init__(
self,
step_idx: int,
*,
stream_output: bool = False,
stream_instructions: bool = False,
stream_partials: bool = False,
stream_delta: bool = False,
):
self.step_idx = step_idx
self.stream_output = stream_output
self.stream_instructions = stream_instructions
self.stream_partials = stream_partials
self.stream_delta = stream_delta
self.buffer = ""
self.step_json: Optional[str] = None
self._sent_instructions = ""
self._sent_output_json: Optional[str] = None
[docs]
def on_chunk(self, chunk: StreamContent) -> list[StreamContent]:
"""Process one incoming chunk and return the chunks to emit."""
if chunk.type == "restart":
self.buffer = ""
self._sent_instructions = ""
self._sent_output_json = None
return [chunk]
if chunk.name != "response":
# Auxiliary provider streams (e.g. Gemini thoughts).
return [chunk] if self.stream_output else []
if chunk.type == "partial":
return self._on_partial(chunk)
if chunk.type == "complete":
self.step_json = self.buffer
return []
[docs]
def step_completed(self, step_output) -> list[StreamContent]:
"""Per-step closing chunks, once the parsed ``StepOutput`` is known."""
events: list[StreamContent] = []
if self.stream_instructions and step_output.instructions:
events.append(
StreamContent(
type="complete",
name="instructions",
value=step_output.instructions,
)
)
if self.stream_partials:
events.append(
StreamContent(
type="complete",
name=f"step{self.step_idx}",
value=None if self.stream_delta else self.step_json,
)
)
return events
def _on_partial(self, chunk: StreamContent) -> list[StreamContent]:
self.buffer += chunk.value or ""
events: list[StreamContent] = []
if self.stream_partials:
events.append(
StreamContent(
type="partial",
name=f"step{self.step_idx}",
value=chunk.value if self.stream_delta else self.buffer,
)
)
if self.stream_instructions or self.stream_output:
parsed = safe_parse_json(self.buffer)
if isinstance(parsed, dict):
events += self._demux_fields(parsed)
return events
def _demux_fields(self, parsed: dict) -> list[StreamContent]:
"""Route the fields of the partially parsed ``StepOutput`` to streams."""
events: list[StreamContent] = []
instructions = parsed.get("instructions")
if (
self.stream_instructions
and isinstance(instructions, str)
and instructions != self._sent_instructions
):
if self.stream_delta and instructions.startswith(self._sent_instructions):
value = instructions[len(self._sent_instructions) :]
else:
if self.stream_delta and self._sent_instructions:
# Non-monotonic change: tell the client to start over.
events.append(StreamContent(type="restart", name="instructions"))
value = instructions
if value:
events.append(
StreamContent(type="partial", name="instructions", value=value)
)
self._sent_instructions = instructions
# The output field streams under the main stream name as full JSON;
# an output produced alongside tool calls is forwarded as-is and may
# be superseded by the final complete.
output_value = parsed.get("output")
if self.stream_output and output_value is not None:
out_json = (
output_value
if isinstance(output_value, str)
else json.dumps(output_value)
)
if out_json != self._sent_output_json:
events.append(
StreamContent(type="partial", name="response", value=out_json)
)
self._sent_output_json = out_json
return events
def _default_template(filename: str) -> Template:
with open(os.path.join(os.path.dirname(__file__), filename), "r") as f:
return Template(f.read())
[docs]
class Agent:
def __init__(
self,
llm_client: BaseLlmClient,
*,
kernel: Optional[FunctionKernel] = None,
run_context: Optional[RunContext] = None,
prompt_template: Optional[Template] = None,
step_template: Optional[Template] = None,
allowed_tools: Optional[list[str]] = None,
on_step: Optional[Callable[[dict], None]] = None,
debug: bool = False,
):
"""Build an agent.
Every step sends one conversation to the model: a system message
rendered from ``prompt_template`` (the task, the context variables and
the tool descriptions), then the chat ``history`` handed to
:meth:`prompt_stream`, then a user message rendered from
``step_template`` (the request this invocation answers, the steps
executed so far and the instruction to produce the next one). Both
templates receive the same variables: ``prompt``, ``data``,
``tool_descriptions``, ``steps``, ``current_step``, ``max_steps``,
``last_step`` (whether this is the final step, on which tool calls are
not run) and ``request``. The defaults are
``default_prompt_template.j2`` and ``default_step_template.j2``
beside this module.
``allowed_tools`` restricts the agent to a subset of the kernel's
tools, given as tool URIs (``python://web.crawl``, or ``rest://api.*``
for a whole server). The tools outside it are neither described to the
model nor callable. ``None`` (the default) allows every registered
tool; an empty list allows none.
``on_step`` is an observer called once per completed reasoning step
with that step's record (index, instructions, tool calls with their
arguments, results and durations, and the step's output). It exists so
an agent's internal tool calls can be recorded — the workflow engine
turns them into task rows — without putting them on the public stream,
which every SSE consumer sees. It is called synchronously and its
exceptions are swallowed: a broken observer must never break a run.
"""
self.debug = debug
self.kernel = kernel
self.allowed_tools = allowed_tools
self.on_step = on_step
if not run_context:
run_context = RunContext()
self.run_context = run_context
self.llm_client = llm_client
self.prompt_template = prompt_template or _default_template(
"default_prompt_template.j2"
)
self.step_template = step_template or _default_template(
"default_step_template.j2"
)
[docs]
async def prompt_stream(
self,
prompt: str,
response_model: Optional[Type[BaseModel]] = None,
max_steps: int = 10,
*,
stream_output: bool = False,
stream_instructions: bool = False,
stream_partials: bool = False,
stream_delta: bool = False,
history: Optional[Sequence[ChatMessage]] = None,
inputs: Optional[Mapping[str, Any]] = None,
request: Optional[str] = None,
) -> AsyncGenerator[StreamContent, None]:
"""Run the agent loop, streaming progress as :class:`StreamContent`.
The agent iterates up to ``max_steps`` times. On each step the LLM
returns a ``StepOutput`` with optional ``tool_calls`` and an optional
final ``output``. Tool calls are executed through the
``FunctionKernel`` and their results are fed back into the prompt so
the model can reason over them on the next step. The loop stops once
the model returns an ``output`` without requesting further tool calls,
or when ``max_steps`` is reached.
The final chunk is always ``complete``/``response`` carrying the final
output (JSON for structured outputs, the plain string otherwise, or
``None`` when no output was produced). The flags gate the optional
progress streams; their semantics (including the naming and the
full-JSON structured output stream) are documented on
:class:`kavalai.workflow.models.AgentNode`. Stream names here are
unscoped (``response``, ``instructions``, ``step<N>``); the workflow
engine prefixes them with the node name.
Args:
prompt: The task description for the agent.
response_model: Optional Pydantic model describing the structured
final output. When omitted, a plain string is produced.
max_steps: Maximum number of reasoning/tool-calling iterations.
stream_output: Stream the step's ``output`` field as it is written.
stream_instructions: Stream each step's ``instructions`` field.
stream_partials: Stream each step's raw model output.
stream_delta: Emit deltas instead of full accumulated values on the
``instructions`` and ``step<N>`` streams.
history: Earlier turns of the conversation, sent with every step
between the system prompt and the step message. The current
user message belongs in ``prompt``, not here — the model
would read it twice.
inputs: The context variables of this invocation: listed in the
system prompt and the values ``input_args`` of a tool call
resolve against. Defaults to the run context's ``data``. The
workflow engine passes the node's resolved ``inputs``, so an
agent node sees what its author declared and not every value
the run has produced so far.
request: The request this invocation answers, stated in every
step message. ``history`` ends on the turn before it, so
without it a model reading the conversation answers that turn.
The workflow engine passes the run's user message.
"""
StepOutput = get_step_output_type(response_model or str)
history_messages = list(history or [])
variables = dict(inputs) if inputs is not None else self.run_context.data
# Per-invocation working memory: tool call results keyed by call_id,
# referenced via `planner_context_args`. Created fresh for each
# invocation (up to `max_steps`) and discarded afterwards, unlike
# `self.run_context` which is passed in at construction.
planner_context: dict[str, Any] = {}
# Record of executed steps, rendered back into the prompt template.
steps: list[dict] = []
final_output: Optional[BaseModel] = None
for step_idx in range(max_steps):
template_vars = dict(
prompt=prompt,
data=variables,
tool_descriptions=(
await self.kernel.get_tool_descriptions(self.allowed_tools)
if self.kernel
else ""
),
steps=steps,
current_step=step_idx,
max_steps=max_steps,
last_step=step_idx == max_steps - 1,
request=request,
)
rendered_prompt = self.prompt_template.render(**template_vars)
rendered_step = self.step_template.render(**template_vars)
if self.debug:
print(rendered_prompt)
print(rendered_step)
chat_history = ChatHistory(
messages=[
ChatMessage(role="system", content=rendered_prompt),
*history_messages,
ChatMessage(role="user", content=rendered_step),
]
)
logger.info(f"Agent step {step_idx}/{max_steps}")
# Always consume the step in raw-delta mode: the demuxer needs the
# full raw text to safe-parse, and its buffer doubles as the value
# to validate on completion.
streamer = await self.llm_client.stream_chat_completions(
chat_history=chat_history,
response_model=StepOutput,
stream_delta=True,
)
demux = StepStreamDemuxer(
step_idx,
stream_output=stream_output,
stream_instructions=stream_instructions,
stream_partials=stream_partials,
stream_delta=stream_delta,
)
async for chunk in streamer:
for event in demux.on_chunk(chunk):
yield event
if demux.step_json is None or not demux.step_json.strip():
logger.warning("LLM returned no step output, stopping.")
break
step_output = StepOutput.model_validate(safe_parse_json(demux.step_json))
for event in demux.step_completed(step_output):
yield event
# Ensure every tool call has a stable id for context lookups.
for idx, tool_call in enumerate(step_output.tool_calls):
if not tool_call.call_id:
tool_call.call_id = f"tool_call_{step_idx}_{idx}"
step_record: dict[str, Any] = {
"index": step_idx,
"instructions": step_output.instructions,
"tool_calls": [],
"output": to_plain(step_output.output),
}
# The last step has no next step to read a result in: its tool calls
# are not run, and the step message says so, so the model answers.
if step_idx == max_steps - 1 and step_output.tool_calls:
logger.warning(
f"Agent asked for {len(step_output.tool_calls)} tool call(s) on its last step; not run"
)
elif step_output.tool_calls and self.kernel:
results = await asyncio.gather(
*[
self._call_tool(tc, planner_context, variables)
for tc in step_output.tool_calls
]
)
for tool_call, args, result, duration in results:
planner_context[tool_call.call_id] = result
step_record["tool_calls"].append(
{
"name": tool_call.name,
"args": args,
"call_id": tool_call.call_id,
"output": to_plain(result),
"duration_seconds": duration,
}
)
steps.append(step_record)
self._publish_step(step_record)
if step_output.output is not None:
final_output = step_output.output
# Stop once the model produced an answer without more tool calls.
if not step_output.tool_calls:
break
else:
if final_output is None:
logger.warning(
f"Agent used its {max_steps} steps without producing an output"
)
value = None
if final_output is not None:
value = (
final_output
if isinstance(final_output, str)
else final_output.model_dump_json()
)
yield StreamContent(type="complete", name="response", value=value)
[docs]
async def prompt(
self,
prompt: str,
response_model: Optional[Type[BaseModel]] = None,
max_steps: int = 10,
*,
history: Optional[Sequence[ChatMessage]] = None,
inputs: Optional[Mapping[str, Any]] = None,
request: Optional[str] = None,
) -> str | BaseModel:
"""Run the agent loop and return the final output (blocking wrapper).
Drains :meth:`prompt_stream` with all progress streams disabled.
Args:
prompt: The task description for the agent.
response_model: Optional Pydantic model describing the structured
final output. When omitted, a plain string is returned.
max_steps: Maximum number of reasoning/tool-calling iterations.
history: Earlier turns of the conversation, as on
:meth:`prompt_stream`.
inputs: The context variables, as on :meth:`prompt_stream`.
request: The request this invocation answers, as on
:meth:`prompt_stream`.
Returns:
The structured ``response_model`` instance, or a string when no
``response_model`` is provided. ``None`` if no output was produced.
"""
final_value = None
async for chunk in self.prompt_stream(
prompt,
response_model=response_model,
max_steps=max_steps,
history=history,
inputs=inputs,
request=request,
):
if chunk.type == "complete" and chunk.name == "response":
final_value = chunk.value
if final_value is None:
return None
if response_model:
return response_model.model_validate_json(final_value)
return final_value
def _resolve_args(
self,
tool_call: ToolCall,
planner_context: dict[str, Any],
variables: Optional[Mapping[str, Any]] = None,
) -> dict:
"""Resolve a ToolCall's argument sources into a single argument dict.
Arguments are merged with precedence ``literal_args`` >
``planner_context_args`` > ``input_args``. ``planner_context_args``
resolves against the per-invocation ``planner_context`` (results of
previous tool calls); ``input_args`` against ``variables``, the context
variables the prompt listed (``self.run_context.data`` when not given).
"""
if variables is None:
variables = self.run_context.data
def parse(field_name: str, value: str) -> dict:
if not value:
return {}
try:
return json.loads(value)
except json.JSONDecodeError:
logger.error(f"Failed to parse {field_name} as JSON: {value}")
return {}
literal_args = parse("literal_args", tool_call.literal_args)
context_keys = parse("planner_context_args", tool_call.planner_context_args)
context_args = {
arg_name: planner_context.get(ctx_key)
for arg_name, ctx_key in context_keys.items()
}
input_keys = parse("input_args", tool_call.input_args)
input_args = {
arg_name: variables.get(input_key)
for arg_name, input_key in input_keys.items()
}
return {**input_args, **context_args, **literal_args}
def _publish_step(self, step_record: dict) -> None:
"""Hand a completed step to the ``on_step`` observer, if there is one.
Never raises: an observer that fails is a logging problem, not a
reason to abandon the run.
"""
if self.on_step is None:
return
try:
self.on_step(step_record)
except Exception as e: # pragma: no cover - defensive
logger.warning(f"Agent on_step observer failed: {e}")
async def _call_tool(
self,
tool_call: ToolCall,
planner_context: dict[str, Any],
variables: Mapping[str, Any],
) -> tuple[ToolCall, dict, Any, float]:
"""Resolve arguments and execute a single tool call via the kernel.
Returns ``(tool_call, args, result, duration_seconds)``. Execution
errors are captured and returned as the result so the model can
self-correct. The duration is measured here rather than around the
caller's ``asyncio.gather`` — that would time the slowest call in the
batch and attribute it to all of them.
"""
args = self._resolve_args(tool_call, planner_context, variables)
start = time.perf_counter()
if not _is_tool_allowed(tool_call.name, self.allowed_tools):
# The tool was never described to the model; tell it so it can
# pick one it is actually allowed to call.
logger.warning(f"Tool {tool_call.name} is not allowed for this agent")
return (
tool_call,
args,
f"Error: tool {tool_call.name} is not available.",
time.perf_counter() - start,
)
logger.info(f"Calling tool {tool_call.name} with {args}")
try:
result = await self.kernel.call_tool(
tool_uri=tool_call.name, arguments=args
)
except Exception as e:
logger.error(f"Tool {tool_call.name} failed: {e}")
result = f"Error: {e}"
return tool_call, args, result, time.perf_counter() - start
if __name__ == "__main__": # pragma: no cover - manual demo
from kavalai.llm_clients.openai_client import OpenAIClient
from kavalai.functionkernel import FunctionKernel, pythontool
import datetime
@pythontool
def get_time():
return datetime.datetime.now(datetime.timezone.utc).isoformat()
kernel = FunctionKernel()
kernel.register_python_tool("get_time", get_time)
llm_client = OpenAIClient("gpt-5.6-luna")
agent = Agent(llm_client=llm_client, kernel=kernel, debug=True)
result = asyncio.run(agent.prompt("Greet the user based on current time!"))
print(result)