Skip to content

Commit 5a83bfc

Browse files
committed
cleanup simplify
1 parent 2102ad4 commit 5a83bfc

File tree

6 files changed

+131
-75
lines changed

6 files changed

+131
-75
lines changed

agent-todo/src/agents/agent.ts

+4-8
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,9 @@ import {
77
childExecute,
88
agentInfo,
99
} from "@restackio/ai/agent";
10-
import { z } from "zod";
1110
import * as functions from "../functions";
12-
import { executeTodoWorkflow, ExecuteTodoSchema } from "../workflows/executeTodo";
13-
14-
const CreateTodoSchema = z.object({
15-
todoTitle: z.string().min(1),
16-
});
11+
import { executeTodoWorkflow } from "../workflows/executeTodo";
12+
import { CreateTodoSchema, ExecuteTodoSchema } from "../functions/toolTypes";
1713

1814
export type EndEvent = {
1915
end: boolean;
@@ -54,7 +50,7 @@ export async function agentTodo(): Promise<agentTodoOutput> {
5450
agentId: agentInfo().workflowId,
5551
runId: agentInfo().runId,
5652
eventName: "createTodo",
57-
eventInput: { function_arguments }
53+
eventInput: function_arguments
5854
});
5955
break;
6056

@@ -75,7 +71,7 @@ export async function agentTodo(): Promise<agentTodoOutput> {
7571

7672
onEvent(createTodoEvent, async (data: any) => {
7773
try {
78-
const parsedArgs = CreateTodoSchema.parse(data.function_arguments);
74+
const parsedArgs = CreateTodoSchema.parse(data);
7975
const stepResult = await step<typeof functions>({}).createTodo(parsedArgs);
8076

8177
await step<typeof functions>({}).sendEvent({

agent-todo/src/functions/createTodo.ts

+2-9
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,7 @@
11
import { log } from "@restackio/ai/function";
2-
import { z } from "zod";
2+
import { CreateTodoInputType, CreateTodoOutput } from "./toolTypes";
33

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> => {
4+
export const createTodo = async ({ todoTitle }: CreateTodoInputType): Promise<CreateTodoOutput> => {
125
const todo_id = `todo-${Math.floor(Math.random() * 10000)}`;
136
log.info("createTodo", { todo_id, todoTitle });
147
return `Created the todo '${todoTitle}' with id: ${todo_id}`;

agent-todo/src/functions/getTools.ts

+3-3
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,17 @@
11
import { zodFunction } from "openai/helpers/zod";
22
import { executeTodoWorkflow } from "../workflows/executeTodo";
33
import { createTodo } from "./createTodo";
4-
import { CreateTodoInput, ExecuteTodoInput } from "./toolTypes";
4+
import { CreateTodoSchema, ExecuteTodoSchema } from "./toolTypes";
55

66
export const getTools = async () => {
77
const tools = [
88
zodFunction({
99
name: createTodo.name,
10-
parameters: CreateTodoInput,
10+
parameters: CreateTodoSchema,
1111
}),
1212
zodFunction({
1313
name: executeTodoWorkflow.name,
14-
parameters: ExecuteTodoInput,
14+
parameters: ExecuteTodoSchema,
1515
}),
1616
];
1717
return tools;

agent-todo/src/functions/llmChat.ts

+99-33
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import {
99
import { z } from "zod";
1010

1111
import { openaiClient } from "../utils/client";
12+
import { CreateTodoSchema, ExecuteTodoSchema } from "./toolTypes";
1213

1314
export type Message =
1415
| ChatCompletionSystemMessageParam
@@ -48,11 +49,71 @@ export const LLMChatOutputSchema = z.object({
4849
export type LLMChatInput = z.infer<typeof LLMChatInputSchema>;
4950
export type LLMChatOutput = z.infer<typeof LLMChatOutputSchema>;
5051

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": {...}}`;
52+
// Convert Zod schema to OpenAI function parameter schema
53+
function zodToJsonSchema(schema: z.ZodType): any {
54+
if (schema instanceof z.ZodObject) {
55+
const shape = schema._def.shape();
56+
const properties: Record<string, any> = {};
57+
const required: string[] = [];
58+
59+
Object.entries(shape).forEach(([key, value]) => {
60+
const zodField = value as z.ZodType;
61+
properties[key] = zodToJsonSchema(zodField);
62+
63+
if (!zodField.isOptional()) {
64+
required.push(key);
65+
}
66+
});
67+
68+
return {
69+
type: "object",
70+
properties,
71+
...(required.length > 0 ? { required } : {}),
72+
};
73+
} else if (schema instanceof z.ZodString) {
74+
return { type: "string" };
75+
} else if (schema instanceof z.ZodNumber) {
76+
return { type: "number" };
77+
} else if (schema instanceof z.ZodBoolean) {
78+
return { type: "boolean" };
79+
} else if (schema instanceof z.ZodArray) {
80+
return {
81+
type: "array",
82+
items: zodToJsonSchema(schema._def.type),
83+
};
84+
} else if (schema instanceof z.ZodOptional) {
85+
return zodToJsonSchema(schema._def.innerType);
86+
} else {
87+
return { type: "string" }; // Default fallback
88+
}
89+
}
90+
91+
// Define the available functions
92+
const availableFunctions = [
93+
{
94+
schema: CreateTodoSchema,
95+
name: "createTodo",
96+
description: CreateTodoSchema.description || "Creates a new todo item"
97+
},
98+
{
99+
schema: ExecuteTodoSchema,
100+
name: "executeTodoWorkflow",
101+
description: ExecuteTodoSchema.description || "Executes a todo item"
102+
}
103+
];
104+
105+
// Convert to OpenAI function format
106+
const functionsForOpenAI = availableFunctions.map(fn => ({
107+
name: fn.name,
108+
description: fn.description,
109+
parameters: zodToJsonSchema(fn.schema),
110+
}));
111+
112+
// Base system message
113+
const baseSystemMessage = `You are a helpful assistant that can create and execute todos.
114+
When a user asks to "do something" or "complete a task", you should:
115+
1. First create the todo using createTodo
116+
2. Then execute it using executeTodoWorkflow with the todoId from the previous step`;
56117

57118
export const llmChat = async ({
58119
systemContent = "",
@@ -62,52 +123,57 @@ export const llmChat = async ({
62123
try {
63124
const openai = openaiClient({});
64125

65-
// Combine system prompts if provided
126+
// Combine system messages
66127
const finalSystemContent = systemContent
67-
? `${structuredOutputPrompt}\n\n${systemContent}`
68-
: structuredOutputPrompt;
128+
? `${baseSystemMessage}\n\n${systemContent}`
129+
: baseSystemMessage;
69130

70-
// Set up response format for JSON
71-
const responseFormat = {
72-
type: "json_object" as const
73-
};
74-
75-
// Chat parameters
131+
// Chat parameters with tools
76132
const chatParams: ChatCompletionCreateParamsNonStreaming = {
77133
messages: [{ role: "system", content: finalSystemContent }, ...messages],
78134
model,
79-
response_format: responseFormat,
135+
tools: functionsForOpenAI.map(fn => ({ type: "function", function: fn })),
136+
tool_choice: "auto",
80137
};
81138

82139
log.debug("OpenAI chat completion params", { chatParams });
83140

84141
const completion = await openai.chat.completions.create(chatParams);
85142
const message = completion.choices[0].message;
86-
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
143+
144+
// Ensure we have a string content or default to empty string
145+
const messageContent = message.content || "";
146+
147+
// Parse function call if available
148+
let structuredData: StructuredResponse | undefined;
149+
if (message.tool_calls && message.tool_calls.length > 0) {
150+
const toolCall = message.tool_calls[0];
151+
if (toolCall.type === "function") {
152+
try {
153+
const functionArguments = JSON.parse(toolCall.function.arguments);
154+
structuredData = {
155+
type: "function_call" as const,
156+
function_name: toolCall.function.name,
157+
function_arguments: functionArguments,
158+
};
159+
log.debug("Function call detected", { structuredData });
160+
} catch (error) {
161+
log.error("Failed to parse function arguments", {
162+
arguments: toolCall.function.arguments
101163
});
102164
}
103-
} catch (error) {
104-
log.error("Failed to parse JSON response", { content: message.content });
105165
}
166+
} else if (messageContent) {
167+
// Handle regular text response
168+
structuredData = {
169+
type: "text" as const,
170+
content: messageContent,
171+
};
106172
}
107173

108174
return {
109175
role: "assistant",
110-
content: message.content,
176+
content: messageContent,
111177
structured_data: structuredData,
112178
};
113179
} catch (error) {

agent-todo/src/functions/toolTypes.ts

+19-6
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,27 @@
11
import { z } from "zod";
22

3-
export const CreateTodoInput = z.object({
4-
todoTitle: z.string(),
5-
});
3+
// Todo creation schemas and types
4+
export const CreateTodoSchema = z.object({
5+
todoTitle: z.string().min(1).describe("The title of the todo to create"),
6+
todoDescription: z.string().optional().describe("Optional detailed description of the todo"),
7+
}).describe("Creates a new todo but doesn't execute it. Use when a user wants to add a task to their list.");
8+
9+
export type CreateTodoInputType = z.infer<typeof CreateTodoSchema>;
10+
export type CreateTodoOutput = string;
11+
12+
// Todo execution schemas and types
13+
export const ExecuteTodoSchema = z.object({
14+
todoId: z.string().min(1).describe("The ID of the todo to execute"),
15+
todoTitle: z.string().optional().describe("The title of the todo to execute (if ID not available)"),
16+
}).describe("Executes a specific todo from the list. Use when a user wants to complete or perform a task.");
617

7-
export type CreateTodoInputType = z.infer<typeof CreateTodoInput>;
18+
export type ExecuteTodoInputType = z.infer<typeof ExecuteTodoSchema>;
819

9-
export const ExecuteTodoInput = z.object({
20+
export const ExecuteTodoOutputSchema = z.object({
1021
todoId: z.string(),
1122
todoTitle: z.string(),
23+
details: z.string(),
24+
status: z.string(),
1225
});
1326

14-
export type ExecuteTodoInputType = z.infer<typeof ExecuteTodoInput>;
27+
export type ExecuteTodoOutput = z.infer<typeof ExecuteTodoOutputSchema>;

agent-todo/src/workflows/executeTodo.ts

+4-16
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,9 @@
11
import { log, sleep, step } from "@restackio/ai/workflow";
22
import * as functions from "../functions";
3-
import { z } from "zod";
3+
import { ExecuteTodoSchema, ExecuteTodoOutput, ExecuteTodoOutputSchema } from "../functions/toolTypes";
44

5-
export const ExecuteTodoSchema = z.object({
6-
todoTitle: z.string().min(1),
7-
todoId: z.string().min(1),
8-
});
9-
10-
export const ExecuteTodoOutputSchema = z.object({
11-
todoId: z.string(),
12-
todoTitle: z.string(),
13-
details: z.string(),
14-
status: z.string(),
15-
});
16-
17-
export type Input = z.infer<typeof ExecuteTodoSchema>;
18-
export type Output = z.infer<typeof ExecuteTodoOutputSchema>;
5+
export type Input = ExecuteTodoOutput;
6+
export type Output = ExecuteTodoOutput;
197

208
export async function executeTodoWorkflow({ todoTitle, todoId }: Input): Promise<Output> {
219
try {
@@ -61,4 +49,4 @@ export async function executeTodoWorkflow({ todoTitle, todoId }: Input): Promise
6149
status: "failed"
6250
};
6351
}
64-
}
52+
}

0 commit comments

Comments
 (0)