Skip to content

Commit 2102ad4

Browse files
martinibachmartinibach
committed
structured output and event based triggers
1 parent 2462d2a commit 2102ad4

File tree

7 files changed

+276
-122
lines changed

7 files changed

+276
-122
lines changed

‎agent-todo/src/agents/agent.ts

+96-58
Original file line numberDiff line numberDiff line change
@@ -5,16 +5,25 @@ import {
55
log,
66
step,
77
childExecute,
8+
agentInfo,
89
} from "@restackio/ai/agent";
10+
import { z } from "zod";
911
import * as functions from "../functions";
10-
import { executeTodoWorkflow } from "../workflows/executeTodo";
12+
import { executeTodoWorkflow, ExecuteTodoSchema } from "../workflows/executeTodo";
13+
14+
const CreateTodoSchema = z.object({
15+
todoTitle: z.string().min(1),
16+
});
1117

1218
export type EndEvent = {
1319
end: boolean;
1420
};
1521

22+
// Define events
1623
export const messagesEvent = defineEvent<functions.Message[]>("messages");
1724
export const endEvent = defineEvent("end");
25+
export const createTodoEvent = defineEvent("createTodo");
26+
export const executeTodoWorkflowEvent = defineEvent("executeTodoWorkflow");
1827

1928
type agentTodoOutput = {
2029
messages: functions.Message[];
@@ -24,82 +33,111 @@ export async function agentTodo(): Promise<agentTodoOutput> {
2433
let endReceived = false;
2534
let agentMessages: functions.Message[] = [];
2635

27-
const tools = await step<typeof functions>({}).getTools();
28-
2936
onEvent(messagesEvent, async ({ messages }: { messages: functions.Message[] }) => {
3037
agentMessages.push(...messages);
3138

3239
const result = await step<typeof functions>({}).llmChat({
3340
messages: agentMessages,
34-
tools,
41+
systemContent: "You are a helpful assistant that can create and execute todos.",
42+
model: "gpt-4.1-mini"
3543
});
3644

3745
agentMessages.push(result);
3846

39-
if (result.tool_calls) {
40-
log.info("result.tool_calls", { result });
41-
for (const toolCall of result.tool_calls) {
42-
switch (toolCall.function.name) {
43-
case "createTodo":
44-
log.info("createTodo", { toolCall });
45-
const toolResult = await step<typeof functions>({}).createTodo(
46-
JSON.parse(toolCall.function.arguments)
47-
);
48-
49-
agentMessages.push({
50-
role: "tool",
51-
tool_call_id: toolCall.id,
52-
content: toolResult,
53-
});
54-
55-
const toolChatResult = await step<typeof functions>({}).llmChat({
56-
messages: agentMessages,
57-
tools,
58-
});
59-
60-
agentMessages.push(toolChatResult);
61-
62-
break;
63-
case "executeTodoWorkflow":
64-
log.info("executeTodoWorkflow", { toolCall });
65-
const workflowId = `executeTodoWorkflow-${new Date().getTime()}`;
66-
const workflowResult = await childExecute({
67-
child: executeTodoWorkflow,
68-
childId: workflowId,
69-
input: JSON.parse(toolCall.function.arguments),
70-
taskQueue: "todo-workflows",
71-
});
72-
73-
agentMessages.push({
74-
role: "tool",
75-
tool_call_id: toolCall.id,
76-
content: JSON.stringify(workflowResult),
77-
});
78-
79-
const toolWorkflowResult = await step<typeof functions>({}).llmChat(
80-
{
81-
messages: agentMessages,
82-
tools,
83-
}
84-
);
85-
86-
agentMessages.push(toolWorkflowResult);
87-
88-
break;
89-
default:
90-
break;
91-
}
47+
if (result.structured_data?.type === "function_call") {
48+
const { function_name, function_arguments } = result.structured_data;
49+
log.info(function_name, { function_arguments });
50+
51+
switch (function_name) {
52+
case "createTodo":
53+
await step<typeof functions>({}).sendEvent({
54+
agentId: agentInfo().workflowId,
55+
runId: agentInfo().runId,
56+
eventName: "createTodo",
57+
eventInput: { function_arguments }
58+
});
59+
break;
60+
61+
case "executeTodoWorkflow":
62+
const args = ExecuteTodoSchema.parse(function_arguments);
63+
await step<typeof functions>({}).sendEvent({
64+
agentId: agentInfo().workflowId,
65+
runId: agentInfo().runId,
66+
eventName: "executeTodoWorkflow",
67+
eventInput: { workflowId: `executeTodoWorkflow-${Date.now()}`, args }
68+
});
69+
break;
9270
}
9371
}
72+
9473
return agentMessages;
9574
});
9675

76+
onEvent(createTodoEvent, async (data: any) => {
77+
try {
78+
const parsedArgs = CreateTodoSchema.parse(data.function_arguments);
79+
const stepResult = await step<typeof functions>({}).createTodo(parsedArgs);
80+
81+
await step<typeof functions>({}).sendEvent({
82+
agentId: agentInfo().workflowId,
83+
runId: agentInfo().runId,
84+
eventName: "messages",
85+
eventInput: {
86+
messages: [{
87+
role: "system",
88+
content: `Function executed: ${stepResult}`
89+
}]
90+
}
91+
});
92+
} catch (error: any) {
93+
log.error("Error in createTodo", { error: error.toString() });
94+
agentMessages.push({
95+
role: "system",
96+
content: `Error handling createTodo: ${error.message}`
97+
});
98+
}
99+
});
100+
101+
onEvent(executeTodoWorkflowEvent, async (data: any) => {
102+
try {
103+
const { workflowId, args } = data;
104+
const workflowResult = await childExecute({
105+
child: executeTodoWorkflow,
106+
childId: workflowId,
107+
input: args,
108+
taskQueue: "todo-workflows"
109+
});
110+
111+
await step<typeof functions>({}).sendEvent({
112+
agentId: agentInfo().workflowId,
113+
runId: agentInfo().runId,
114+
eventName: "messages",
115+
eventInput: {
116+
messages: [{
117+
role: "system",
118+
content: `Todo workflow executed successfully! Status: ${workflowResult.status} Details: ${workflowResult.details} ID: ${workflowResult.todoId}`
119+
}]
120+
}
121+
});
122+
} catch (workflowError: any) {
123+
log.error("Workflow execution failed", {
124+
error: workflowError.toString(),
125+
stack: workflowError.stack
126+
});
127+
128+
agentMessages.push({
129+
role: "system",
130+
content: `The workflow execution failed: ${workflowError.message}`
131+
});
132+
}
133+
});
134+
97135
onEvent(endEvent, async () => {
98136
endReceived = true;
99137
});
100138

101139
await condition(() => endReceived);
102-
103140
log.info("end condition met");
141+
104142
return { messages: agentMessages };
105143
}

‎agent-todo/src/functions/createTodo.ts

+9-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,14 @@
11
import { log } from "@restackio/ai/function";
2+
import { z } from "zod";
23

3-
export const createTodo = async ({ todoTitle }: { todoTitle: string }) => {
4+
export const CreateTodoSchema = z.object({
5+
todoTitle: z.string().min(1),
6+
});
7+
8+
export type CreateTodoInput = z.infer<typeof CreateTodoSchema>;
9+
export type CreateTodoOutput = string;
10+
11+
export const createTodo = async ({ todoTitle }: CreateTodoInput): Promise<CreateTodoOutput> => {
412
const todo_id = `todo-${Math.floor(Math.random() * 10000)}`;
513
log.info("createTodo", { todo_id, todoTitle });
614
return `Created the todo '${todoTitle}' with id: ${todo_id}`;

‎agent-todo/src/functions/index.ts

+1-1
Original file line numberDiff line numberDiff line change
@@ -2,5 +2,5 @@ export * from "./llmChat";
22
export * from "./getTools";
33
export * from "./createTodo";
44
export * from "./getRandom";
5-
export * from "./toolTypes";
65
export * from "./getResult";
6+
export * from "./sendEvent";

‎agent-todo/src/functions/llmChat.ts

+81-24
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,11 @@ import { FunctionFailure, log } from "@restackio/ai/function";
22
import {
33
ChatCompletionCreateParamsNonStreaming,
44
ChatCompletionSystemMessageParam,
5-
ChatCompletionTool,
6-
ChatCompletionToolMessageParam,
75
ChatCompletionUserMessageParam,
86
ChatCompletionAssistantMessageParam,
9-
ChatCompletionMessage,
7+
ChatCompletionToolMessageParam,
108
} from "openai/resources/chat/completions";
9+
import { z } from "zod";
1110

1211
import { openaiClient } from "../utils/client";
1312

@@ -17,43 +16,101 @@ export type Message =
1716
| ChatCompletionAssistantMessageParam
1817
| ChatCompletionToolMessageParam;
1918

20-
export type OpenAIChatInput = {
21-
systemContent?: string;
22-
model?: string;
23-
messages: Message[];
24-
tools?: ChatCompletionTool[];
25-
};
19+
// Define schema for structured response format
20+
export const ResponseSchema = z.discriminatedUnion("type", [
21+
z.object({
22+
type: z.literal("text"),
23+
content: z.string()
24+
}),
25+
z.object({
26+
type: z.literal("function_call"),
27+
function_name: z.string(),
28+
function_arguments: z.record(z.any())
29+
})
30+
]);
31+
32+
export type StructuredResponse = z.infer<typeof ResponseSchema>;
33+
34+
// Input schema
35+
export const LLMChatInputSchema = z.object({
36+
systemContent: z.string().optional().default(""),
37+
model: z.string().optional().default("gpt-4.1-mini"),
38+
messages: z.array(z.any()),
39+
});
40+
41+
// Output schema
42+
export const LLMChatOutputSchema = z.object({
43+
role: z.literal("assistant"),
44+
content: z.string().nullable(),
45+
structured_data: ResponseSchema.optional(),
46+
});
47+
48+
export type LLMChatInput = z.infer<typeof LLMChatInputSchema>;
49+
export type LLMChatOutput = z.infer<typeof LLMChatOutputSchema>;
50+
51+
// System prompt for structured output
52+
const structuredOutputPrompt = `You are a helpful assistant that can create and execute todos.
53+
Respond with valid JSON in one of these formats:
54+
- For general questions: {"type": "text", "content": "Your response"}
55+
- For actions: {"type": "function_call", "function_name": "createTodo or executeTodoWorkflow", "function_arguments": {...}}`;
2656

2757
export const llmChat = async ({
2858
systemContent = "",
2959
model = "gpt-4.1-mini",
3060
messages,
31-
tools,
32-
}: OpenAIChatInput): Promise<ChatCompletionMessage> => {
61+
}: LLMChatInput): Promise<LLMChatOutput> => {
3362
try {
3463
const openai = openaiClient({});
3564

65+
// Combine system prompts if provided
66+
const finalSystemContent = systemContent
67+
? `${structuredOutputPrompt}\n\n${systemContent}`
68+
: structuredOutputPrompt;
69+
70+
// Set up response format for JSON
71+
const responseFormat = {
72+
type: "json_object" as const
73+
};
74+
75+
// Chat parameters
3676
const chatParams: ChatCompletionCreateParamsNonStreaming = {
37-
messages: [
38-
...(systemContent
39-
? [{ role: "system" as const, content: systemContent }]
40-
: []),
41-
...(messages ?? []),
42-
],
77+
messages: [{ role: "system", content: finalSystemContent }, ...messages],
4378
model,
44-
tools,
79+
response_format: responseFormat,
4580
};
4681

47-
log.debug("OpenAI chat completion params", {
48-
chatParams,
49-
});
82+
log.debug("OpenAI chat completion params", { chatParams });
5083

5184
const completion = await openai.chat.completions.create(chatParams);
52-
5385
const message = completion.choices[0].message;
5486

55-
return message;
87+
// Parse structured data from JSON response
88+
let structuredData;
89+
if (message.content) {
90+
try {
91+
const parsedData = JSON.parse(message.content);
92+
// Validate against our schema
93+
const validationResult = ResponseSchema.safeParse(parsedData);
94+
if (validationResult.success) {
95+
structuredData = validationResult.data;
96+
log.debug("Structured response validated", { structuredData });
97+
} else {
98+
log.error("Invalid JSON structure", {
99+
errors: validationResult.error.errors,
100+
content: message.content
101+
});
102+
}
103+
} catch (error) {
104+
log.error("Failed to parse JSON response", { content: message.content });
105+
}
106+
}
107+
108+
return {
109+
role: "assistant",
110+
content: message.content,
111+
structured_data: structuredData,
112+
};
56113
} catch (error) {
57-
throw FunctionFailure.nonRetryable(`Error OpenAI chat: ${error}`);
114+
throw FunctionFailure.nonRetryable(`Error in OpenAI chat: ${error}`);
58115
}
59116
};

‎agent-todo/src/functions/sendEvent.ts

+27
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
import { client } from "../client";
2+
import { log } from "@restackio/ai/function";
3+
4+
export type SendEventInput = {
5+
agentId: string;
6+
runId: string;
7+
eventName: string;
8+
eventInput?: Record<string, any>;
9+
};
10+
11+
export async function sendEvent(input: SendEventInput): Promise<void> {
12+
try {
13+
await client.sendAgentEvent({
14+
event: {
15+
name: input.eventName,
16+
input: input.eventInput,
17+
},
18+
agent: {
19+
agentId: input.agentId,
20+
runId: input.runId,
21+
},
22+
});
23+
} catch (error) {
24+
log.error("Error sending agent event", { error, input });
25+
throw error;
26+
}
27+
}

0 commit comments

Comments
 (0)