Skip to content

Commit

Permalink
feat(tools): add llm and calculator tools (#132)
Browse files Browse the repository at this point in the history
Signed-off-by: Lukáš Janeček <lukas.janecek1@ibm.com>
Co-authored-by: Lukáš Janeček <lukas.janecek1@ibm.com>
  • Loading branch information
xjacka and Lukáš Janeček authored Dec 17, 2024
1 parent e899308 commit 1797a53
Showing 4 changed files with 76 additions and 14 deletions.
23 changes: 11 additions & 12 deletions pnpm-lock.yaml

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

33 changes: 32 additions & 1 deletion src/runs/execution/tools/helpers.ts
Original file line number Diff line number Diff line change
@@ -33,13 +33,16 @@ import { LLMChatTemplates } from 'bee-agent-framework/adapters/shared/llmChatTem
import { DuckDuckGoSearchTool } from 'bee-agent-framework/tools/search/duckDuckGoSearch';
import { SearchToolOptions, SearchToolOutput } from 'bee-agent-framework/tools/search/base';
import { PromptTemplate } from 'bee-agent-framework/template';
import { CalculatorTool } from 'bee-agent-framework/tools/calculator';
import { LLMTool } from 'bee-agent-framework/tools/llm';

import { AgentContext } from '../execute.js';
import { getRunVectorStores } from '../helpers.js';
import { CodeInterpreterTool as CodeInterpreterUserTool } from '../../../tools/entities/tool/code-interpreter-tool.entity.js';
import { ApiTool as ApiCallUserTool } from '../../../tools/entities/tool/api-tool.entity.js';
import { createCodeLLM } from '../factory.js';
import { createChatLLM, createCodeLLM } from '../factory.js';
import { RedisCache } from '../cache.js';
import { getDefaultModel } from '../constants.js';

import { createPythonStorage } from './python-tool-storage.js';
import { FunctionTool, FunctionToolOutput } from './function.js';
@@ -128,6 +131,22 @@ export async function getTools(run: LoadedRun, context: AgentContext): Promise<F
);
if (wikipediaUsage) tools.push(new WikipediaSimilaritySearchTool());

const llmUsage = run.tools.find(
(tool): tool is SystemUsage => tool.type === ToolType.SYSTEM && tool.toolId === SystemTools.LLM
);
if (llmUsage)
tools.push(
new LLMTool({
llm: createChatLLM({ model: getDefaultModel() })
})
);

const calculatorUsage = run.tools.find(
(tool): tool is SystemUsage =>
tool.type === ToolType.SYSTEM && tool.toolId === SystemTools.CALCULATOR
);
if (calculatorUsage) tools.push(new CalculatorTool());

const weatherUsage = run.tools.find(
(tool): tool is SystemUsage =>
tool.type === ToolType.SYSTEM && tool.toolId === SystemTools.WEATHER
@@ -342,6 +361,16 @@ export async function createToolCall(
toolId: SystemTools.READ_FILE,
input: await tool.parse(input)
});
} else if (tool instanceof LLMTool) {
return new SystemCall({
toolId: SystemTools.LLM,
input: await tool.parse(input)
});
} else if (tool instanceof CalculatorTool) {
return new SystemCall({
toolId: SystemTools.CALCULATOR,
input: await tool.parse(input)
});
} else if (tool instanceof FunctionTool) {
return new FunctionCall({ name: tool.name, arguments: JSON.stringify(input) });
} else if (tool instanceof CustomTool || tool instanceof ApiCallTool) {
@@ -387,6 +416,8 @@ export async function finalizeToolCall(
toolCall.output = result;
break;
}
case SystemTools.CALCULATOR:
case SystemTools.LLM:
case SystemTools.READ_FILE: {
if (!(result instanceof StringToolOutput)) throw new TypeError();
toolCall.output = result.result;
4 changes: 3 additions & 1 deletion src/tools/entities/tool-calls/system-call.entity.ts
Original file line number Diff line number Diff line change
@@ -25,7 +25,9 @@ export enum SystemTools {
WIKIPEDIA = 'wikipedia',
WEATHER = 'weather',
ARXIV = 'arxiv',
READ_FILE = 'read_file'
READ_FILE = 'read_file',
LLM = 'llm',
CALCULATOR = 'calculator'
}

@Embeddable({ discriminatorValue: ToolType.SYSTEM })
30 changes: 30 additions & 0 deletions src/tools/tools.service.ts
Original file line number Diff line number Diff line change
@@ -19,6 +19,8 @@ import { CustomTool, CustomToolCreateError } from 'bee-agent-framework/tools/cus
import dayjs from 'dayjs';
import mime from 'mime/lite';
import { WikipediaTool } from 'bee-agent-framework/tools/search/wikipedia';
import { LLMTool } from 'bee-agent-framework/tools/llm';
import { CalculatorTool } from 'bee-agent-framework/tools/calculator';
import { Tool as FrameworkTool } from 'bee-agent-framework/tools/base';
import { ZodTypeAny } from 'zod';
import { zodToJsonSchema } from 'zod-to-json-schema';
@@ -72,6 +74,8 @@ import { createCodeInterpreterConnectionOptions } from '@/runs/execution/tools/h
import { ReadFileTool } from '@/runs/execution/tools/read-file-tool.js';
import { snakeToCamel } from '@/utils/strings.js';
import { createSearchTool } from '@/runs/execution/tools/search-tool';
import { createChatLLM } from '@/runs/execution/factory.js';
import { getDefaultModel } from '@/runs/execution/constants.js';

type SystemTool = Pick<FrameworkTool, 'description' | 'name' | 'inputSchema'> & {
type: ToolType;
@@ -462,6 +466,10 @@ function getSystemTools() {
});
const fileSearch = new FileSearchTool({ vectorStores: [], maxNumResults: 0 });
const readFile = new ReadFileTool({ files: [], fileSize: 0 });
const llmTool = new LLMTool({
llm: createChatLLM({ model: getDefaultModel() })
});
const calculatorTool = new CalculatorTool();

const systemTools = new Map<string, SystemTool>();

@@ -569,6 +577,26 @@ function getSystemTools() {
userDescription:
'Execute Python code for various tasks, including data analysis, file processing, and visualizations. Supports the installation of any library such as NumPy, Pandas, SciPy, and Matplotlib. Users can create new files or convert existing files, which are then made available for download.'
});
systemTools.set(SystemTools.LLM, {
type: ToolType.SYSTEM,
id: SystemTools.LLM,
createdAt: new Date('2024-12-12'),
...llmTool,
inputSchema: llmTool.inputSchema.bind(llmTool),
isExternal: false,
userDescription:
'Uses expert LLM to work with data in the existing conversation (classification, entity extraction, summarization, ...)'
});
systemTools.set(SystemTools.CALCULATOR, {
type: ToolType.SYSTEM,
id: SystemTools.CALCULATOR,
createdAt: new Date('2024-12-12'),
...calculatorTool,
inputSchema: calculatorTool.inputSchema.bind(calculatorTool),
isExternal: false,
userDescription:
'A calculator tool that performs basic arithmetic operations like addition, subtraction, multiplication, and division. Only use the calculator tool if you need to perform a calculation.'
});

return systemTools;
}
@@ -591,6 +619,8 @@ export async function listTools({
allSystemTools.get(SystemTools.WIKIPEDIA),
allSystemTools.get(SystemTools.WEATHER),
allSystemTools.get(SystemTools.ARXIV),
allSystemTools.get(SystemTools.LLM),
allSystemTools.get(SystemTools.CALCULATOR),
allSystemTools.get('read_file')
]
: [];

0 comments on commit 1797a53

Please sign in to comment.