@empyria/restate 0.4.11 → 0.4.12

package/README.md

@@ -133,12 +133,13 @@ address, so the deployment must survive the process stopping and come back at th
 
 `defineAgent` options beyond the basics:
 
-| Option                   | Does                                                                                                                                                                                                                                       |
-| ------------------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
-| `tools` as a function    | `tools: (input) => ({...})` builds the tools from the validated request, for tools that depend on it. It runs outside any `ctx.run` on every execution and replay, so it must be pure. An invalid map fails that request terminally (500). |
-| `respond(result, input)` | Shapes what `ask` returns from the loop result `{text, output?, steps, totalTokens, trace}`. Pair it with `responseSchema`, which validates the shaped value before Restate records it.                                                    |
-| `label(input)`           | A terminal failure then reads `[label] message`, so a failed run says what it was working on. Code and metadata are kept; cancellations, timeouts and retryable errors are untouched.                                                      |
-| `limits.maxInputBytes`   | Largest request (JSON bytes, default 64 KB); a larger one is refused terminally (413).                                                                                                                                                     |
+| Option                   | Does                                                                                                                                                                                                                                                                                                                                                                  |
+| ------------------------ | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `tools` as a function    | `tools: (input) => ({...})` builds the tools from the validated request, for tools that depend on it. It runs outside any `ctx.run` on every execution and replay, so it must be pure. An invalid map fails that request terminally (500).                                                                                                                            |
+| `output` as a function   | `output: (input) => schema \| undefined` chooses the answer's JSON Schema per request (same rules as `tools`: pure, no I/O); `undefined` means a free-text answer for that request. With any `output`, the model's answer is checked against the schema: a wrong one is handed back to the model to correct, and a model that never gets it right ends at `maxSteps`. |
+| `respond(result, input)` | Shapes what `ask` returns from the loop result `{text, output?, steps, totalTokens, trace}`. Pair it with `responseSchema`, which validates the shaped value before Restate records it.                                                                                                                                                                               |
+| `label(input)`           | A terminal failure then reads `[label] message`, so a failed run says what it was working on. Code and metadata are kept; cancellations, timeouts and retryable errors are untouched.                                                                                                                                                                                 |
+| `limits.maxInputBytes`   | Largest request (JSON bytes, default 64 KB); a larger one is refused terminally (413).                                                                                                                                                                                                                                                                                |
 
 **Keying block tools.** A block tool calls another Restate building block. Its `key(args, meta)` and `input(args, meta)`
 hooks decide the target key and the parameter from the model's arguments and a per-call idempotency key

package/lib/agent/Agent.js

@@ -68,9 +68,11 @@ export const AskOutputSchema = {
  * @property {object} [input] Input JSON Schema. Defaults to {@link AskInputSchema}.
  * @property {(input: any) => string} [prompt] Builds the user prompt from the validated
  *   input. Defaults to `input.prompt`.
- * @property {object} [output] Output JSON Schema. When given, the model must answer through
- *   the `final_answer` tool and `ask` returns its arguments. Defaults to
- *   {@link AskOutputSchema} (`{text, steps}`).
+ * @property {object | ((input: any) => object | undefined)} [output] Output JSON Schema. When given, the model
+ *   must answer through the `final_answer` tool; its arguments are checked against the schema (a wrong answer is
+ *   handed back to the model to correct) and `ask` returns them. A function chooses the schema per request, like
+ *   `tools`: it runs on every execution and replay, so it must be pure and do no I/O; `undefined` means a free-text
+ *   answer for that request. Defaults to {@link AskOutputSchema} (`{text, steps}`).
  * @property {(result: {text: string, output?: unknown, steps: number, totalTokens: number, trace: Array<object>}, input: any) => any} [respond]
  *   Shapes what `ask` returns from the loop's result (`output` is set when an `output` schema is
  *   given; `trace` is the per-step summary). Without it, `ask` returns `output`, or `{text, steps}`.
@@ -162,7 +164,12 @@ export function defineAgent(spec) {
 		maxInputBytes = DEFAULT_MAX_INPUT_BYTES,
 	} = limits
 
-	const outputSchema = respond ? (responseSchema ?? {}) : (output ?? AskOutputSchema)
+	// A per-request `output` cannot be known here; the loop checks the answer against the schema it was given.
+	const outputSchema = respond
+		? (responseSchema ?? {})
+		: typeof output === 'function'
+			? {}
+			: (output ?? AskOutputSchema)
 
 	const ask = withValidation(input, outputSchema, async (ctx, validatedInput) => {
 		const inputBytes = jsonBytes(validatedInput)
@@ -193,6 +200,8 @@ export function defineAgent(spec) {
 		const requestTools =
 			typeof tools === 'function' ? await resolveTools(name, tools, validatedInput) : tools
 
+		const requestOutput = typeof output === 'function' ? output(validatedInput) : output
+
 		const context = loadContext
 			? await ctx.run('load-context', async () => {
 					const { systemPrompt = '', skills = [] } =
@@ -217,7 +226,7 @@ export function defineAgent(spec) {
 			skills: context.skills,
 			loadSkill: skillLoader({ loadSkill, loadContext, input: validatedInput }),
 			loadHistory: session?.load,
-			outputSchema: output,
+			outputSchema: requestOutput,
 			maxSteps,
 			maxTokens,
 			retry: llmRetry,
@@ -226,7 +235,7 @@ export function defineAgent(spec) {
 		if (session) await session.save(result.messages, maxHistoryTurns)
 
 		if (respond) return respond(result, validatedInput)
-		return output ? result.output : { text: result.text, steps: result.steps }
+		return requestOutput ? result.output : { text: result.text, steps: result.steps }
 	}
 
 	const options = { retryPolicy }

package/lib/agent/Loop.js

@@ -1,11 +1,16 @@
 import { generateText } from 'ai'
 import { TerminalError } from '@restatedev/restate-sdk'
+import { createValidator } from '@empyria/common'
 import { DEFAULT_LLM_RETRY, toRestateLLMError } from './Errors.js'
 import { cachedSkillFetcher, expandSkillRefs } from './Skills.js'
 import { FINAL_ANSWER_TOOL, buildToolDefs, executeTool } from './Tools.js'
 
 const FINAL_ANSWER_NUDGE = `Respond by calling the ${FINAL_ANSWER_TOOL} tool with your answer.`
 
+/** The text handed back to the model when its final answer does not fit the schema (kept short). */
+const rejection = (message) =>
+	`Your answer was not accepted because it does not match the required format: ${String(message).slice(0, 1500)}\nCall ${FINAL_ANSWER_TOOL} again with a corrected answer.`
+
 /**
  * Drives an AI SDK Core tool-calling loop ONE round at a time, from inside a Restate
  * handler — deliberately NOT `generateText`'s own multi-step loop wrapped in a single
@@ -36,7 +41,8 @@ const FINAL_ANSWER_NUDGE = `Respond by calling the ${FINAL_ANSWER_TOOL} tool wit
  *   retry?: import('@restatedev/restate-sdk').RunOptions<any>,
  * }} params `skills` is the index only (name + description); `loadSkill` fetches a body
  *   when the model asks for it. With `outputSchema`, the model must answer through the
- *   `final_answer` tool and the loop returns that tool's input as `output`.
+ *   `final_answer` tool and the loop returns that tool's input as `output`, after checking it against the
+ *   schema: an answer that does not fit is handed back to the model to correct, not accepted.
  * @returns {Promise<{text: string, output?: unknown, steps: number, totalTokens: number, trace: Array<StepTrace>, messages: Array<object>}>}
  *   `messages` are the new messages of this turn only (prompt, responses, tool results).
  *   `trace` has one small entry per LLM step. The same facts are recorded in the journal with each
@@ -72,6 +78,7 @@ export async function runAgentLoop(
 	const messages = [{ role: 'user', content: prompt }]
 	let totalTokens = 0
 	const fetchSkill = cachedSkillFetcher(loadSkill)
+	const validateOutput = outputSchema ? createValidator(outputSchema) : undefined
 	const trace = []
 
 	for (let step = 0; step < maxSteps; step++) {
@@ -149,14 +156,26 @@ export async function runAgentLoop(
 
 		for (const [index, call] of result.toolCalls.entries()) {
 			if (outputSchema && call.toolName === FINAL_ANSWER_TOOL) {
-				return {
-					text: result.text,
-					output: call.input,
-					steps: step + 1,
-					totalTokens,
-					trace,
-					messages,
+				let output
+				try {
+					output = validateOutput(call.input)
+				} catch (error) {
+					// The model's arguments are not checked by the provider. Tell it what is wrong and let it
+					// answer again; this uses a step, so a model that cannot get it right ends at maxSteps.
+					messages.push({
+						role: 'tool',
+						content: [
+							{
+								type: 'tool-result',
+								toolCallId: call.toolCallId,
+								toolName: call.toolName,
+								output: { type: 'error-text', value: rejection(error.message) },
+							},
+						],
+					})
+					continue
 				}
+				return { text: result.text, output, steps: step + 1, totalTokens, trace, messages }
 			}
 			const output = await executeTool(ctx, call, { tools, loadSkill, step, index })
 			messages.push({

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@empyria/restate",
-	"version": "0.4.11",
+	"version": "0.4.12",
 	"description": "Restate.dev helpers for the Empyria nanoservice framework",
 	"license": "MIT",
 	"author": "Imre Fazekas <imre.fazekas@icloud.com>",

package/test/agent/Agent.test.js

@@ -111,6 +111,50 @@ describe('defineAgent: ask', () => {
 		await expect(answer({ score: 'high' })).rejects.toThrow(TerminalError)
 	})
 
+	test('output can be chosen per request; undefined means a free-text answer', async () => {
+		const schemas = {
+			score: {
+				type: 'object',
+				properties: { score: { type: 'number' } },
+				required: ['score'],
+			},
+			name: { type: 'object', properties: { name: { type: 'string' } }, required: ['name'] },
+		}
+		const agent = (model) =>
+			defineAgent({
+				name: 'A',
+				model,
+				input: { type: 'object', properties: { want: { type: 'string' } } },
+				prompt: () => 'go',
+				output: ({ want }) => schemas[want],
+			})
+		const answering = (input) =>
+			new MockLanguageModelV4({
+				doGenerate: toolCallResult([{ toolName: 'final_answer', input }]),
+			})
+
+		expect(
+			await agent(answering({ score: 3 })).service.ask(fakeCtx(), { want: 'score' }),
+		).toEqual({ score: 3 })
+		expect(
+			await agent(answering({ name: 'x' })).service.ask(fakeCtx(), { want: 'name' }),
+		).toEqual({
+			name: 'x',
+		})
+		// A wrong answer for the request's schema is corrected, not accepted.
+		const model = new MockLanguageModelV4({
+			doGenerate: sequence(
+				toolCallResult([{ toolName: 'final_answer', input: { name: 5 } }]),
+				toolCallResult([{ toolName: 'final_answer', input: { name: 'ok' } }]),
+			),
+		})
+		expect(await agent(model).service.ask(fakeCtx(), { want: 'name' })).toEqual({ name: 'ok' })
+		// No schema for this request: the plain answer.
+		expect(await agent(okModel()).service.ask(fakeCtx(), { want: 'none' })).toMatchObject({
+			steps: 1,
+		})
+	})
+
 	test('a custom input schema with a prompt builder', async () => {
 		const model = okModel()
 		const agent = defineAgent({

package/test/agent/Loop.test.js

@@ -154,6 +154,51 @@ describe('runAgentLoop', () => {
 		expect(JSON.stringify(model.doGenerateCalls[1].prompt)).toContain('final_answer')
 	})
 
+	test('an answer that does not fit the schema goes back to the model, which corrects it', async () => {
+		const model = new MockLanguageModelV4({
+			doGenerate: sequence(
+				toolCallResult([{ toolName: 'final_answer', input: { score: 'high' } }]),
+				toolCallResult([{ toolName: 'final_answer', input: { score: 7 } }]),
+			),
+		})
+
+		const result = await runAgentLoop(fakeCtx(), {
+			model,
+			system: '',
+			prompt: 'rate it',
+			outputSchema: {
+				type: 'object',
+				properties: { score: { type: 'number' } },
+				required: ['score'],
+			},
+		})
+
+		expect(result).toMatchObject({ output: { score: 7 }, steps: 2 })
+		const second = JSON.stringify(model.doGenerateCalls[1].prompt)
+		expect(second).toContain('does not match the required format')
+		expect(second).toContain('final_answer')
+	})
+
+	test('a model that never gets the answer right ends at maxSteps, terminally', async () => {
+		const model = new MockLanguageModelV4({
+			doGenerate: toolCallResult([{ toolName: 'final_answer', input: { score: 'high' } }]),
+		})
+		const caught = await runAgentLoop(fakeCtx(), {
+			model,
+			system: '',
+			prompt: 'rate it',
+			maxSteps: 3,
+			outputSchema: {
+				type: 'object',
+				properties: { score: { type: 'number' } },
+				required: ['score'],
+			},
+		}).catch((e) => e)
+		expect(caught).toBeInstanceOf(TerminalError)
+		expect(caught.message).toContain('maxSteps')
+		expect(model.doGenerateCalls).toHaveLength(3)
+	})
+
 	test('prepends loadHistory’s messages to every model call, without journaling them', async () => {
 		const ctx = fakeCtx()
 		const model = new MockLanguageModelV4({ doGenerate: textResult('ok') })