Skip to content

Commit 3003b2c

Browse files
olzhik11olzhik11
authored
feat: add structured output, dont reset msg content (#957)
* feat: add structured output, dont reset msg content * feat: refactor, update ui * feat: add structured output conditionally * feat: move narrow type to chat specifically * feat: broaden type
1 parent f14feb9 commit 3003b2c

File tree

12 files changed

+299
-89
lines changed

12 files changed

+299
-89
lines changed

‎frontend/app/api/projects/[projectId]/chat/route.ts‎

Lines changed: 14 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,14 @@
1-
import { handleChatGeneration, PlaygroundParamsSchema } from "@/lib/actions/chat";
1+
import { NextResponse } from "next/server";
2+
import { prettifyError, ZodError } from "zod/v4";
3+
4+
import { handleChatGeneration } from "@/lib/actions/chat";
25
import { parseSystemMessages } from "@/lib/playground/utils";
36

47
export async function POST(req: Request, props: { params: Promise<{ projectId: string }> }) {
58
try {
69
const body = await req.json();
710
const { projectId } = await props.params;
811

9-
// Convert form messages to ModelMessages
1012
const convertedMessages = body.messages ? parseSystemMessages(body.messages) : [];
1113

1214
const params = {
@@ -15,27 +17,22 @@ export async function POST(req: Request, props: { params: Promise<{ projectId: s
1517
projectId,
1618
};
1719

18-
const parseResult = PlaygroundParamsSchema.safeParse(params);
19-
20-
if (!parseResult.success) {
21-
return new Response(JSON.stringify(parseResult.error), { status: 400 });
22-
}
23-
2420
const result = await handleChatGeneration({
2521
...params,
2622
abortSignal: req.signal,
2723
});
2824

29-
return new Response(JSON.stringify(result));
30-
} catch (e) {
31-
return new Response(
32-
JSON.stringify({
33-
error: e instanceof Error ? e.message : "Internal server error.",
34-
details: e instanceof Error ? e.name : "Unknown error",
35-
}),
25+
return NextResponse.json(result);
26+
} catch (error) {
27+
if (error instanceof ZodError) {
28+
return NextResponse.json({ error: prettifyError(error) }, { status: 400 });
29+
}
30+
31+
return NextResponse.json(
3632
{
37-
status: 500,
38-
}
33+
error: error instanceof Error ? error.message : "Internal server error.",
34+
},
35+
{ status: 500 }
3936
);
4037
}
4138
}

‎frontend/components/playground/messages/llm-select.tsx‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -55,12 +55,12 @@ const LlmSelect = ({ apiKeys, disabled, onChange, value, className }: LlmSelectN
5555
return (
5656
<DropdownMenu>
5757
<DropdownMenuTrigger value={value} asChild>
58-
<Button disabled={disabled} className={cn("focus-visible:ring-0", className)} variant="outline">
59-
<span className="mr-2">{providerIconMap[value.split(":")[0] as Provider]}</span>
60-
<span className="truncate mr-2 py-0.5">
58+
<Button disabled={disabled} className={cn("focus-visible:ring-0 text-xs px-2", className)} variant="outline">
59+
<span className="mr-1">{providerIconMap[value.split(":")[0] as Provider]}</span>
60+
<span className="truncate mr-1 py-0.5">
6161
{providers.flatMap((p) => p.models).find((m) => m.id === value)?.label ?? "Select model"}
6262
</span>
63-
<ChevronDown className="ml-auto" size={16} />
63+
<ChevronDown className="ml-auto w-3.5 h-3.5" size={16} />
6464
</Button>
6565
</DropdownMenuTrigger>
6666
<DropdownMenuContent align="start">

‎frontend/components/playground/messages/message.tsx‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -69,11 +69,7 @@ const Message = ({ insert, remove, update, index, deletable = true }: MessagePro
6969

7070
const handleUpdateRole =
7171
(onChange: ControllerRenderProps["onChange"]) => (value: PlaygroundForm["messages"]["0"]["role"]) => {
72-
if (value === "system") {
73-
update(index, { content: [{ type: "text", text: "" }], role: value });
74-
} else {
75-
onChange(value);
76-
}
72+
onChange(value);
7773
};
7874

7975
return (

‎frontend/components/playground/messages/params-popover.tsx‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ const ParamsPopover = ({ className }: ParamsPopoverProps) => {
3030
variant="outline"
3131
className={cn(className, "w-8 h-8 self-end")}
3232
>
33-
<SlidersHorizontal className="w-4 h-4" />
33+
<SlidersHorizontal className="w-3.5 h-3.5" />
3434
</Button>
3535
</PopoverTrigger>
3636
</TooltipTrigger>

‎frontend/components/playground/messages/structured-output-sheet.tsx‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
import { TooltipPortal } from "@radix-ui/react-tooltip";
2+
import { BracesIcon, X } from "lucide-react";
3+
import { PropsWithChildren, useCallback } from "react";
4+
import { Controller, useFormContext } from "react-hook-form";
5+
6+
import { Button } from "@/components/ui/button";
7+
import CodeHighlighter from "@/components/ui/code-highlighter/index";
8+
import { Sheet, SheetContent, SheetHeader, SheetTitle, SheetTrigger } from "@/components/ui/sheet";
9+
import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
10+
import { PlaygroundForm } from "@/lib/playground/types";
11+
import { cn } from "@/lib/utils";
12+
13+
const exampleStructuredOutput = {
14+
type: "object",
15+
properties: {
16+
name: {
17+
type: "string",
18+
description: "The name of the person",
19+
},
20+
age: {
21+
type: "number",
22+
description: "The age of the person",
23+
},
24+
},
25+
required: ["name", "age"],
26+
additionalProperties: false,
27+
};
28+
29+
export default function StructuredOutputSheet({
30+
children,
31+
className,
32+
}: PropsWithChildren<{
33+
className?: string;
34+
}>) {
35+
const {
36+
control,
37+
watch,
38+
setValue,
39+
formState: { errors },
40+
} = useFormContext<PlaygroundForm>();
41+
42+
const structuredOutput = watch("structuredOutput");
43+
const model = watch("model");
44+
45+
const renderTrigger = useCallback(() => {
46+
if (!structuredOutput) {
47+
return (
48+
<SheetTrigger asChild>
49+
<Button
50+
disabled={!model}
51+
variant="outline"
52+
size="icon"
53+
className={cn("focus-visible:ring-0 w-8 h-8 p-2", className)}
54+
>
55+
<BracesIcon className="w-3.5 h-3.5" />
56+
</Button>
57+
</SheetTrigger>
58+
);
59+
}
60+
61+
return (
62+
<div className="flex flex-row [&>*:first-child]:border-r-0 [&>*:first-child]:rounded-l [&>*:first-child]:rounded-r-none [&>*:last-child]:rounded-r [&>*:last-child]:rounded-l-none">
63+
<SheetTrigger asChild>
64+
<Button
65+
disabled={!model}
66+
variant="outlinePrimary"
67+
size="icon"
68+
className={cn("focus-visible:ring-0 w-8 h-8 p-2", className)}
69+
>
70+
<BracesIcon className="w-3.5 h-3.5" />
71+
</Button>
72+
</SheetTrigger>
73+
<Button
74+
onClick={() => setValue("structuredOutput", undefined)}
75+
className="w-8 h-8"
76+
variant="outlinePrimary"
77+
size="icon"
78+
>
79+
<X className="h-3.5 w-3.5" />
80+
</Button>
81+
</div>
82+
);
83+
}, [className, model, setValue, structuredOutput]);
84+
85+
return (
86+
<Sheet>
87+
<Tooltip>
88+
<TooltipTrigger asChild>{children || renderTrigger()}</TooltipTrigger>
89+
<TooltipPortal>
90+
<TooltipContent>Structured Output</TooltipContent>
91+
</TooltipPortal>
92+
</Tooltip>
93+
<SheetContent side="right" className="min-w-[50vw] w-full flex flex-col gap-4 p-4">
94+
<SheetHeader>
95+
<SheetTitle>Structured Output</SheetTitle>
96+
</SheetHeader>
97+
<div className="flex flex-col gap-2 flex-1 min-h-0">
98+
<div className="flex items-center gap-2 justify-between">
99+
<span className="text-sm">JSON Schema</span>
100+
<Button
101+
onClick={() => setValue("structuredOutput", JSON.stringify(exampleStructuredOutput, null, 2))}
102+
className="text-primary text-sm p-0"
103+
variant="ghost"
104+
>
105+
<span className="text-sm">Insert example</span>
106+
</Button>
107+
</div>
108+
<div className="p-1 flex flex-1 overflow-hidden">
109+
<Controller
110+
render={({ field: { onChange } }) => (
111+
<CodeHighlighter
112+
onChange={(v) => onChange(v)}
113+
codeEditorClassName="rounded-b"
114+
className={cn("rounded h-full", {
115+
"border border-destructive/75": errors.structuredOutput?.message,
116+
})}
117+
defaultMode="json"
118+
modes={["JSON"]}
119+
value={structuredOutput ?? ""}
120+
/>
121+
)}
122+
name="structuredOutput"
123+
control={control}
124+
rules={{
125+
validate: (value) => {
126+
try {
127+
if (!value) {
128+
return true;
129+
}
130+
JSON.parse(value);
131+
return true;
132+
} catch (e) {
133+
return "Invalid JSON structure";
134+
}
135+
},
136+
}}
137+
/>
138+
</div>
139+
<span className="text-xs text-secondary-foreground">
140+
Define a JSON Schema to structure the model&apos;s output.
141+
</span>
142+
</div>
143+
</SheetContent>
144+
</Sheet>
145+
);
146+
}

‎frontend/components/playground/messages/tools-sheet.tsx‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@ export default function ToolsSheet({
8484
size="icon"
8585
className={cn("focus-visible:ring-0 w-8 h-8 p-2", className)}
8686
>
87-
<Bolt className="w-4 h-4" />
87+
<Bolt className="w-3.5 h-3.5" />
8888
</Button>
8989
</SheetTrigger>
9090
);
@@ -95,16 +95,16 @@ export default function ToolsSheet({
9595
<SheetTrigger asChild>
9696
<Button
9797
disabled={!model}
98-
variant="outline"
98+
variant="outlinePrimary"
9999
size="icon"
100100
className={cn("focus-visible:ring-0 h-8 w-fit p-2", className)}
101101
>
102-
<Bolt className="w-4 h-4" />
103-
<span className="ml-2">{pluralize(toolsCount, "tool", "tools")}</span>
102+
<Bolt className="w-3.5 h-3.5" />
103+
<span className="ml-1 text-xs ">{pluralize(toolsCount, "tool", "tools")}</span>
104104
</Button>
105105
</SheetTrigger>
106-
<Button onClick={() => setValue("tools", "")} className="w-8 h-8" variant="outline" size="icon">
107-
<X className="h-3.5 w-3.5" />
106+
<Button onClick={() => setValue("tools", "")} className="w-8 h-8" variant="outlinePrimary" size="icon">
107+
<X className="w-3.5 h-3.5" />
108108
</Button>
109109
</div>
110110
);

‎frontend/components/playground/playground-panel.tsx‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import { prettifyError } from "zod/v4";
1111
import Messages from "@/components/playground/messages";
1212
import LlmSelect from "@/components/playground/messages/llm-select";
1313
import ParamsPopover from "@/components/playground/messages/params-popover";
14+
import StructuredOutputSheet from "@/components/playground/messages/structured-output-sheet";
1415
import ToolsSheet from "@/components/playground/messages/tools-sheet";
1516
import PlaygroundHistoryTable from "@/components/playground/playground-history-table";
1617
import { usePlaygroundOutput } from "@/components/playground/playground-output";
@@ -85,6 +86,7 @@ export default function PlaygroundPanel({
8586
providerOptions: form.providerOptions,
8687
tools: form.tools,
8788
toolChoice: form.toolChoice,
89+
structuredOutput: form.structuredOutput,
8890
}),
8991
});
9092

@@ -161,25 +163,26 @@ export default function PlaygroundPanel({
161163
/>
162164
<ParamsPopover />
163165
<ToolsSheet />
166+
<StructuredOutputSheet />
164167
<Button
165168
variant={history ? "outlinePrimary" : "outline"}
166169
size="sm"
167170
onClick={() => setHistory(!history)}
168171
className="h-8 w-fit px-2"
169172
>
170-
<History className="w-4 h-4 mr-1" />
173+
<History className="w-3.5 h-3.5 mr-1" />
171174
History
172175
</Button>
173176
{isLoading ? (
174177
<Button variant="outlinePrimary" onClick={abortRequest} className="ml-auto h-8 w-fit px-2">
175-
<Square className="w-4 h-4 mr-2" />
176-
<span className="mr-2">Stop</span>
178+
<Square className="w-3.5 h-3.5 mr-1" />
179+
<span className="mr-2 text-xs">Stop</span>
177180
<Loader className="animate-spin w-4 h-4" />
178181
</Button>
179182
) : (
180183
<Button onClick={handleSubmit(submit)} className="ml-auto h-8 w-fit px-2">
181-
<PlayIcon className="w-4 h-4 mr-2" />
182-
<span className="mr-2">Run</span>
184+
<PlayIcon className="w-3.5 h-3.5 mr-1" />
185+
<span className="mr-2 text-xs">Run</span>
183186
<div className="text-center text-xs opacity-75">⌘ + ⏎</div>
184187
</Button>
185188
)}

‎frontend/components/playground/playground.tsx‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ export default function Playground({ playground }: { playground: PlaygroundType
6868
: getDefaultThinkingModelProviderOptions(playground.modelId as PlaygroundForm["model"]),
6969
tools: JSON.stringify(playground.tools),
7070
toolChoice: playground.toolChoice as PlaygroundForm["toolChoice"],
71+
structuredOutput: playground.outputSchema ?? undefined,
7172
});
7273
}
7374
resetOutput();
@@ -87,6 +88,7 @@ export default function Playground({ playground }: { playground: PlaygroundType
8788
maxTokens: form.maxTokens,
8889
temperature: form.temperature,
8990
providerOptions: form.providerOptions,
91+
outputSchema: form.structuredOutput,
9092
}),
9193
});
9294
} catch (e) {

‎frontend/components/playground/utils.tsx‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,7 @@ export const getPlaygroundConfig = (
217217
modelId: string;
218218
maxTokens?: number;
219219
temperature?: number;
220+
outputSchema?: string;
220221
} => {
221222
const model = get(span, ["attributes", "gen_ai.response.model"]) as string | undefined;
222223

@@ -229,7 +230,12 @@ export const getPlaygroundConfig = (
229230
const toolChoice = get(span, ["attributes", "ai.prompt.toolChoice"]);
230231
const parsedToolChoice = parseToolChoiceFromSpan(toolChoice);
231232

232-
const referenceModel = model && existingModels.find((existingModel) => model.includes(existingModel));
233+
const outputSchema = get(span, ["attributes", "gen_ai.request.structured_output_schema"]) as string | undefined;
234+
235+
const referenceModel =
236+
model &&
237+
(existingModels.find((existingModel) => model === existingModel) ||
238+
existingModels.filter((existingModel) => model.includes(existingModel)).sort((a, b) => b.length - a.length)[0]);
233239
const foundModel = models.find((m) => m.name === referenceModel)?.id;
234240

235241
const result = {
@@ -238,6 +244,7 @@ export const getPlaygroundConfig = (
238244
toolChoice: parsedToolChoice || (parsedTools ? "auto" : undefined),
239245
maxTokens: get(span, ["attributes", "gen_ai.request.max_tokens"], defaultMaxTokens),
240246
temperature: get(span, ["attributes", "gen_ai.request.temperature"], defaultTemperature),
247+
outputSchema,
241248
};
242249

243250
return pickBy(result, (value) => value !== undefined) as typeof result;

0 commit comments

Comments
 (0)