Skip to content

Commit 90564c0

Browse files
NitayRabiNitayRabi
authored
Merge pull request #4 from code-rabi/feature-generate-object
feat: generate object
2 parents f21f36f + 73aeb92 commit 90564c0

6 files changed

Lines changed: 564 additions & 2 deletions

File tree

‎README.md‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,53 @@ const result = await rlm.completion(
8282

8383
The LLM will know it can access `context.users`, `context.settings`, etc. with full type awareness.
8484

85+
### Structured Output with Zod (`generateObject`)
86+
87+
If you want schema-validated JSON output directly (without REPL/code execution), use `generateObject`.
88+
RLLM will retry when output is invalid JSON or fails Zod validation.
89+
90+
```typescript
91+
import { z } from 'zod';
92+
import { createRLLM } from 'rllm';
93+
94+
const rlm = createRLLM({ model: 'gpt-4o-mini' });
95+
96+
const OutputSchema = z.object({
97+
summary: z.string(),
98+
keyPoints: z.array(z.string()),
99+
confidence: z.number().min(0).max(1),
100+
});
101+
const InputSchema = z.object({
102+
reportText: z.string(),
103+
locale: z.string(),
104+
});
105+
106+
const result = await rlm.generateObject(
107+
"Summarize this report and provide key points with confidence",
108+
{
109+
input: {
110+
reportText: hugeDocument,
111+
locale: "en-US",
112+
},
113+
inputSchema: InputSchema,
114+
outputSchema: OutputSchema,
115+
},
116+
{
117+
maxRetries: 2, // total attempts = 3
118+
onRetry: (event) => {
119+
console.log(`Retry ${event.attempt}/${event.maxRetries + 1}: ${event.errorType}`);
120+
},
121+
}
122+
);
123+
124+
console.log(result.object.summary);
125+
console.log(result.attempts, result.usage.tokenUsage.totalTokens);
126+
```
127+
128+
`generateObject` differs from `completion()`:
129+
- `generateObject` asks for one JSON object and validates it against your schema.
130+
- `completion()` runs the full recursive REPL workflow where the model writes and executes JS code.
131+
85132
The LLM will write code like:
86133
```javascript
87134
// LLM-generated code runs in V8 isolate
@@ -153,6 +200,7 @@ Defaults:
153200
| Method | Description |
154201
|--------|-------------|
155202
| `rlm.completion(prompt, options)` | Full RLM completion with code execution |
203+
| `rlm.generateObject(prompt, { input?, inputSchema?, outputSchema }, options?)` | Structured output with Zod validation + retries |
156204
| `rlm.chat(messages)` | Direct LLM chat |
157205
| `rlm.getClient()` | Get underlying LLM client |
158206

@@ -164,6 +212,23 @@ Defaults:
164212
| `context` | `string \| T` | The context data available to LLM-generated code |
165213
| `contextSchema` | `ZodType<T>` | Optional Zod schema describing context structure |
166214

215+
### `GenerateObjectOptions`
216+
217+
| Option | Type | Description |
218+
|--------|------|-------------|
219+
| `maxRetries` | `number` | Retries after first attempt (default `2`) |
220+
| `temperature` | `number` | Optional generation temperature |
221+
| `maxTokens` | `number` | Optional max completion tokens |
222+
| `onRetry` | `(event) => void` | Called when parse/validation fails and a retry is scheduled |
223+
224+
### `GenerateObject` schema config
225+
226+
| Field | Type | Description |
227+
|-------|------|-------------|
228+
| `input` | `TInput` | Optional structured input value |
229+
| `inputSchema` | `ZodType<TInput>` | Optional input schema used for pre-validation + prompt typing |
230+
| `outputSchema` | `ZodType<TOutput>` | Required output schema used for retry validation |
231+
167232
### Sandbox Bindings
168233

169234
The V8 isolate provides these bindings to LLM-generated code:

‎src/index.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,4 +57,10 @@ export type {
5757
RLMEventType,
5858
RLMEvent,
5959
RLMEventCallback,
60+
GenerateObjectErrorType,
61+
GenerateObjectRetryEvent,
62+
GenerateObjectOptions,
63+
GenerateObjectSchemas,
64+
GenerateObjectUsage,
65+
GenerateObjectResult,
6066
} from "./types.js";

‎src/prompts.ts‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
*/
77

88
import type { ZodType } from "zod";
9+
import type { GenerateObjectRetryEvent } from "./types.js";
910

1011
/**
1112
* Main RLM system prompt - instructs the LLM on how to use the REPL environment
@@ -347,3 +348,65 @@ export function buildUserPrompt(
347348

348349
return { role: "user", content };
349350
}
351+
352+
/**
353+
* Build messages for schema-constrained JSON generation.
354+
*/
355+
export function buildGenerateObjectMessages(
356+
prompt: string,
357+
outputSchema: ZodType,
358+
input?: unknown,
359+
inputSchema?: ZodType
360+
): Array<{ role: "system" | "user"; content: string }> {
361+
const outputSchemaDescription = zodSchemaToTypeDescription(outputSchema);
362+
const inputSchemaDescription = inputSchema ? zodSchemaToTypeDescription(inputSchema) : null;
363+
const inputSection = input !== undefined
364+
? (
365+
"Input value (JSON):\n" +
366+
`\`\`\`json\n${JSON.stringify(input, null, 2)}\n\`\`\`\n\n`
367+
)
368+
: "";
369+
const inputTypeSection = inputSchemaDescription
370+
? (
371+
"Input TypeScript type:\n" +
372+
`\`\`\`typescript\ntype Input = ${inputSchemaDescription}\n\`\`\`\n\n`
373+
)
374+
: "";
375+
376+
return [
377+
{
378+
role: "system",
379+
content:
380+
"You generate structured data. Return exactly one valid JSON object that matches the provided schema. " +
381+
"Do not include markdown, code fences, comments, or any extra text before/after the JSON.",
382+
},
383+
{
384+
role: "user",
385+
content:
386+
`Task:\n${prompt}\n\n` +
387+
inputTypeSection +
388+
inputSection +
389+
"Target TypeScript type:\n" +
390+
`\`\`\`typescript\ntype Output = ${outputSchemaDescription}\n\`\`\`\n\n` +
391+
"Return only the JSON object.",
392+
},
393+
];
394+
}
395+
396+
/**
397+
* Build retry feedback after a failed parse/validation attempt.
398+
*/
399+
export function buildGenerateObjectRetryPrompt(
400+
event: GenerateObjectRetryEvent
401+
): { role: "user"; content: string } {
402+
const issues = event.validationIssues?.length
403+
? `\nValidation issues:\n- ${event.validationIssues.join("\n- ")}`
404+
: "";
405+
406+
return {
407+
role: "user",
408+
content:
409+
`Previous attempt ${event.attempt} failed (${event.errorType}): ${event.errorMessage}.${issues}\n\n` +
410+
"Please try again and return only one corrected JSON object with no surrounding text.",
411+
};
412+
}

‎src/rlm.test.ts‎

Lines changed: 190 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,190 @@
1+
import { describe, it, expect, vi } from "vitest";
2+
import { z } from "zod";
3+
import { RLLM } from "./rlm.js";
4+
import type { ChatMessage, TokenUsage } from "./types.js";
5+
6+
interface MockCompletionResponse {
7+
content: string;
8+
usage: TokenUsage;
9+
}
10+
11+
function createTestRLLMWithMock(responses: MockCompletionResponse[]): {
12+
rllm: RLLM;
13+
completeMock: ReturnType<typeof vi.fn>;
14+
} {
15+
const rllm = new RLLM({
16+
client: {
17+
provider: "openai",
18+
model: "gpt-4o-mini",
19+
apiKey: "test-key",
20+
},
21+
});
22+
23+
const completeMock = vi.fn().mockImplementation(async () => {
24+
const next = responses.shift();
25+
if (!next) {
26+
throw new Error("No mock response configured");
27+
}
28+
return {
29+
message: { role: "assistant", content: next.content },
30+
usage: next.usage,
31+
finishReason: "stop",
32+
};
33+
});
34+
35+
(
36+
rllm as unknown as {
37+
client: {
38+
complete: (options: { messages: ChatMessage[] }) => Promise<unknown>;
39+
};
40+
}
41+
).client = { complete: completeMock };
42+
43+
return { rllm, completeMock };
44+
}
45+
46+
describe("RLLM.generateObject", () => {
47+
it("returns typed object on first valid attempt", async () => {
48+
const outputSchema = z.object({
49+
name: z.string(),
50+
count: z.number(),
51+
});
52+
const inputSchema = z.object({
53+
report: z.string(),
54+
});
55+
56+
const { rllm } = createTestRLLMWithMock([
57+
{
58+
content: '{"name":"ok","count":3}',
59+
usage: { promptTokens: 11, completionTokens: 7, totalTokens: 18 },
60+
},
61+
]);
62+
63+
const result = await rllm.generateObject(
64+
"Generate object",
65+
{
66+
input: { report: "hello" },
67+
inputSchema,
68+
outputSchema,
69+
}
70+
);
71+
72+
expect(result.object).toEqual({ name: "ok", count: 3 });
73+
expect(result.attempts).toBe(1);
74+
expect(result.rawResponse).toBe('{"name":"ok","count":3}');
75+
expect(result.usage.totalCalls).toBe(1);
76+
expect(result.usage.tokenUsage).toEqual({
77+
promptTokens: 11,
78+
completionTokens: 7,
79+
totalTokens: 18,
80+
});
81+
});
82+
83+
it("retries after invalid JSON and succeeds", async () => {
84+
const outputSchema = z.object({
85+
city: z.string(),
86+
});
87+
const onRetry = vi.fn();
88+
89+
const { rllm, completeMock } = createTestRLLMWithMock([
90+
{
91+
content: '{"city":"Tel Aviv"',
92+
usage: { promptTokens: 5, completionTokens: 4, totalTokens: 9 },
93+
},
94+
{
95+
content: '{"city":"Tel Aviv"}',
96+
usage: { promptTokens: 6, completionTokens: 4, totalTokens: 10 },
97+
},
98+
]);
99+
100+
const result = await rllm.generateObject("Return city", { outputSchema }, {
101+
maxRetries: 2,
102+
onRetry,
103+
});
104+
105+
expect(result.object).toEqual({ city: "Tel Aviv" });
106+
expect(result.attempts).toBe(2);
107+
expect(result.usage.totalCalls).toBe(2);
108+
expect(result.usage.tokenUsage).toEqual({
109+
promptTokens: 11,
110+
completionTokens: 8,
111+
totalTokens: 19,
112+
});
113+
114+
expect(onRetry).toHaveBeenCalledTimes(1);
115+
expect(onRetry.mock.calls[0]?.[0].errorType).toBe("json_parse");
116+
expect(completeMock).toHaveBeenCalledTimes(2);
117+
});
118+
119+
it("retries after schema mismatch and succeeds", async () => {
120+
const outputSchema = z.object({
121+
status: z.enum(["ok", "error"]),
122+
count: z.number(),
123+
});
124+
const onRetry = vi.fn();
125+
126+
const { rllm } = createTestRLLMWithMock([
127+
{
128+
content: '{"status":"ok","count":"3"}',
129+
usage: { promptTokens: 8, completionTokens: 5, totalTokens: 13 },
130+
},
131+
{
132+
content: '{"status":"ok","count":3}',
133+
usage: { promptTokens: 9, completionTokens: 5, totalTokens: 14 },
134+
},
135+
]);
136+
137+
const result = await rllm.generateObject("Return status and count", { outputSchema }, {
138+
maxRetries: 2,
139+
onRetry,
140+
});
141+
142+
expect(result.object).toEqual({ status: "ok", count: 3 });
143+
expect(result.attempts).toBe(2);
144+
expect(onRetry).toHaveBeenCalledTimes(1);
145+
expect(onRetry.mock.calls[0]?.[0].errorType).toBe("schema_validation");
146+
expect(onRetry.mock.calls[0]?.[0].validationIssues?.length).toBeGreaterThan(0);
147+
});
148+
149+
it("throws actionable error after exhausting retries", async () => {
150+
const outputSchema = z.object({
151+
id: z.string(),
152+
});
153+
154+
const { rllm } = createTestRLLMWithMock([
155+
{
156+
content: '{"id":123}',
157+
usage: { promptTokens: 3, completionTokens: 2, totalTokens: 5 },
158+
},
159+
{
160+
content: '{"id":456}',
161+
usage: { promptTokens: 3, completionTokens: 2, totalTokens: 5 },
162+
},
163+
]);
164+
165+
await expect(
166+
rllm.generateObject("Return id", { outputSchema }, { maxRetries: 1 })
167+
).rejects.toThrow(/generateObject failed after 2 attempt/);
168+
});
169+
170+
it("fails fast when input does not satisfy inputSchema", async () => {
171+
const outputSchema = z.object({ answer: z.string() });
172+
const inputSchema = z.object({ age: z.number() });
173+
const { rllm, completeMock } = createTestRLLMWithMock([
174+
{
175+
content: '{"answer":"ok"}',
176+
usage: { promptTokens: 1, completionTokens: 1, totalTokens: 2 },
177+
},
178+
]);
179+
180+
await expect(
181+
rllm.generateObject("Use input", {
182+
input: { age: "not-a-number" } as unknown as { age: number },
183+
inputSchema,
184+
outputSchema,
185+
})
186+
).rejects.toThrow(/input failed inputSchema validation/);
187+
188+
expect(completeMock).toHaveBeenCalledTimes(0);
189+
});
190+
});

0 commit comments

Comments
 (0)