iachat/spec/enterprise/services/captain/tool_registry_service_spec.rb
Pranav d657f35a76
feat: Introduce the concept of tool registry within Captain (#11516)
This PR introduces the concept of a tool registry. The implementation is
straightforward: you can define a tool by creating a class with a
function name. The function name gets registered in the registry and can
be referenced during LLM calls. When the LLM invokes a tool using the
registered name, the registry locates and executes the appropriate tool.
If the LLM calls an unregistered tool, the registry returns an error
indicating that the tool is not defined.
2025-05-19 15:26:38 -07:00

88 lines
1.9 KiB
Ruby

require 'rails_helper'
# Test tool implementation
class TestTool < Captain::Tools::BaseService
def name
'test_tool'
end
def description
'A test tool for specs'
end
def parameters
{
type: 'object',
properties: {
test_param: {
type: 'string'
}
}
}
end
def execute(*args)
args
end
end
RSpec.describe Captain::ToolRegistryService do
let(:assistant) { create(:captain_assistant) }
let(:service) { described_class.new(assistant) }
describe '#initialize' do
it 'initializes with empty tools and registered_tools' do
expect(service.tools).to be_empty
expect(service.registered_tools).to be_empty
end
end
describe '#register_tool' do
let(:tool_class) { TestTool }
it 'registers a new tool' do
service.register_tool(tool_class)
expect(service.tools['test_tool']).to be_a(TestTool)
expect(service.registered_tools).to include(
{
type: 'function',
function: {
name: 'test_tool',
description: 'A test tool for specs',
parameters: {
type: 'object',
properties: {
test_param: {
type: 'string'
}
}
}
}
}
)
end
end
describe 'method_missing' do
let(:tool_class) { TestTool }
before do
service.register_tool(tool_class)
end
context 'when method corresponds to a registered tool' do
it 'executes the tool with given arguments' do
result = service.test_tool(test_param: 'arg1')
expect(result).to eq([{ test_param: 'arg1' }])
end
end
context 'when method does not correspond to a registered tool' do
it 'raises NoMethodError' do
expect { service.unknown_tool }.to raise_error(NoMethodError)
end
end
end
end