@dataclass
class Engine:
"""Generates implementations via an agent loop and commits them to the store.
`completion` is a LiteLLM-compatible callable; tests inject a fake. A single shared
engine backs all `@jiti` functions so the cycle guard and store stay consistent across
a cascade.
"""
store: JitiStore
model: str = DEFAULT_MODEL
max_tokens: int = DEFAULT_MAX_TOKENS
max_turns: int = DEFAULT_MAX_TURNS
style: str = STYLE_GUIDE
test_guide: str = TEST_GUIDE
quality_threshold: int = DEFAULT_QUALITY_THRESHOLD
max_refactor: int = DEFAULT_MAX_REFACTOR
test_paths: tuple[str, ...] | None = None
"""Where to find gates: None scans the tree, a tuple narrows it, () disables discovery."""
execution_timeout: float = DEFAULT_EXECUTION_TIMEOUT
"""Max IDLE seconds (no LLM call in flight) for one in-process execution — a test,
gate, or tool experiment. Cascaded generation resets the clock on every model call,
so deep cascades run as long as they need; only a hung candidate trips this."""
frozen: bool = False
"""Refuse to generate. A cache miss raises `FrozenError` instead of calling the LLM —
use in deployments where only committed code should run. The default engine also honors
`JITI_FROZEN=1` (env wins if both are off, an explicit `frozen=True` always freezes)."""
provider_retries: int = DEFAULT_NUM_RETRIES
"""How many times litellm retries a transient provider error (429, 5xx, network blip)
on each turn before giving up. Set to 0 to disable; the default lets a single hiccup
pass without bubbling out of the user's `@jiti` call."""
completion: Completion | None = None
_in_progress: set[str] = field(default_factory=set)
_discovered: bool = field(default=False)
_llm: LiteLLMClient = field(init=False)
def __post_init__(self) -> None:
if self.completion is None:
self._llm = LiteLLMClient(num_retries=self.provider_retries)
else:
self._llm = LiteLLMClient(self.completion, num_retries=self.provider_retries)
def discover(self) -> None:
"""Import the project's test modules once so their gates register before generation."""
if self._discovered:
return
self._discovered = True
import_test_modules(self.test_paths)
def _require_unfrozen(self, key: str, action: Action) -> None:
if not self.frozen:
return
verb = "generated" if action is Action.GENERATE else "regenerated (spec changed)"
raise FrozenError(
f"{key} needs to be {verb}, but the engine is frozen. Generate it in "
"development (unset JITI_FROZEN) and commit `.jiti/`, or run `jiti merge` "
"before the freeze."
)
def implement(
self,
declaration: Declaration,
args: tuple[object, ...],
kwargs: dict[str, object],
*,
target: RoutingTarget | None = None,
) -> Any:
resolution = self.store.resolve(declaration)
if resolution.action is Action.CONFLICT:
raise ConflictError(
f"{declaration.key}: the implementation was hand-edited and the declaration "
"has since changed. Reconcile them before running."
)
if resolution.action in (Action.GENERATE, Action.REGENERATE):
self._require_unfrozen(declaration.key, resolution.action)
self._generate(declaration, args, kwargs, target)
return self.store.load(declaration)
def _generate(
self,
declaration: Declaration,
args: tuple[object, ...],
kwargs: dict[str, object],
target: RoutingTarget | None,
) -> None:
key = declaration.key
if key in self._in_progress:
raise GenerationCycleError(
f"generation cycle: {key} is needed to generate itself (mutually recursive stubs)."
)
self._in_progress.add(key)
depth = len(self._in_progress)
log_start(key, depth)
started = perf_counter()
recorder = Recorder()
try:
gates = self._prepare_gates(declaration)
context = CallContext(
declaration,
args,
kwargs,
import_path=_import_path(declaration),
gates=gates,
timeout=self.execution_timeout,
recorder=recorder,
target=target,
)
run = self._run_agent(
context,
self._system_blocks(),
_task_prompt(declaration),
IMPL_TOOLS,
threshold=self.quality_threshold,
max_refactor=self.max_refactor,
)
if context.passing is None:
raise GenerationError(
f"{key}: the agent finished without an implementation that passes validation."
)
impl, tests = context.passing
self.store.write(declaration, impl, _committed_tests(declaration, tests))
log_done(key, depth, perf_counter() - started, run.cost, run.llm_seconds, run.llm_calls)
record_generation(run.cost)
finally:
recorder.write(transcript_path(self.store.root, declaration.module, declaration.name))
self._in_progress.discard(key)
def generate_test(self, test: Declaration, target: Declaration) -> None:
"""Generate a jiti-test body from `target`'s interface (TDD), validated by ruff + ty."""
if (section := self.store.read_test_section(test)) and section.spec_hash == test.spec_hash:
return
self._require_unfrozen(test.key, Action.GENERATE)
key = test.key
if key in self._in_progress:
raise GenerationCycleError(f"generation cycle: {key} is needed to generate itself.")
self._in_progress.add(key)
depth = len(self._in_progress)
log_start(key, depth)
started = perf_counter()
recorder = Recorder()
try:
context = CallContext(
test,
(),
{},
import_path=_import_path(test),
timeout=self.execution_timeout,
recorder=recorder,
)
task = _test_task_prompt(test, target)
# threshold=0 → no refactor pass for tests (red→green→refactor is for the impl).
run = self._run_agent(
context, self._test_system_blocks(), task, TEST_TOOLS, threshold=0, max_refactor=0
)
if context.passing is None:
raise GenerationError(f"{key}: the agent finished without a passing test.")
body, _ = context.passing
self.store.write_test(test, body)
log_done(key, depth, perf_counter() - started, run.cost, run.llm_seconds, run.llm_calls)
record_generation(run.cost)
finally:
recorder.write(transcript_path(self.store.root, test.module, test.name))
self._in_progress.discard(key)
def run_test(self, test: Declaration, target: Declaration) -> types.FunctionType:
"""Ensure a jiti-test is generated, then return its committed body to run."""
self.generate_test(test, target)
return self.store.load_test(test)
def _prepare_gates(self, declaration: Declaration) -> tuple[Gate, ...]:
# Human gates run as-is; jiti-test gates are generated (test-mode) then loaded so both
# run against the candidate via the same rebinding path.
runnable: list[Gate] = []
for gate in declaration.gates:
if gate.kind == "human" or gate.test is None:
runnable.append(gate)
continue
test_decl = introspect(gate.test)
loaded = self.run_test(test_decl, declaration)
runnable.append(Gate(gate.name, "jiti", gate.spec, test=loaded, target=gate.target))
return tuple(runnable)
def _run_agent(
self,
context: CallContext,
system: list[dict[str, Any]],
task: str,
tools: list[dict[str, Any]],
*,
threshold: int,
max_refactor: int,
) -> AgentRun:
depth = len(self._in_progress)
messages: list[dict[str, Any]] = [{"role": "user", "content": task}]
total_cost = 0.0
llm_seconds = 0.0
refactors = 0
turn = 0
for turn in range(1, self.max_turns + 1):
started = perf_counter()
response = self._llm.complete(
model=self.model,
max_tokens=self.max_tokens,
system=system,
tools=tools,
messages=messages,
)
usage = getattr(response, "usage", None)
key = context.declaration.key
elapsed = perf_counter() - started
llm_seconds += elapsed
log_llm_call(key, turn, depth, elapsed, usage, self.model)
spent = response.cost if response.cost is not None else cost(self.model, usage) or 0.0
total_cost += spent
if context.recorder is not None:
context.recorder.turn(turn, elapsed, usage, spent, response.content)
messages.append({"role": "assistant", "content": response.content})
tool_uses = [block for block in response.content if _is_tool_use(block)]
if not tool_uses:
return AgentRun(total_cost, llm_seconds, turn)
results = [_tool_result(context, block) for block in tool_uses]
if context.passing is not None:
if context.quality >= threshold or refactors >= max_refactor:
# green and polished enough — skip the model's wrap-up turn
return AgentRun(total_cost, llm_seconds, turn)
refactors += 1
results.append(_refactor_nudge(context.quality, threshold))
messages.append({"role": "user", "content": results})
return AgentRun(total_cost, llm_seconds, turn)
def _system_blocks(self) -> list[dict[str, Any]]:
# Style and test guidance are separate cached blocks so jiti's mechanical rules stay
# cacheable independent of whichever guides are in effect.
blocks = [_cached(SYSTEM_PROMPT)]
if self.style.strip():
house_style = f"Follow this house style in the code you write:\n\n{self.style}"
blocks.append(_cached(house_style))
if self.test_guide.strip():
guidance = f"Follow this guidance when writing tests:\n\n{self.test_guide}"
blocks.append(_cached(guidance))
return blocks
def _test_system_blocks(self) -> list[dict[str, Any]]:
blocks = [_cached(TEST_MODE_PROMPT)]
if self.test_guide.strip():
blocks.append(_cached(f"Follow this guidance when writing tests:\n\n{self.test_guide}"))
return blocks