Skip to content

Commit 4288e91

Browse files
bfollingtonbfollington
authored
Refactor spellcaster routes (commontoolsinc#498)
* Refactor spellcaster routes Implement fast path: imagine * Format pass
1 parent 4f5e887 commit 4288e91

File tree

8 files changed

+579
-333
lines changed

8 files changed

+579
-333
lines changed

‎typescript/packages/toolshed/routes/ai/spell/fulfill.ts‎ renamed to ‎typescript/packages/toolshed/routes/ai/spell/handlers/fulfill.ts‎

Lines changed: 99 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,85 @@
1-
import { getAllBlobs, getBlob } from "@/routes/ai/spell/behavior/effects.ts";
2-
import { generateText } from "@/lib/llm.ts";
3-
import { performSearch } from "@/routes/ai/spell/behavior/search.ts";
4-
import { checkSchemaMatch } from "@/lib/schema-match.ts";
1+
import * as HttpStatusCodes from "stoker/http-status-codes";
2+
import { z } from "zod";
3+
import { getAllBlobs } from "@/routes/ai/spell/behavior/effects.ts";
4+
5+
import type { AppRouteHandler } from "@/lib/types.ts";
6+
import type { FulfillSchemaRoute } from "@/routes/ai/spell/spell.routes.ts";
7+
import { Spell } from "@/routes/ai/spell/spell.ts";
8+
import { performSearch } from "../behavior/search.ts";
59
import { Logger } from "@/lib/prefixed-logger.ts";
6-
import {
7-
ProcessSchemaRequest,
8-
ProcessSchemaResponse,
9-
} from "@/routes/ai/spell/spell.handlers.ts";
10+
import { candidates } from "@/routes/ai/spell/caster.ts";
11+
import { CasterSchemaRoute } from "@/routes/ai/spell/spell.routes.ts";
12+
import { processSpellSearch } from "@/routes/ai/spell/behavior/spell-search.ts";
13+
import { captureException } from "@sentry/deno";
14+
import { areSchemaCompatible } from "@/routes/ai/spell/schema-compatibility.ts";
15+
16+
import { generateText } from "@/lib/llm.ts";
1017
import {
1118
decomposeSchema,
1219
findExactMatches,
1320
findFragmentMatches,
1421
reassembleFragments,
1522
SchemaFragment,
1623
} from "@/routes/ai/spell/schema.ts";
24+
import { extractJSON } from "@/routes/ai/spell/json.ts";
25+
26+
export const FulfillSchemaRequestSchema = z.object({
27+
schema: z.record(
28+
z
29+
.string()
30+
.or(
31+
z.number().or(z.boolean().or(z.array(z.any()).or(z.record(z.any())))),
32+
),
33+
).openapi({
34+
example: {
35+
title: { type: "string" },
36+
url: { type: "string" },
37+
},
38+
}),
39+
tags: z.array(z.string()).optional(),
40+
many: z.boolean().optional(),
41+
prompt: z.string().optional(),
42+
options: z
43+
.object({
44+
format: z.enum(["json", "yaml"]).optional(),
45+
validate: z.boolean().optional(),
46+
maxExamples: z.number().default(5).optional(),
47+
exact: z.boolean().optional(),
48+
})
49+
.optional(),
50+
});
51+
52+
export const FulfillSchemaResponseSchema = z.object({
53+
result: z.union([z.record(z.any()), z.array(z.record(z.any()))]),
54+
metadata: z.object({
55+
processingTime: z.number(),
56+
schemaFormat: z.string(),
57+
fragments: z.array(
58+
z.object({
59+
matches: z.array(
60+
z.object({
61+
key: z.string(),
62+
data: z.record(z.any()),
63+
similarity: z.number(),
64+
}),
65+
),
66+
path: z.array(z.string()),
67+
schema: z.record(z.any()),
68+
}),
69+
),
70+
reassembledExample: z.record(z.any()),
71+
tagMatchInfo: z.object({
72+
usedTags: z.any(),
73+
matchRanks: z.array(z.object({
74+
path: z.any(),
75+
matches: z.any(),
76+
})),
77+
}),
78+
}),
79+
});
80+
81+
export type FulfillSchemaRequest = z.infer<typeof FulfillSchemaRequestSchema>;
82+
export type FulfillSchemaResponse = z.infer<typeof FulfillSchemaResponseSchema>;
1783

1884
function calculateTagRank(
1985
data: Record<string, unknown>,
@@ -34,10 +100,10 @@ function calculateTagRank(
34100
}
35101

36102
export async function processSchema(
37-
body: ProcessSchemaRequest,
103+
body: FulfillSchemaRequest,
38104
logger: Logger,
39105
startTime: number,
40-
): Promise<ProcessSchemaResponse> {
106+
): Promise<FulfillSchemaResponse> {
41107
const tags = body.tags || [];
42108
logger.info(
43109
{ schema: body.schema, many: body.many, options: body.options, tags },
@@ -162,28 +228,6 @@ export async function processSchema(
162228
);
163229

164230
let result: Record<string, unknown> | Array<Record<string, unknown>>;
165-
function extractJSON(
166-
text: string,
167-
): Record<string, unknown> | Array<Record<string, unknown>> {
168-
try {
169-
// Try to extract from markdown code block first
170-
const markdownMatch = text.match(/```(?:json)?\s*([\s\S]*?)\s*```/);
171-
if (markdownMatch) {
172-
return JSON.parse(markdownMatch[1].trim());
173-
}
174-
175-
// If not in markdown, try to find JSON-like content
176-
const jsonMatch = text.match(/\{[\s\S]*\}|\[[\s\S]*\]/);
177-
if (jsonMatch) {
178-
return JSON.parse(jsonMatch[0].trim());
179-
}
180-
181-
// If no special formatting, try parsing the original text
182-
return JSON.parse(text.trim());
183-
} catch (error) {
184-
return {};
185-
}
186-
}
187231

188232
try {
189233
logger.debug("Parsing LLM response");
@@ -318,3 +362,26 @@ Respond with ${
318362
many ? "an array of valid JSON objects" : "a single valid JSON object"
319363
}.`;
320364
}
365+
366+
export const fulfill: AppRouteHandler<FulfillSchemaRoute> = async (c) => {
367+
const logger: Logger = c.get("logger");
368+
const body = (await c.req.json()) as FulfillSchemaRequest;
369+
const startTime = performance.now();
370+
371+
try {
372+
const response = await processSchema(body, logger, startTime);
373+
374+
logger.info(
375+
{ processingTime: response.metadata.processingTime },
376+
"Request completed",
377+
);
378+
return c.json(response, HttpStatusCodes.OK);
379+
} catch (error) {
380+
logger.error({ error }, "Error processing schema");
381+
captureException(error);
382+
return c.json(
383+
{ error: "Failed to process schema" },
384+
HttpStatusCodes.INTERNAL_SERVER_ERROR,
385+
);
386+
}
387+
};

‎typescript/packages/toolshed/routes/ai/spell/handlers/imagine.ts‎

Lines changed: 176 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,176 @@
1+
import { generateText } from "@/lib/llm.ts";
2+
import { Logger } from "@/lib/prefixed-logger.ts";
3+
import { extractJSON } from "@/routes/ai/spell/json.ts";
4+
import * as HttpStatusCodes from "stoker/http-status-codes";
5+
import { z } from "zod";
6+
import { captureException } from "@sentry/deno";
7+
import type { AppRouteHandler } from "@/lib/types.ts";
8+
import type { ImagineDataRoute } from "@/routes/ai/spell/spell.routes.ts";
9+
10+
export const ImagineDataRequestSchema = z.object({
11+
schema: z.record(
12+
z
13+
.string()
14+
.or(
15+
z.number().or(z.boolean().or(z.array(z.any()).or(z.record(z.any())))),
16+
),
17+
)
18+
.describe("JSON schema format to conform to")
19+
.openapi({
20+
example: {
21+
title: { type: "string" },
22+
url: { type: "string" },
23+
},
24+
}),
25+
model: z.string().default("claude-3-7-sonnet").describe(
26+
"The LLM to use for data generation",
27+
).openapi({ example: "claude-3-7-sonnet" }),
28+
prompt: z.string().optional().describe(
29+
"Guide data generation with a prompt",
30+
).openapi({ example: "Make it about cats" }),
31+
options: z
32+
.object({
33+
many: z.boolean().default(false).describe(
34+
"Whether to generate multiple results",
35+
),
36+
})
37+
.optional(),
38+
});
39+
40+
export const ImagineDataResponseSchema = z.object({
41+
result: z.union([z.record(z.any()), z.array(z.record(z.any()))]),
42+
metadata: z.object({
43+
processingTime: z.number(),
44+
}),
45+
});
46+
47+
export type ImagineDataRequest = z.infer<typeof ImagineDataRequestSchema>;
48+
export type ImagineDataResponse = z.infer<typeof ImagineDataResponseSchema>;
49+
50+
export async function processSchema(
51+
body: ImagineDataRequest,
52+
logger: Logger,
53+
startTime: number,
54+
): Promise<ImagineDataResponse> {
55+
logger.info(
56+
{ schema: body.schema, options: body.options },
57+
"Starting schema processing request",
58+
);
59+
60+
logger.debug("Constructing prompt with reassembled examples");
61+
const prompt = constructSchemaPrompt(
62+
body.schema,
63+
body.prompt,
64+
body?.options?.many,
65+
);
66+
67+
logger.info({ prompt }, "Sending request to LLM");
68+
const llmStartTime = performance.now();
69+
const llmResponse = await generateText({
70+
model: "claude-3-7-sonnet",
71+
system: body?.options?.many
72+
? "Generate realistic example data that fits the provided schema. Return valid JSON array with multiple objects. Each object must match the schema exactly and respect all descriptions and constraints."
73+
: "Generate realistic example data that fits the provided schema. Return a valid JSON object that matches the schema exactly and respects all descriptions and constraints.",
74+
stream: false,
75+
messages: [{ role: "user", content: prompt }],
76+
});
77+
logger.info(
78+
{ llmTime: Math.round(performance.now() - llmStartTime) },
79+
"Received LLM response",
80+
);
81+
82+
let result: Record<string, unknown> | Array<Record<string, unknown>>;
83+
84+
try {
85+
logger.debug("Parsing LLM response");
86+
result = extractJSON(llmResponse);
87+
logger.debug({ extractedJSON: result }, "Extracted JSON from response");
88+
89+
if (body?.options?.many && !Array.isArray(result)) {
90+
logger.debug("Converting single object to array for many=true");
91+
result = [result];
92+
}
93+
logger.info(
94+
{
95+
resultType: body?.options?.many ? "array" : "object",
96+
resultSize: body?.options?.many ? (result as Array<unknown>).length : 1,
97+
},
98+
"Successfully parsed LLM response",
99+
);
100+
} catch (error) {
101+
logger.error(
102+
{ error, response: llmResponse },
103+
"Failed to parse LLM response",
104+
);
105+
throw new Error("Failed to parse LLM response as JSON");
106+
}
107+
108+
const totalTime = Math.round(performance.now() - startTime);
109+
logger.info(
110+
{ totalTime },
111+
"Completed schema processing request",
112+
);
113+
114+
return {
115+
result,
116+
metadata: {
117+
processingTime: totalTime,
118+
},
119+
};
120+
}
121+
122+
function constructSchemaPrompt(
123+
schema: Record<string, unknown>,
124+
userPrompt?: string,
125+
many?: boolean,
126+
): string {
127+
const schemaStr = JSON.stringify(schema, null, 2);
128+
129+
return `# TASK
130+
${
131+
many
132+
? `Generate multiple objects that fit the requested schema based on the references provided.`
133+
: `Fit data into the requested schema based on the references provided.`
134+
}
135+
136+
# SCHEMA
137+
${schemaStr}
138+
139+
# INSTRUCTIONS
140+
1. ${
141+
many
142+
? `Generate an array of objects that strictly follow the schema structure`
143+
: `Generate an object that strictly follows the schema structure`
144+
}
145+
2. Return ONLY valid JSON ${many ? "array" : "object"} matching the schema
146+
147+
${userPrompt ? `# ADDITIONAL REQUIREMENTS\n${userPrompt}\n\n` : ""}
148+
149+
# RESPONSE FORMAT
150+
Respond with ${
151+
many ? "an array of valid JSON objects" : "a single valid JSON object"
152+
}.`;
153+
}
154+
155+
export const imagine: AppRouteHandler<ImagineDataRoute> = async (c) => {
156+
const logger: Logger = c.get("logger");
157+
const body = (await c.req.json()) as ImagineDataRequest;
158+
const startTime = performance.now();
159+
160+
try {
161+
const response = await processSchema(body, logger, startTime);
162+
163+
logger.info(
164+
{ processingTime: response.metadata.processingTime },
165+
"Request completed",
166+
);
167+
return c.json(response, HttpStatusCodes.OK);
168+
} catch (error) {
169+
logger.error({ error }, "Error processing schema");
170+
captureException(error);
171+
return c.json(
172+
{ error: "Failed to process schema" },
173+
HttpStatusCodes.INTERNAL_SERVER_ERROR,
174+
);
175+
}
176+
};

0 commit comments

Comments
 (0)