🐝 Bee Framework Agent

1import { Actor, log } from 'apify';
2import { BeeAgent } from 'bee-agent-framework/agents/bee/agent';
3import { UnconstrainedMemory } from 'bee-agent-framework/memory/unconstrainedMemory';
4import { z } from 'zod';
5import { LangChainChatModel } from 'bee-agent-framework/adapters/langchain/backend/chat';
6import { ChatOpenAI } from '@langchain/openai';
7import { OpenAIChatModel } from 'bee-agent-framework/adapters/openai/backend/chat';
8import { CalculatorSumTool } from './tools/calculator.js';
9import { InstagramScrapeTool } from './tools/instagram.js';
10import { StructuredOutputGenerator } from './structured_response_generator.js';
11import { beeOutputTotalTokens, chargeForActorStart, chargeForModelTokens } from './ppe_utils.js';
12
13// This is an ESM project, and as such, it requires you to specify extensions in your relative imports.
14// Read more about this here: https://nodejs.org/docs/latest-v18.x/api/esm.html#mandatory-file-extensions
15// Note that we need to use `.js` even when inside TS files
16// import { router } from './routes.js';
17
18// Actor input schema
19interface Input {
20    query: string;
21    modelName: string;
22    debug?: boolean;
23}
24
25// The init() call configures the Actor for its environment. It's recommended to start every Actor with an init().
26await Actor.init();
27
28// Handle input
29const {
30    // The query default value is provided only for template testing purposes.
31    // You can remove it.
32    query,
33    modelName,
34    debug,
35} = await Actor.getInput() as Input;
36if (debug) {
37    log.setLevel(log.LEVELS.DEBUG);
38}
39if (!query) {
40    throw new Error('An agent query is required.');
41}
42
43/**
44 * Actor code
45*/
46// Charge for Actor start
47await chargeForActorStart();
48
49// Create a ReAct agent that can use tools.
50// See https://i-am-bee.github.io/bee-agent-framework/#/agents?id=bee-agent
51// In order to use PPE, the LangChain adapter must be used
52// otherwise, the token usage is not tracked.
53log.debug(`Using model: ${modelName}`);
54const llm = new LangChainChatModel(
55    new ChatOpenAI({ model: modelName }),
56);
57// The LangChain adapter does not work with the structured output generation
58// for some reason.
59// Create a separate LLM for structured output generation.
60const llmStructured = new OpenAIChatModel(modelName);
61const agent = new BeeAgent({
62    llm,
63    memory: new UnconstrainedMemory(),
64    tools: [new CalculatorSumTool(),
65        new InstagramScrapeTool()],
66});
67
68// Store tool messages for later structured output generation.
69// This can be removed if you don't need structured output.
70const structuredOutputGenerator = new StructuredOutputGenerator(llmStructured);
71
72// Prompt the agent with the query.
73// Debug log agent status updates, e.g., thoughts, tool calls, etc.
74const response = await agent
75    .run({ prompt: query })
76    .observe((emitter) => {
77        emitter.on('update', async ({ update }) => {
78            log.debug(`Agent (${update.key}) 🤖 : ${update.value}`);
79
80            // Save tool messages for later structured output generation.
81            // This can be removed if you don't need structured output.
82            if (['tool_name', 'tool_output', 'tool_input'].includes(update.key as string)) {
83                structuredOutputGenerator.processToolMessage(
84                    update.key as 'tool_name' | 'tool_output' | 'tool_input',
85                    update.value,
86                );
87            }
88            // End of tool message saving.
89        });
90    });
91
92const tokensTotal = beeOutputTotalTokens(response);
93await chargeForModelTokens(modelName, tokensTotal);
94
95log.info(`Agent 🤖 : ${response.result.text}`);
96
97// Hacky way to get the structured output.
98// Using the stored tool messages and the user query to create a structured output.
99const structuredResponse = await structuredOutputGenerator.generateStructuredOutput(query,
100    z.object({
101        totalLikes: z.number(),
102        totalComments: z.number(),
103        mostPopularPosts: z.array(z.object({
104            url: z.string(),
105            likes: z.number(),
106            comments: z.number(),
107            timestamp: z.string(),
108            caption: z.string().nullable().optional(),
109            alt: z.string().nullable().optional(),
110        })),
111    }));
112log.debug(`Structured response: ${JSON.stringify(structuredResponse)}`);
113// Since the token usage tracking does not work with the Bee LLM, we will
114// just charge the same amount of tokens as the total tokens used by the agent for the
115// structured output generation - which is mostly the tool calls passed to the structured output generator.
116await chargeForModelTokens(modelName, tokensTotal);
117// End of structured output generation.
118
119// Push results to the dataset.
120await Actor.pushData({
121    query,
122    response: response.result.text,
123    // This can be removed if you don't need structured output.
124    structuredResponse: structuredResponse.object,
125});
126log.info('Pushed the data into the dataset!');
127
128// Gracefully exit the Actor process. It's recommended to quit all Actors with an exit().
129await Actor.exit();