I’m using Vercel’s AI SDK in my agent.ts
file to handle chat interactions and execute specific tools based on user inputs. The file streams AI-generated responses and triggers tool calls like create_streamlit_app
for generating Streamlit code. I need guidance on which functions or methods within the package can be used to programmatically initiate tool calls and manage message processing effectively. Any insights or examples would be greatly appreciated.
agent.ts
import { streamObject, CoreUserMessage, ToolCall, ToolResult } from 'ai'
import { generateCode } from './tools'
import { CSVAnalysis, StreamChunk } from './types'
// ... existing code ...
export class GruntyAgent {
// ... existing properties and constructor ...
async *chat(
chatId: string,
userId: string,
latestMessage: string,
tools: Tool[],
temperature: number,
maxTokens: number,
csvAnalysis?: CSVAnalysis
): AsyncGenerator<StreamChunk> {
// ... existing code ...
const streamResult = await streamObject({
model: modelClient as LanguageModelV1,
output: 'array',
system: this.roleDescription,
schema: messageSchema,
messages: sanitizedMessages,
maxTokens: maxTokens,
onFinish: async () => {
await streamResult.usage.then(recordTokenUsage);
await this.storeMessage(
chatId,
userId,
latestMessage,
accumulatedResponse.trim(),
this.totalTokens,
toolCalls,
toolResults
)
}
})
for await (const events of streamResult.partialObjectStream) {
for (const event of events) {
if (event.type === 'text') {
accumulatedResponse += event.content;
yield event as StreamChunk;
if (event.content.includes('create a Streamlit app')) {
console.log('Keyword detected for Streamlit app creation.');
const toolCallEvent = {
toolCallId: 'create-streamlit-app',
toolName: 'create_streamlit_app',
args: { query: 'Analyze CSV and create Streamlit app' }
} as ToolCall<string, any>;
toolCalls.push(toolCallEvent);
console.log('Tool Call Manually Triggered:', toolCallEvent);
try {
const codeQuery = `
${toolCallEvent.args.query}
Use the following CSV analysis to inform your code:
${JSON.stringify(csvAnalysis, null, 2)}
`;
const { generatedCode, codeTokenCount } = await generateCode(codeQuery);
this.codeTokens += codeTokenCount;
const toolResultEvent = {
toolCallId: toolCallEvent.toolCallId,
toolName: toolCallEvent.toolName,
result: generatedCode,
args: toolCallEvent.args
} as ToolResult<string, any, any>;
toolResults.push(toolResultEvent);
console.log('Tool Result Manually Triggered:', toolResultEvent);
} catch (error) {
console.error('Error generating code:', error);
}
}
}
// ... existing code ...
}
}
}
// ... existing methods ...
}
tools.ts
import { Anthropic } from '@anthropic-ai/sdk'
import { z } from 'zod'
import { Tool } from './types'
const codeGenerationAnthropicAgent = new Anthropic({
apiKey: process.env.ANTHROPIC_API_KEY,
})
const streamlitAppSchema = z.object({
query: z.string().min(1, 'Query cannot be empty').describe(
'Explain the requirements for the Streamlit code you want to generate.'
),
})
export const tools: Tool[] = [
{
name: 'create_streamlit_app',
description: 'Generates Python (Streamlit) code based on a given query',
inputSchema: streamlitAppSchema,
parameters: streamlitAppSchema,
execute: async (input) => {
const { query } = streamlitAppSchema.parse(input)
return generateCode(query)
},
},
]
export async function generateCode(
query: string
): Promise<{ generatedCode: string; codeTokenCount: number }> {
if (!query || !query.trim()) {
throw new Error('Query cannot be empty or just whitespace.')
}
console.log('Sending query to LLM:', query)
try {
const response = await codeGenerationAnthropicAgent.messages.create({
model: 'claude-3-5-sonnet-20240620',
max_tokens: 2000,
system: 'You are a Python code generation assistant specializing in Streamlit apps.',
messages: [{ role: 'user', content: query }],
})
if (Array.isArray(response.content) && response.content.length > 0) {
const generatedCode =
response.content[0].type === 'text'
? response.content[0].text.replace(/^```python/, '').replace(/```$/, '')
: ''
return {
generatedCode,
codeTokenCount: response.usage.input_tokens + response.usage.output_tokens,
}
} else {
console.error('Unexpected response format:', response)
throw new Error('Unexpected response format from code generation API')
}
} catch (error) {
console.error('Error generating code:', error)
throw new Error('Failed to generate code. Please check the query and try again.')
}
}
These snippets illustrate how the agent.ts
file handles chat interactions and triggers tool calls, while tools.ts defines the tool and its execution logic.