JAZ

A Python framework for building and optimizing LLM-agents through intelligent, iterative execution.

Overview

JAZ provides a unified framework for creating agents that combine deterministic Python code with LLM reasoning. Agents execute in a REPL loop where the LLM generates code, observes results, and iterates until the task is complete.

Prerequisites

• Python ≥ 3.12 — 3.14+ is recommended for the full feature set. The t-string prompt syntax (invoke(task=t"Analyze the {data}")) and the interactive jaz console need PEP 750 t-strings, hence Python ≥ 3.14. Everything else works on 3.12/3.13; the equivalent invoke(task="Analyze the `data`", data=data) works on every supported version.

Installation

pip install jaz-lang

The distribution is named jaz-lang; the import name is jaz:

import jaz

Optional extras:

pip install "jaz-lang[tracing]"   # OpenTelemetry tracing
pip install "jaz-lang[copilot]"   # GitHub Copilot backend — no API key, billed to AI credits

To pin a specific release, or to install one that has not reached PyPI, use the distribution repo and a release tag:

pip install "jaz-lang @ git+https://github.com/jaz-lang/jaz.git@v0.2.0a4"

API Keys

JAZ reads provider API keys from the environment, under each provider's standard variable name. Set them via any of:

• Shell profile — add export OPENAI_API_KEY=sk-... to ~/.zshrc or ~/.bashrc (simplest, available globally)
• .env file — create a .env at the project root with OPENAI_API_KEY=sk-... per line, then load it with:
  • direnv — add dotenv to an .envrc file to auto-load on cd (brew install direnv)

No API key at hand? The copilot backend (the copilot extra) bills a Copilot subscription's GitHub AI credits instead. Authentication checks an explicit github_token first, then COPILOT_GITHUB_TOKEN / GH_TOKEN / GITHUB_TOKEN in the environment (in that order), then your signed-in gh/copilot CLI. An exported token therefore takes precedence over your CLI login:

import jaz
from jaz.llm import CopilotLLM

jaz.configure(llm=CopilotLLM(model="gpt-5-mini"))   # ids are Copilot's own, unprefixed

For CLI authentication, run gh auth login, or start copilot and enter /login.

Like every backend other than LiteLLM it sits outside the stable API surface — see Custom & Local LLM Backends.

Quick Start

from jaz import invoke
from jaz.hooks import ReturnType

# Simple task
result = invoke(
    ReturnType(int),
    task="Calculate the factorial of 10",
)

# With custom inputs
result = invoke(
    ReturnType(list),
    task="Double each number in the list",
    numbers=[1, 2, 3, 4, 5]
)

Later snippets assume the Quick Start imports (from jaz import invoke, from jaz.hooks import ReturnType, import jaz) and show only what each adds.

Every input is a keyword argument — the prompt included. invoke()'s positional arguments are hooks (ReturnType(int), a budget, …) and an optional ConfigOverride, so invoke("Calculate the factorial of 10") is an error, not a shorthand.

The jaz console

Installing puts a jaz command on PATH (it needs Python ≥ 3.14 to run): a real Python REPL — variables persist, _ holds the last result — with conversational sugar layered on top. Run jaz, then enter these lines at its prompt:

> what is 2 + 2?
n: int <- how many days in a leap year?
?n
% switch to gpt-5-mini
import math

> asks the agent; n: int <- captures a typed result; ?n shows what the agent sees for n; % asks the jaz helper to change settings, with confirmation. Plain Python still runs. Enter the sigil lines without trailing # annotations: they become prompt text, or can leave a ? expression waiting for more input.

A session with no model configured opens a setup wizard first: pick a provider and paste its API key, or pick github copilot — no provider key, just the copilot extra and GitHub authentication. The choice is validated with one live request before anything is kept. Backend and model defaults are saved to ~/.jaz/settings.json; provider API keys entered in the wizard go in ~/.jaz/credentials.json.

Flags configure one session without touching that default — jaz --backend copilot --model gpt-5-mini, jaz --max-cost 0.50, jaz --log run.log — and jaz -c "…" runs a single prompt and exits. In-session, setup_llm() re-runs the wizard and switch_model("openai/gpt-5") changes model (keeping the backend) after validating it.

Key Features

Recursive Agent Composition

Agents can invoke nested agents with separate budgets and automatic cost tracking:

from jaz.hooks import RecursionLimit, ReturnType

with RecursionLimit(max_depth=3):  # optional cap; recursion is unbounded by default
    result = invoke(
        ReturnType(list),
        task="""
        For each item, use jaz.invoke to process it.
        Combine and return the results.
        """,
        items=data,
    )

Custom Tools

A tool is just a function you pass in. Give it as a keyword argument and it binds to that name in the agent's REPL; its signature and docstring render in the prompt, so the docstring is the description the agent reads:

def my_tool(x: int) -> int:
    """Double a number."""
    return x * 2

result = invoke(ReturnType(int), task="Use my_tool on 5", my_tool=my_tool)

The agent sees `my_tool(x: int) -> int`: Double a number. — type hints and a one-line docstring are what make a tool usable. Nothing is registered up front, so tools sit alongside data in the same call:

def web_search(query: str) -> list[str]:
    """Search the web and return result snippets."""
    ...

result = invoke(
    ReturnType(dict),
    task="Look up each name and summarize what you find",
    names=["ada", "grace"],
    web_search=web_search,
)

To group related tools under one name, pass an object instead: its public methods render as a catalog under that name, and the agent calls them as utils.my_tool(...).

class Utils:
    """Utility helpers."""

    def my_tool(self, x: int) -> int:
        """Double a number."""
        return x * 2

result = invoke(ReturnType(int), task="Use utils.my_tool on 5", utils=Utils())

To make a tool propagate automatically to nested invoke calls, bind it with jaz.scope instead of passing it per call:

from jaz import scope

with scope(my_tool=my_tool):
    invoke(ReturnType(int), task="Use my_tool on 5")  # nested invokes inherit `my_tool`

Hooks for Extensibility

Use hooks to add logging, tracing, workflow strategies, and more:

from jaz.hooks import BudgetPool, FileLogger

with FileLogger("agent.log"), BudgetPool(cost_budget=1.0):
    result = invoke(ReturnType(str), task="Build a web scraper")

Budget Control

Pool budgets (LLM cost / call count, shared across the whole invoke/recursion tree) are enforced by the opt-in BudgetPool; per-level execution limits are configuration:

from jaz.hooks import BudgetPool, IterationLimit, RecursionLimit

# Pool budgets: tracked + enforced only while the hook is active.
with BudgetPool(cost_budget=0.50, calls_budget=100):  # USD / call-count
    result = invoke(ReturnType(str), task="Complex task")

# Per-level execution limits are hooks too:
with IterationLimit(max_iterations=20), RecursionLimit(max_depth=2):
    result = invoke(ReturnType(str), task="Complex task")

Configuration

Override defaults globally or per-invoke:

Each group takes the configured component, whose constructor is that group's settings. Setting a group replaces it — a component states itself completely, so there is no partial update of one setting.

from jaz.llm import LiteLLM
from jaz.repl.python_repl import PythonREPL

# Per-invoke override
with jaz.ConfigOverride(llm=LiteLLM(model="openai/gpt-5-mini", temperature=0.7, max_tokens=2000)):
    result = invoke(ReturnType(str), task="Creative task")

# Global configuration
jaz.configure(llm=LiteLLM(model="openai/gpt-5-mini"), repl=PythonREPL(exec_timeout=60))

Custom & Local LLM Backends

JAZ talks to LLMs through a small BaseLLM layer. Currently, LiteLLM is the only backend covered by the stable API, and the default — one backend that routes to every provider LiteLLM supports (OpenAI, Anthropic, Gemini, Bedrock, Vertex, …). Other backends ship in the package but can change or be removed. You can add your own backend without forking.

A config selects one backend: the BaseLLM you pass to llm=, defaulting to LiteLLM. Its constructor takes everything it needs — the model id (a LiteLLM route like openai/gpt-5-mini), the backend's own settings, and any per-request defaults. A config file can name a backend instead of constructing one, but the name lives in whatever reads the file: JAZ's own loaders (the eval harness, the console) each keep their own tag → class mapping, and core carries neither the mapping nor a name on the class. Because litellm is the default, a config need only name the model.

Local OpenAI-compatible server (Ollama, vLLM, LM Studio, llama.cpp, …) — no code, just config:

# route the litellm backend to a local OpenAI-compatible server
# (api_key is required by the openai/ route even for a local server — a dummy value is fine;
#  it rides request_defaults into litellm.completion)
jaz.configure(
    llm=LiteLLM(model="openai/llama3", api_base="http://localhost:11434/v1", api_key="dummy"),
)

A custom backend — subclass BaseLLM and pass an instance (a name is only needed if a config file must select it, and it is that file's reader that holds the name):

import os
from jaz.llm import BaseLLM, CompletionResponse, Usage, Choice, Message

class MyLLM(BaseLLM):
    def __init__(self, base_url: str | None = None, **retry):
        super().__init__(**retry)             # forwards the retry_* settings
        self.base_url = base_url or os.environ["MYBACKEND_API_BASE"]

    def complete(self, model, messages, **kwargs):
        # ... call self.base_url; raise jaz.llm.RateLimitError / AuthenticationError / ...
        # on API errors so retry classification works.
        response = CompletionResponse(
            choices=[Choice(message=Message("assistant", "..."))],
            usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2),
            model=model,
        )
        return self.finalize(response, model)  # token accounting + cost, for free

jaz.configure(llm=MyLLM(model="my-model", base_url="http://localhost:8000"))

One class owns the whole job: the API call, retry, cost accounting and model metadata. finalize() turns your wire response into the normalized jaz.llm.LLMResponse and prices it from the bundled table (models absent from it simply report cost=None); the retry wrappers and the non-retryable-error classification come from the base.

When a config file names your backend, the split between "settings for the backend object" and "params for the request" is taken from your __init__ signature, so base_url reaches MyLLM(...) while model rides each call — a backend declares its construction keys just by declaring __init__. Whoever loads that file supplies the name → class mapping (jaz.instantiate.build_config(..., resolvers={"llm": my_lookup})); JAZ core holds no registry, and a backend class carries no name of its own. Constructing it yourself needs no name at all: ConfigOverride(llm=my_llm).

OpenAILLM and AnthropicLLM ship in the package but are in no loader's tag map, so backend: openai names nothing — reach those providers through litellm's openai/… / anthropic/… routes. Being nameable and being stable are separate axes: a name can resolve without being covered by the stable API. To add or replace a backend, subclass BaseLLM; add it to your loader's map only if a config file needs to name it.

The configured model id is sent verbatim: JAZ neither strips nor rewrites a provider/ prefix, so write the id the backend's own documentation uses. A backend is chosen by configuring it, not by prefixing the model — OpenAILLM(model="litellm/gpt-4o") sends that string to OpenAI, which rejects it.

Architecture

jaz/
├── invoke.py       # Public invoke() API
├── config.py       # Configuration system
├── budget.py       # Cost tracking
├── repl/           # Python REPL implementations
├── hooks/          # Hook system, event orchestration, and built-in hooks
├── llm/            # LLM backends (LiteLLM by default)
└── protocol/       # Wire-format codec between the LLM and the REPL

Core Concepts

• Agent loop: Orchestrates LLM queries and REPL execution
• REPL: Executes agent-generated code with safety controls
• Hooks: Event-based system for logging, budget control, workflow capture, and extensibility
• Tools: Plain functions passed as inputs; their signature and docstring become the description the agent reads

Built-in Hooks

Hook	Purpose
ReturnType	Declare + enforce the invoke's return type
ValidateReturn / ValidateREPLCode	Validate the return value / veto REPL code before it runs
BudgetPool	Shared LLM cost / call-count budget with hard-stop enforcement
IterationLimit / RecursionLimit	Per-level turn cap / invoke-nesting cap
BudgetForcing	Refuse early finishes so the agent keeps working
ContextWindowWarning	Warn the agent as its prompt nears the model's window
PrintLogger / FileLogger	Log events to console / file
TrajectoryRecorder	Write the run as an ATIF trajectory (replay/cost source)
TrajectoryDirectoryRecorder	Write the run as a browsable Markdown directory, live (tail -f a running agent)
TrajectoryReplay	Resume a run from an ATIF trace: saved responses replay, then live calls take over
RolloutRecorder	Record token-native rollouts for training
JaegerTracing / LangfuseTracing	OpenTelemetry tracing presets

Writing Custom Hooks

Hooks observe agent execution via events (from jaz.hooks.events) and influence it by returning effects (from jaz.hooks.effects):

from jaz import invoke
from jaz.hooks import Hook
from jaz.hooks.effects import AddMessages, Effect
from jaz.hooks.events import LLMQueryEnter

class ConciseHook(Hook):
    """Instruct the agent to be concise after iteration 3."""

    # Override the typed per-event handler for the event you care about — no
    # isinstance/dispatch boilerplate. (A cross-cutting observer that wants *every*
    # event overrides `on_any` instead.)
    def on_llm_query_enter(self, event: LLMQueryEnter) -> list[Effect]:
        if event.iteration >= 3:  # iterations are 0-based
            return [AddMessages([{"role": "user", "content": "Be concise."}])]
        return []

# Hooks activate via context managers and propagate to nested invoke() calls
with ConciseHook():
    result = invoke(task="Solve this step by step")

Hooks compose naturally as context managers:

# Stack multiple hooks
with ConciseHook(), FileLogger("agent.log"):
    result = invoke(...)

Event Types

Events fire around three spans — the whole Invoke, each turn's LLMQuery, and each turn's REPLExec — and every span walks the same four stages. All are importable from jaz.hooks.events.

Stage	Observes	Accepts	Fires
*Enter	the proposal	edit effects + Abort	always
*Send	the committed input	supply effects + Abort	iff the input committed
*Complete	the raw result	modify effects + Abort	iff a raw result exists
*Exit	the outcome union (Completed | Aborted | Failed)	nothing (observation-only)	whenever the span opened

Plus LLMQueryRetry, fired per retry attempt of an LLM call. Every event carries contextual data (e.g. LLMQueryEnter.model, LLMQueryEnter.iteration) and a timestamp (its emission time — durations are arithmetic over two events' timestamps).

Effect Types

Effects are returned by hooks to influence execution. All are importable from jaz.hooks.effects; each composes at a specific stage (an effect returned anywhere else raises InvalidEffectError).

Effect	Stage	Influence
AddInputs / DropInputs	InvokeEnter	Add/un-pass invoke inputs (prompt + REPL)
DisableRecursion	InvokeEnter	Withhold the recursive jaz.invoke tool
AddMessages / DropMessages	LLMQueryEnter	Edit the messages sent to the model (transient or persistent)
AddVariables / DropVariables	REPLExecEnter	Bind/unbind REPL namespace names for the turn
SupplyLLMResponse	LLMQuerySend	Supply the response, skipping the API call
SupplyExecResult	REPLExecSend	Supply a result, skipping execution
ModifyExecResult	REPLExecComplete / InvokeComplete	Replace the raw result
Abort	any Enter/Send/Complete	Abort the invoke; its error propagates out of invoke()
BlackboardWrite	any event	Write to the per-invoke cross-hook blackboard

Langfuse Tracing Quick Start

from jaz import invoke
from jaz.hooks import LangfuseTracing, ReturnType

with LangfuseTracing():
    result = invoke(ReturnType(int), task="Calculate 2+2")

Environment variables:

• LANGFUSE_PUBLIC_KEY
• LANGFUSE_SECRET_KEY
• LANGFUSE_HOST (optional, default: https://cloud.langfuse.com)

Contributing

This repository is generated: each release publishes a cleaned copy of the package, so changes are not made here. Contributor setup and code conventions live with the source.

License

See LICENSE file.