Skip to content

Commit

Permalink
Add toolbox for tool augmented conversations.
Browse files Browse the repository at this point in the history
  • Loading branch information
ioquatix committed Nov 26, 2024
1 parent 67c4425 commit 59ebaf3
Show file tree
Hide file tree
Showing 4 changed files with 194 additions and 0 deletions.
2 changes: 2 additions & 0 deletions lib/async/ollama.rb
Original file line number Diff line number Diff line change
Expand Up @@ -5,3 +5,5 @@

require_relative "ollama/version"
require_relative "ollama/client"

require_relative "ollama/conversation"
62 changes: 62 additions & 0 deletions lib/async/ollama/conversation.rb
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
# frozen_string_literal: true

# Released under the MIT License.
# Copyright, 2024, by Samuel Williams.

require_relative "generate"
require_relative "toolbox"

module Async
module Ollama
class Conversation
def initialize(client, model: "llama3", context: nil)
@client = client

@toolbox = Toolbox.new

@state = Generate.new(client.with(path: "/api/generate"), value: {
model: model,
context: context,
})
end

attr :toolbox

def context
@state.context
end

def call(prompt)
@state = @state.generate(prompt)

while true
response = @state.response
if message = tool?(response)
result = @toolbox.call(message)
@state = @state.generate(result)
else
return response
end
end
end

def start(prompt)
call(prompt + "\n\n" + @toolbox.explain)
end

private

def tool?(response)
if response.start_with?('{')
begin
return JSON.parse(response, symbolize_names: true)
rescue => error
Console.debug(self, error)
end
end

return false
end
end
end
end
89 changes: 89 additions & 0 deletions lib/async/ollama/toolbox.rb
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
# frozen_string_literal: true

# Released under the MIT License.
# Copyright, 2024, by Samuel Williams.

module Async
module Ollama
class Tool
def initialize(name, explain, &block)
@name = name
@explain = explain
@block = block
end

attr :name
attr :explain

def call(message)
@block.call(message)
end
end

class Toolbox
PROMPT = "You have access to the following tools, which you can invoke by replying with a single line of valid JSON:\n\n"

USAGE = <<~EOF
- You are an AI assistant with access to specific tools. Use these tools to enhance your ability to answer user queries accurately.
- When you need to use a tool to answer the user's query, respond **only** with the JSON invocation.
- Example: {"tool":"ruby", "code": "5+5"}
- **Do not** include any explanations, greetings, or additional text when invoking a tool.
- If you are dealing with numbers, ensure you provide them as Integers or Floats, not Strings.
- After invoking a tool:
1. You will receive the tool's result as the next input.
2. Use the result to formulate a direct, user-friendly response that answers the original query.
3. Assume the user is unaware of the tool invocation or its result, so clearly summarize the answer without referring to the tool usage or the response it generated.
- Continue the conversation naturally after providing the answer. Ensure your responses are concise and user-focused.
## Example Flow:
- User: "Why doesn't 5 + 5 equal 11?"
- Assistant (invokes tool): {"tool": "ruby", "code": "5+5"}
- (Tool Result): 10
- Assistant: "The result of 5 + 5 is 10, because addition follows standard arithmetic rules."
Remember, invoke tools silently and use their results seamlessly in your responses to the user.
EOF

def initialize
@tools = {}
end

attr :tools

def register(name, explain, &block)
@tools[name] = Tool.new(name, explain, &block)
end

def call(message)
name = message[:tool]

if tool = @tools[name]
result = tool.call(message)

return {result: result}.to_json
else
raise ArgumentError.new("Unknown tool: #{name}")
end
rescue => error
{error: error.message}.to_json
end

def explain
buffer = String.new
buffer << PROMPT

@tools.each do |name, tool|
buffer << tool.explain
end

buffer << "\n" << USAGE << "\n"

return buffer
end
end
end
end
41 changes: 41 additions & 0 deletions test/async/ollama/conversation.rb
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
# frozen_string_literal: true

# Released under the MIT License.
# Copyright, 2024, by Samuel Williams.

require "async/ollama"
require "sus/fixtures/async/reactor_context"

describe Async::Ollama::Client do
include Sus::Fixtures::Async::ReactorContext

attr :client

before do
@client = Async::Ollama::Client.open
end

after do
@client.close
end

let(:conversation) {Async::Ollama::Conversation.new(@client)}

it "can invoke a tool" do
invoked = false
conversation.toolbox.register("sum", '{"tool":"sum", "arguments":[...]}') do |message|
invoked = message

message[:arguments].sum
end

response = conversation.start("You are a mathematician.")
inform response

response = conversation.call("Can you add the first 5 prime numbers?")
inform response

expect(invoked).to be == {tool: "sum", arguments: [2, 3, 5, 7, 11]}
expect(response).to be =~ /28/
end
end

0 comments on commit 59ebaf3

Please sign in to comment.