brute 3.2.1 → 4.0.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- checksums.yaml +4 -4
- data/lib/brute/hooks.rb +77 -0
- data/lib/brute/message_transport/open_router.rb +77 -3
- data/lib/brute/middleware/025_skills.rb +91 -0
- data/lib/brute/middleware/070_tool_pipeline.rb +110 -13
- data/lib/brute/middleware/open_router.rb +79 -4
- data/lib/brute/prompt_template.rb +125 -0
- data/lib/brute/prompts/base.rb +105 -0
- data/lib/brute/prompts/skills.rb +51 -13
- data/lib/brute/prompts/text/skills/default.erb +14 -0
- data/lib/brute/skill.rb +211 -80
- data/lib/brute/turn/agent_pipeline.rb +27 -1
- data/lib/brute/turn/pipeline.rb +15 -0
- data/lib/brute/version.rb +1 -1
- data/lib/brute/version.rb.erb +5 -0
- metadata +6 -1
checksums.yaml
CHANGED
|
@@ -1,7 +1,7 @@
|
|
|
1
1
|
---
|
|
2
2
|
SHA256:
|
|
3
|
-
metadata.gz:
|
|
4
|
-
data.tar.gz:
|
|
3
|
+
metadata.gz: 65f8fa383e8c2d5793cc4ec00023ee6d2117e75535a81676617232b94748f154
|
|
4
|
+
data.tar.gz: 88f3794d39b79c3c7c03c96f54383fa8647850be1cfc77ae8877fb6ea42aa365
|
|
5
5
|
SHA512:
|
|
6
|
-
metadata.gz:
|
|
7
|
-
data.tar.gz:
|
|
6
|
+
metadata.gz: 57b849d89d406b624fe26dcbc93aa84af2f9b90b12160fb57608938ec53fd7109c87aa3d2c1357620c78f88b9702f797ff106432840a5e4ccaca343f49e7035a
|
|
7
|
+
data.tar.gz: d73709348a3f205f034389e86950a5ce171d0fd0bf62a9c9e84ad587d3ac322aad211b94d059a93b9b19e16a0c1b6211baebd7f274b88435c1cf61a9a6373007
|
data/lib/brute/hooks.rb
ADDED
|
@@ -0,0 +1,77 @@
|
|
|
1
|
+
# frozen_string_literal: true
|
|
2
|
+
|
|
3
|
+
require "bundler/setup"
|
|
4
|
+
require "brute"
|
|
5
|
+
|
|
6
|
+
module Brute
|
|
7
|
+
# Pub/sub registry for agent lifecycle hooks, subscribed on the builder:
|
|
8
|
+
#
|
|
9
|
+
# Brute.agent
|
|
10
|
+
# .use(Brute::Middleware::MaxIterations)
|
|
11
|
+
# .run(->(env) { env[:messages].assistant("done") })
|
|
12
|
+
# .on(:before_llm) { |env| ... }
|
|
13
|
+
# .on(:approve_tool) { |call| call[:name] != "exec" }
|
|
14
|
+
#
|
|
15
|
+
# Emission points and payloads:
|
|
16
|
+
#
|
|
17
|
+
# :turn_start, :turn_end → the turn env (AgentPipeline#start; turn_end
|
|
18
|
+
# fires from an ensure, so it also fires on error)
|
|
19
|
+
# :before_llm, :after_llm → the turn env, around every LLM call
|
|
20
|
+
# :before_tool → call env {name:, arguments:, result:, events:,
|
|
21
|
+
# metadata:, turn_env:} — mutate :arguments to
|
|
22
|
+
# rewrite the call, or set :result (or return a
|
|
23
|
+
# value) to skip execution entirely ("respond")
|
|
24
|
+
# :approve_tool → call env — a false return denies the call; a
|
|
25
|
+
# String return denies it with that message
|
|
26
|
+
# :after_tool → call env — mutate :result
|
|
27
|
+
#
|
|
28
|
+
# Subscribers run inline (tool events may fire from parallel threads).
|
|
29
|
+
# Exceptions propagate to the caller — layers that want fail-open semantics
|
|
30
|
+
# rescue in their own subscriber.
|
|
31
|
+
class Hooks
|
|
32
|
+
EVENTS = %i[turn_start turn_end before_llm after_llm before_tool approve_tool after_tool].freeze
|
|
33
|
+
|
|
34
|
+
def initialize
|
|
35
|
+
@subscribers = Hash.new { |hash, key| hash[key] = [] }
|
|
36
|
+
end
|
|
37
|
+
|
|
38
|
+
def on(event, &block)
|
|
39
|
+
@subscribers[event.to_sym] << block
|
|
40
|
+
self
|
|
41
|
+
end
|
|
42
|
+
|
|
43
|
+
# Fire an event; returns every subscriber's raw result (nils and false
|
|
44
|
+
# included — the deny contract distinguishes them).
|
|
45
|
+
def emit(event, payload)
|
|
46
|
+
@subscribers[event.to_sym].map { |subscriber| subscriber.call(payload) }
|
|
47
|
+
end
|
|
48
|
+
|
|
49
|
+
def any?(event) = @subscribers[event.to_sym].any?
|
|
50
|
+
end
|
|
51
|
+
end
|
|
52
|
+
|
|
53
|
+
__END__
|
|
54
|
+
|
|
55
|
+
describe "brute/hooks" do
|
|
56
|
+
it "emits to subscribers in registration order" do
|
|
57
|
+
hooks = Brute::Hooks.new
|
|
58
|
+
seen = []
|
|
59
|
+
hooks.on(:before_llm) { |p| seen << "a#{p}" }
|
|
60
|
+
hooks.on(:before_llm) { |p| seen << "b#{p}" }
|
|
61
|
+
hooks.emit(:before_llm, 1)
|
|
62
|
+
seen.should == ["a1", "b1"]
|
|
63
|
+
end
|
|
64
|
+
|
|
65
|
+
it "returns raw results, false included (deny contract)" do
|
|
66
|
+
hooks = Brute::Hooks.new
|
|
67
|
+
hooks.on(:approve_tool) { |_call| false }
|
|
68
|
+
hooks.emit(:approve_tool, {}).should == [false]
|
|
69
|
+
end
|
|
70
|
+
|
|
71
|
+
it "answers any? and stays chainable" do
|
|
72
|
+
hooks = Brute::Hooks.new
|
|
73
|
+
hooks.any?(:turn_start).should.be.false
|
|
74
|
+
hooks.on(:turn_start) { nil }.should.equal?(hooks)
|
|
75
|
+
hooks.any?(:turn_start).should.be.true
|
|
76
|
+
end
|
|
77
|
+
end
|
|
@@ -3,28 +3,102 @@
|
|
|
3
3
|
require "bundler/setup"
|
|
4
4
|
require "brute"
|
|
5
5
|
require "brute/message_transport"
|
|
6
|
+
require "json"
|
|
6
7
|
|
|
7
8
|
module Brute
|
|
8
9
|
class MessageTransport
|
|
10
|
+
# MessageTransport for the open_router_enhanced gem
|
|
11
|
+
# (https://github.com/estiens/open_router_enhanced). Brute does not
|
|
12
|
+
# require it — you do:
|
|
13
|
+
#
|
|
14
|
+
# require "open_router"
|
|
9
15
|
class OpenRouter < MessageTransport
|
|
10
16
|
def self.dump(message)
|
|
11
17
|
message.to_h
|
|
12
18
|
end
|
|
13
19
|
|
|
20
|
+
# An OpenRouter::Response's messages (one per choice; in practice
|
|
21
|
+
# OpenRouter returns exactly one).
|
|
22
|
+
def messages
|
|
23
|
+
return @result.choices.map { |choice| choice["message"] || choice[:message] } if @result.respond_to?(:choices)
|
|
24
|
+
|
|
25
|
+
super
|
|
26
|
+
end
|
|
27
|
+
|
|
14
28
|
private
|
|
15
29
|
|
|
16
30
|
def wrap(message)
|
|
17
31
|
# Coerce string keys to symbol keys if necessary
|
|
18
32
|
hash = message.to_h.transform_keys(&:to_sym)
|
|
19
|
-
|
|
33
|
+
hash[:role] = hash[:role].to_sym if hash.key?(:role)
|
|
34
|
+
hash[:tool_calls] = hash[:tool_calls].map { |tc| wrap_tool_call(tc) } if hash[:tool_calls]
|
|
35
|
+
|
|
20
36
|
case hash
|
|
21
37
|
in { role: (:system | :user | :assistant | :tool) }
|
|
22
|
-
#
|
|
23
|
-
Brute::Message.
|
|
38
|
+
# Slice away provider extras (refusal, reasoning, model, ...)
|
|
39
|
+
# that Brute::Message doesn't know.
|
|
40
|
+
Brute::Message.new(**hash.slice(:role, :content, :tool_calls, :tool_call_id))
|
|
24
41
|
else
|
|
25
42
|
raise "Unrecognised message format #{message.inspect}"
|
|
26
43
|
end
|
|
27
44
|
end
|
|
45
|
+
|
|
46
|
+
# An OpenAI-wire tool call ({ id:, type:, function: { name:, arguments: JSON } })
|
|
47
|
+
# -> the flat { id:, name:, arguments: Hash } Brute::Message understands.
|
|
48
|
+
def wrap_tool_call(tool_call)
|
|
49
|
+
tc = tool_call.to_h.transform_keys(&:to_sym)
|
|
50
|
+
return tc unless tc[:function] # already flat { id:, name:, arguments: }
|
|
51
|
+
|
|
52
|
+
function = tc[:function].to_h.transform_keys(&:to_sym)
|
|
53
|
+
arguments = function[:arguments].to_s
|
|
54
|
+
|
|
55
|
+
{
|
|
56
|
+
id: tc[:id],
|
|
57
|
+
name: function[:name],
|
|
58
|
+
arguments: JSON.parse(arguments.empty? ? "{}" : arguments),
|
|
59
|
+
}
|
|
60
|
+
end
|
|
28
61
|
end
|
|
29
62
|
end
|
|
30
63
|
end
|
|
64
|
+
|
|
65
|
+
__END__
|
|
66
|
+
|
|
67
|
+
describe "brute/message_transport/open_router" do
|
|
68
|
+
require "brute/messages"
|
|
69
|
+
|
|
70
|
+
it "dumps a message to the wire format" do
|
|
71
|
+
m = Brute::Message.new(role: :user, content: "hi")
|
|
72
|
+
Brute::MessageTransport::OpenRouter.dump(m).should == { role: :user, content: "hi" }
|
|
73
|
+
end
|
|
74
|
+
|
|
75
|
+
it "wraps a response's choice messages with symbolised roles, dropping provider extras" do
|
|
76
|
+
fake_response = Struct.new(:choices).new([
|
|
77
|
+
{ "message" => { "role" => "assistant", "content" => "hi there",
|
|
78
|
+
"refusal" => nil, "reasoning" => nil, "model" => "openrouter/auto" } },
|
|
79
|
+
])
|
|
80
|
+
|
|
81
|
+
out = Brute::MessageTransport::OpenRouter.new(fake_response).wrap_each.to_a
|
|
82
|
+
|
|
83
|
+
out.size.should == 1
|
|
84
|
+
out.first.role.should == :assistant
|
|
85
|
+
out.first.content.should == "hi there"
|
|
86
|
+
end
|
|
87
|
+
|
|
88
|
+
it "unwraps OpenAI-wire tool calls and parses their JSON arguments" do
|
|
89
|
+
fake_response = Struct.new(:choices).new([
|
|
90
|
+
{ "message" => {
|
|
91
|
+
"role" => "assistant", "content" => nil,
|
|
92
|
+
"tool_calls" => [{ "id" => "tc1", "type" => "function",
|
|
93
|
+
"function" => { "name" => "shell", "arguments" => '{"command":"ls"}' } }],
|
|
94
|
+
} },
|
|
95
|
+
])
|
|
96
|
+
|
|
97
|
+
out = Brute::MessageTransport::OpenRouter.new(fake_response).wrap_each.to_a.first
|
|
98
|
+
|
|
99
|
+
out.tool_call?.should.be.true
|
|
100
|
+
out.tool_calls.first.id.should == "tc1"
|
|
101
|
+
out.tool_calls.first.name.should == "shell"
|
|
102
|
+
out.tool_calls.first.arguments.should == { "command" => "ls" }
|
|
103
|
+
end
|
|
104
|
+
end
|
|
@@ -0,0 +1,91 @@
|
|
|
1
|
+
# frozen_string_literal: true
|
|
2
|
+
|
|
3
|
+
require "bundler/setup"
|
|
4
|
+
require "brute"
|
|
5
|
+
|
|
6
|
+
module Brute
|
|
7
|
+
module Middleware
|
|
8
|
+
# Loads skill objects into the agent context.
|
|
9
|
+
#
|
|
10
|
+
# Skills are handed in as objects — discovery is the caller's job:
|
|
11
|
+
#
|
|
12
|
+
# skills = Brute::Skill.all(cwd: Dir.pwd)
|
|
13
|
+
# agent
|
|
14
|
+
# .use(Brute::Middleware::Skills, skills: skills)
|
|
15
|
+
# .use(Brute::Middleware::SystemPrompt)
|
|
16
|
+
#
|
|
17
|
+
# Per turn:
|
|
18
|
+
# 1. env[:skills] = the objects, for downstream middleware, tools, and
|
|
19
|
+
# the terminal app (prime-agent's resourceLoader.getSkills() analogue)
|
|
20
|
+
# 2. env[:metadata][:skills] = the same objects, so
|
|
21
|
+
# Middleware::SystemPrompt merges them into the prompt ctx and
|
|
22
|
+
# Brute::Prompts::Skills renders the <available_skills> section
|
|
23
|
+
#
|
|
24
|
+
# Place it before Middleware::SystemPrompt in the stack. It never touches
|
|
25
|
+
# env[:messages] itself.
|
|
26
|
+
class Skills
|
|
27
|
+
def initialize(app, skills: [])
|
|
28
|
+
@app = app
|
|
29
|
+
@skills = skills
|
|
30
|
+
end
|
|
31
|
+
|
|
32
|
+
def call(env)
|
|
33
|
+
env[:skills] = @skills
|
|
34
|
+
env[:metadata] ||= {}
|
|
35
|
+
env[:metadata][:skills] ||= @skills
|
|
36
|
+
|
|
37
|
+
@app.call(env)
|
|
38
|
+
end
|
|
39
|
+
end
|
|
40
|
+
end
|
|
41
|
+
end
|
|
42
|
+
|
|
43
|
+
__END__
|
|
44
|
+
|
|
45
|
+
describe "brute/middleware/025_skills" do
|
|
46
|
+
def skill(name)
|
|
47
|
+
Brute::Skill.new(name: name, description: "x", file_path: "/x/#{name}/SKILL.md")
|
|
48
|
+
end
|
|
49
|
+
|
|
50
|
+
def build_middleware(skills: [], &inner)
|
|
51
|
+
Brute::Middleware::Skills.new(inner || ->(env) { env }, skills: skills)
|
|
52
|
+
end
|
|
53
|
+
|
|
54
|
+
it "stashes skill objects in env[:skills]" do
|
|
55
|
+
skills = [skill("debugging")]
|
|
56
|
+
env = { messages: Brute.log, metadata: {} }
|
|
57
|
+
|
|
58
|
+
build_middleware(skills: skills).call(env)
|
|
59
|
+
|
|
60
|
+
env[:skills].should == skills
|
|
61
|
+
end
|
|
62
|
+
|
|
63
|
+
it "mirrors skills into env[:metadata] for the prompt layer" do
|
|
64
|
+
skills = [skill("debugging")]
|
|
65
|
+
env = { messages: Brute.log }
|
|
66
|
+
|
|
67
|
+
build_middleware(skills: skills).call(env)
|
|
68
|
+
|
|
69
|
+
env[:metadata][:skills].should == skills
|
|
70
|
+
end
|
|
71
|
+
|
|
72
|
+
it "does not clobber an explicit metadata[:skills]" do
|
|
73
|
+
explicit = [skill("explicit")]
|
|
74
|
+
env = { messages: Brute.log, metadata: { skills: explicit } }
|
|
75
|
+
|
|
76
|
+
build_middleware(skills: [skill("other")]).call(env)
|
|
77
|
+
|
|
78
|
+
env[:metadata][:skills].should == explicit
|
|
79
|
+
env[:skills].map(&:name).should == ["other"]
|
|
80
|
+
end
|
|
81
|
+
|
|
82
|
+
it "defaults to an empty list and passes control down the chain" do
|
|
83
|
+
called = false
|
|
84
|
+
env = { messages: Brute.log }
|
|
85
|
+
|
|
86
|
+
build_middleware { |e| called = true }.call(env)
|
|
87
|
+
|
|
88
|
+
env[:skills].should == []
|
|
89
|
+
called.should.be.true
|
|
90
|
+
end
|
|
91
|
+
end
|
|
@@ -48,23 +48,56 @@ module Brute
|
|
|
48
48
|
name = tool_call.name.to_sym
|
|
49
49
|
args = tool_call.arguments
|
|
50
50
|
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
|
|
55
|
-
|
|
56
|
-
|
|
57
|
-
|
|
58
|
-
|
|
51
|
+
# Lifecycle hooks (Brute::Hooks): before_tool may rewrite
|
|
52
|
+
# :arguments or short-circuit with a :result; approve_tool
|
|
53
|
+
# denies on a false (or String) return; after_tool may
|
|
54
|
+
# rewrite :result.
|
|
55
|
+
call_env = {
|
|
56
|
+
name: name.to_s,
|
|
57
|
+
arguments: args,
|
|
58
|
+
result: nil,
|
|
59
|
+
events: env[:events],
|
|
60
|
+
metadata: {},
|
|
61
|
+
turn_env: env,
|
|
62
|
+
}
|
|
63
|
+
if (hooks = env[:hooks])
|
|
64
|
+
responses = hooks.emit(:before_tool, call_env).compact
|
|
65
|
+
call_env[:result] = responses.last if call_env[:result].nil? && !responses.empty?
|
|
66
|
+
|
|
67
|
+
if call_env[:result].nil?
|
|
68
|
+
denial = hooks.emit(:approve_tool, call_env).find { |r| r == false || r.is_a?(String) }
|
|
69
|
+
unless denial.nil?
|
|
70
|
+
call_env[:result] = denial.is_a?(String) ? denial : %(Tool call to "#{name}" was denied.)
|
|
71
|
+
end
|
|
59
72
|
end
|
|
73
|
+
end
|
|
60
74
|
|
|
61
|
-
|
|
62
|
-
|
|
63
|
-
|
|
64
|
-
|
|
75
|
+
result = if call_env[:result].nil?
|
|
76
|
+
available_tools[name].call(call_env[:arguments])
|
|
77
|
+
else
|
|
78
|
+
call_env[:result]
|
|
79
|
+
end
|
|
80
|
+
|
|
81
|
+
if (hooks = env[:hooks])
|
|
82
|
+
call_env[:result] = result
|
|
83
|
+
hooks.emit(:after_tool, call_env)
|
|
84
|
+
result = call_env[:result]
|
|
85
|
+
end
|
|
86
|
+
|
|
87
|
+
# Coerce to String so Hash results (e.g. Shell's
|
|
88
|
+
# {stdout:, stderr:, exit_code:}) serialize predictably.
|
|
89
|
+
if result.is_a?(String)
|
|
90
|
+
content = result
|
|
91
|
+
else
|
|
92
|
+
content = result.to_s
|
|
93
|
+
end
|
|
65
94
|
|
|
66
|
-
|
|
95
|
+
# Universal truncation safety net — skip if already truncated
|
|
96
|
+
unless Brute::Truncation.already_truncated?(content)
|
|
97
|
+
content = Brute::Truncation.truncate(content)
|
|
67
98
|
end
|
|
99
|
+
|
|
100
|
+
results << [tool_call, content]
|
|
68
101
|
rescue => e
|
|
69
102
|
# Capture the error as a tool result so the LLM can see it
|
|
70
103
|
# and reason about the failure, rather than crashing the
|
|
@@ -141,6 +174,70 @@ describe "brute/middleware/070_tool_pipeline" do
|
|
|
141
174
|
seen.should == [tool]
|
|
142
175
|
end
|
|
143
176
|
|
|
177
|
+
# --- lifecycle hooks (Brute::Hooks) ---
|
|
178
|
+
|
|
179
|
+
def hook_env(hooks)
|
|
180
|
+
{ messages: Brute.log, events: [], hooks: hooks }
|
|
181
|
+
end
|
|
182
|
+
|
|
183
|
+
it "before_tool may rewrite arguments and short-circuit with a result" do
|
|
184
|
+
tool = { name: "echo", description: "", execute: ->(text:) { "ran:#{text}" } }
|
|
185
|
+
inner = ->(env) do
|
|
186
|
+
env[:messages] << Brute::Message.new(role: :assistant, content: "",
|
|
187
|
+
tool_calls: [{ id: "tc1", name: "echo", arguments: { "text" => "orig" } }])
|
|
188
|
+
end
|
|
189
|
+
hooks = Brute::Hooks.new
|
|
190
|
+
hooks.on(:before_tool) { |call| call[:arguments] = { text: "rewritten" }; nil }
|
|
191
|
+
mw = Brute::Middleware::ToolPipeline.new(inner, tools: [tool])
|
|
192
|
+
env = hook_env(hooks)
|
|
193
|
+
env[:messages].user("hi")
|
|
194
|
+
mw.call(env)
|
|
195
|
+
env[:messages].last.content.should == "ran:rewritten"
|
|
196
|
+
|
|
197
|
+
hooks2 = Brute::Hooks.new
|
|
198
|
+
hooks2.on(:before_tool) { |_call| "canned" }
|
|
199
|
+
mw2 = Brute::Middleware::ToolPipeline.new(inner, tools: [tool])
|
|
200
|
+
env2 = hook_env(hooks2)
|
|
201
|
+
env2[:messages].user("hi")
|
|
202
|
+
mw2.call(env2)
|
|
203
|
+
env2[:messages].last.content.should == "canned" # never executed
|
|
204
|
+
end
|
|
205
|
+
|
|
206
|
+
it "approve_tool denies on false (generic message) or String (custom)" do
|
|
207
|
+
tool = { name: "exec", description: "", execute: ->(**) { "ran" } }
|
|
208
|
+
inner = ->(env) do
|
|
209
|
+
env[:messages] << Brute::Message.new(role: :assistant, content: "",
|
|
210
|
+
tool_calls: [{ id: "tc1", name: "exec", arguments: {} }])
|
|
211
|
+
end
|
|
212
|
+
hooks = Brute::Hooks.new
|
|
213
|
+
hooks.on(:approve_tool) { |_call| false }
|
|
214
|
+
env = hook_env(hooks)
|
|
215
|
+
env[:messages].user("hi")
|
|
216
|
+
Brute::Middleware::ToolPipeline.new(inner, tools: [tool]).call(env)
|
|
217
|
+
env[:messages].last.content.should == %(Tool call to "exec" was denied.)
|
|
218
|
+
|
|
219
|
+
hooks2 = Brute::Hooks.new
|
|
220
|
+
hooks2.on(:approve_tool) { |_call| "denied by policy" }
|
|
221
|
+
env2 = hook_env(hooks2)
|
|
222
|
+
env2[:messages].user("hi")
|
|
223
|
+
Brute::Middleware::ToolPipeline.new(inner, tools: [tool]).call(env2)
|
|
224
|
+
env2[:messages].last.content.should == "denied by policy"
|
|
225
|
+
end
|
|
226
|
+
|
|
227
|
+
it "after_tool may rewrite the result" do
|
|
228
|
+
tool = { name: "echo", description: "", execute: ->(**) { "raw" } }
|
|
229
|
+
inner = ->(env) do
|
|
230
|
+
env[:messages] << Brute::Message.new(role: :assistant, content: "",
|
|
231
|
+
tool_calls: [{ id: "tc1", name: "echo", arguments: {} }])
|
|
232
|
+
end
|
|
233
|
+
hooks = Brute::Hooks.new
|
|
234
|
+
hooks.on(:after_tool) { |call| call[:result] = "rewrote(#{call[:result]})" }
|
|
235
|
+
env = hook_env(hooks)
|
|
236
|
+
env[:messages].user("hi")
|
|
237
|
+
Brute::Middleware::ToolPipeline.new(inner, tools: [tool]).call(env)
|
|
238
|
+
env[:messages].last.content.should == "rewrote(raw)"
|
|
239
|
+
end
|
|
240
|
+
|
|
144
241
|
# --- Universal output truncation ---
|
|
145
242
|
|
|
146
243
|
it "truncates large tool results via Truncation" do
|
|
@@ -4,29 +4,104 @@ module Brute
|
|
|
4
4
|
module Middleware
|
|
5
5
|
module OpenRouter
|
|
6
6
|
class Completion
|
|
7
|
+
# config: keyword arguments for OpenRouter::Client.new
|
|
8
|
+
# (access_token:, request_timeout:, uri_base:, extra_headers:).
|
|
9
|
+
# Defaults to OpenRouter.configuration's global settings.
|
|
10
|
+
# options: keyword arguments for OpenRouter::CompletionOptions.new
|
|
11
|
+
# (model:, temperature:, tools:, ...).
|
|
7
12
|
def initialize(app, config: {}, **options)
|
|
8
13
|
@app = app
|
|
9
|
-
@config =
|
|
10
|
-
@options = ::OpenRouter::CompletionOptions.new(options)
|
|
14
|
+
@config = config
|
|
15
|
+
@options = ::OpenRouter::CompletionOptions.new(**options)
|
|
11
16
|
end
|
|
12
17
|
|
|
13
18
|
def call(env)
|
|
19
|
+
env[:hooks]&.emit(:before_llm, env)
|
|
20
|
+
|
|
14
21
|
messages = Brute::MessageTransport::OpenRouter.dump_all(env[:messages])
|
|
15
22
|
|
|
16
|
-
::OpenRouter::Client.new(
|
|
23
|
+
::OpenRouter::Client.new(**@config).then do |client|
|
|
17
24
|
client.complete(messages, @options).then do |response|
|
|
18
25
|
|
|
26
|
+
# Expose the provider's usage for downstream accounting
|
|
27
|
+
# middleware (goal budgets, autonomous limits, compaction
|
|
28
|
+
# thresholds, usage attribution) — additive metadata only.
|
|
29
|
+
if response.respond_to?(:usage) && response.usage
|
|
30
|
+
(env[:metadata] ||= {})[:last_llm_usage] = response.usage
|
|
31
|
+
end
|
|
32
|
+
|
|
19
33
|
# OpenRouter in fact only returns a single message...
|
|
20
|
-
# https://github.com/estiens/open_router_enhanced/blob/main/lib/open_router/response.rb
|
|
34
|
+
# https://github.com/estiens/open_router_enhanced/blob/main/lib/open_router/response.rb
|
|
21
35
|
Brute::MessageTransport::OpenRouter.wrap_each(response) do |message|
|
|
22
36
|
env[:messages] << message
|
|
23
37
|
end
|
|
24
38
|
end
|
|
25
39
|
end
|
|
26
40
|
|
|
41
|
+
env[:hooks]&.emit(:after_llm, env)
|
|
27
42
|
env
|
|
28
43
|
end
|
|
29
44
|
end
|
|
30
45
|
end
|
|
31
46
|
end
|
|
32
47
|
end
|
|
48
|
+
|
|
49
|
+
__END__
|
|
50
|
+
|
|
51
|
+
describe "brute/middleware/open_router" do
|
|
52
|
+
require "brute/messages"
|
|
53
|
+
|
|
54
|
+
# The repo suite has no open_router gem; stub the two constants the
|
|
55
|
+
# middleware touches (the transport wraps duck-typed responses fine).
|
|
56
|
+
begin
|
|
57
|
+
require "open_router"
|
|
58
|
+
rescue LoadError
|
|
59
|
+
module OpenRouter
|
|
60
|
+
CompletionOptions = Class.new { def initialize(**_opts); end }
|
|
61
|
+
Client = Class.new
|
|
62
|
+
end
|
|
63
|
+
end
|
|
64
|
+
|
|
65
|
+
FakeUsageResponse = Struct.new(:usage) do
|
|
66
|
+
def choices
|
|
67
|
+
[{ "message" => { "role" => "assistant", "content" => "hello" } }]
|
|
68
|
+
end
|
|
69
|
+
end
|
|
70
|
+
|
|
71
|
+
it "records the provider usage into env metadata and appends the message" do
|
|
72
|
+
response = FakeUsageResponse.new({ "prompt_tokens" => 10, "completion_tokens" => 5, "total_tokens" => 15 })
|
|
73
|
+
fake_client = Object.new
|
|
74
|
+
fake_client.define_singleton_method(:complete) { |_messages, _options| response }
|
|
75
|
+
original = OpenRouter::Client.method(:new)
|
|
76
|
+
OpenRouter::Client.define_singleton_method(:new) { |**_config| fake_client }
|
|
77
|
+
begin
|
|
78
|
+
middleware = Brute::Middleware::OpenRouter::Completion.new(->(env) { env })
|
|
79
|
+
env = { messages: Brute.log }
|
|
80
|
+
env[:messages].user("hi")
|
|
81
|
+
middleware.call(env)
|
|
82
|
+
|
|
83
|
+
env[:messages].last.role.should == :assistant
|
|
84
|
+
env[:metadata][:last_llm_usage]["total_tokens"].should == 15
|
|
85
|
+
ensure
|
|
86
|
+
OpenRouter::Client.define_singleton_method(:new, original)
|
|
87
|
+
end
|
|
88
|
+
end
|
|
89
|
+
|
|
90
|
+
it "leaves metadata alone when the response has no usage" do
|
|
91
|
+
response = FakeUsageResponse.new(nil)
|
|
92
|
+
fake_client = Object.new
|
|
93
|
+
fake_client.define_singleton_method(:complete) { |_messages, _options| response }
|
|
94
|
+
original = OpenRouter::Client.method(:new)
|
|
95
|
+
OpenRouter::Client.define_singleton_method(:new) { |**_config| fake_client }
|
|
96
|
+
begin
|
|
97
|
+
middleware = Brute::Middleware::OpenRouter::Completion.new(->(env) { env })
|
|
98
|
+
env = { messages: Brute.log }
|
|
99
|
+
env[:messages].user("hi")
|
|
100
|
+
middleware.call(env)
|
|
101
|
+
|
|
102
|
+
env.key?(:metadata).should.be.false
|
|
103
|
+
ensure
|
|
104
|
+
OpenRouter::Client.define_singleton_method(:new, original)
|
|
105
|
+
end
|
|
106
|
+
end
|
|
107
|
+
end
|
|
@@ -0,0 +1,125 @@
|
|
|
1
|
+
# frozen_string_literal: true
|
|
2
|
+
|
|
3
|
+
require "bundler/setup"
|
|
4
|
+
require "brute"
|
|
5
|
+
require "erb"
|
|
6
|
+
|
|
7
|
+
module Brute
|
|
8
|
+
# An ERB-backed system-prompt object for Middleware::SystemPrompt — the
|
|
9
|
+
# open alternative to Brute::SystemPrompt's built-in section stacks. You
|
|
10
|
+
# bring a template and named values; every keyword becomes an attr_accessor
|
|
11
|
+
# and an ERB local of the same name:
|
|
12
|
+
#
|
|
13
|
+
# prompt = Brute::PromptTemplate.new(
|
|
14
|
+
# "prompt.erb",
|
|
15
|
+
# identity: "You are Pico.",
|
|
16
|
+
# memory: -> { File.read("memory/MEMORY.md") }, # zero-arity proc
|
|
17
|
+
# env: ->(ctx) { Brute::Prompts::Environment.call(ctx) },
|
|
18
|
+
# )
|
|
19
|
+
# prompt.identity = "You are Brute." # attr_accessor per section
|
|
20
|
+
#
|
|
21
|
+
# Brute.agent.use(Brute::Middleware::SystemPrompt, system_prompt: prompt)
|
|
22
|
+
#
|
|
23
|
+
# Proc values are re-evaluated on every prepare (zero-arity procs are
|
|
24
|
+
# called with no arguments, others receive the turn ctx), and a template
|
|
25
|
+
# path is re-read from disk each time — so file-backed sections hot-reload
|
|
26
|
+
# between turns. The prepare(ctx) -> Result(#empty?, #to_s) contract is
|
|
27
|
+
# what Middleware::SystemPrompt expects.
|
|
28
|
+
class PromptTemplate
|
|
29
|
+
Result = Struct.new(:text) do
|
|
30
|
+
def to_s = text.to_s
|
|
31
|
+
def empty? = text.to_s.strip.empty?
|
|
32
|
+
end
|
|
33
|
+
|
|
34
|
+
def initialize(template, **sections)
|
|
35
|
+
@template = template
|
|
36
|
+
@section_keys = []
|
|
37
|
+
sections.each { |key, value| self[key] = value }
|
|
38
|
+
end
|
|
39
|
+
|
|
40
|
+
def [](key)
|
|
41
|
+
public_send(key)
|
|
42
|
+
end
|
|
43
|
+
|
|
44
|
+
def []=(key, value)
|
|
45
|
+
unless respond_to?(key)
|
|
46
|
+
singleton_class.class_eval { attr_accessor key }
|
|
47
|
+
@section_keys << key
|
|
48
|
+
end
|
|
49
|
+
public_send("#{key}=", value)
|
|
50
|
+
end
|
|
51
|
+
|
|
52
|
+
# Called once per turn by Middleware::SystemPrompt.
|
|
53
|
+
def prepare(ctx = {})
|
|
54
|
+
Result.new(render(locals(ctx)))
|
|
55
|
+
end
|
|
56
|
+
|
|
57
|
+
private
|
|
58
|
+
|
|
59
|
+
def locals(ctx)
|
|
60
|
+
@section_keys.to_h do |key|
|
|
61
|
+
value = self[key]
|
|
62
|
+
resolved = value.is_a?(Proc) ? (value.arity.zero? ? value.call : value.call(ctx)) : value
|
|
63
|
+
[key, resolved]
|
|
64
|
+
end.merge(ctx: ctx)
|
|
65
|
+
end
|
|
66
|
+
|
|
67
|
+
def render(values)
|
|
68
|
+
context = binding
|
|
69
|
+
values.each { |key, value| context.local_variable_set(key, value) }
|
|
70
|
+
ERB.new(template_source, trim_mode: "-").result(context)
|
|
71
|
+
end
|
|
72
|
+
|
|
73
|
+
# A path that exists is re-read every time; anything else is treated as
|
|
74
|
+
# an inline ERB source string.
|
|
75
|
+
def template_source
|
|
76
|
+
File.exist?(@template.to_s) ? File.read(@template) : @template.to_s
|
|
77
|
+
end
|
|
78
|
+
end
|
|
79
|
+
end
|
|
80
|
+
|
|
81
|
+
__END__
|
|
82
|
+
|
|
83
|
+
describe "brute/prompt_template" do
|
|
84
|
+
require "tmpdir"
|
|
85
|
+
|
|
86
|
+
it "renders keyword sections as ERB locals" do
|
|
87
|
+
prompt = Brute::PromptTemplate.new("Hello <%= name %>, <%= mood %> today.", name: "Pico", mood: "happy")
|
|
88
|
+
prompt.prepare.to_s.should == "Hello Pico, happy today."
|
|
89
|
+
end
|
|
90
|
+
|
|
91
|
+
it "exposes an attr_accessor per section" do
|
|
92
|
+
prompt = Brute::PromptTemplate.new("<%= name %>", name: "Pico")
|
|
93
|
+
prompt.name.should == "Pico"
|
|
94
|
+
prompt.name = "Brute"
|
|
95
|
+
prompt.prepare.to_s.should == "Brute"
|
|
96
|
+
end
|
|
97
|
+
|
|
98
|
+
it "re-evaluates zero-arity procs on every prepare" do
|
|
99
|
+
count = 0
|
|
100
|
+
prompt = Brute::PromptTemplate.new("<%= tick %>", tick: -> { count += 1 })
|
|
101
|
+
prompt.prepare.to_s.should == "1"
|
|
102
|
+
prompt.prepare.to_s.should == "2"
|
|
103
|
+
end
|
|
104
|
+
|
|
105
|
+
it "passes ctx to procs that take an argument" do
|
|
106
|
+
prompt = Brute::PromptTemplate.new("<%= who %>", who: ->(ctx) { ctx[:agent] })
|
|
107
|
+
prompt.prepare(agent: "pico").to_s.should == "pico"
|
|
108
|
+
end
|
|
109
|
+
|
|
110
|
+
it "re-reads a template file on every prepare" do
|
|
111
|
+
Dir.mktmpdir do |dir|
|
|
112
|
+
path = File.join(dir, "prompt.erb")
|
|
113
|
+
File.write(path, "v1 <%= x %>")
|
|
114
|
+
prompt = Brute::PromptTemplate.new(path, x: "a")
|
|
115
|
+
prompt.prepare.to_s.should == "v1 a"
|
|
116
|
+
File.write(path, "v2 <%= x %>")
|
|
117
|
+
prompt.prepare.to_s.should == "v2 a"
|
|
118
|
+
end
|
|
119
|
+
end
|
|
120
|
+
|
|
121
|
+
it "honours the Result contract (empty?/to_s)" do
|
|
122
|
+
Brute::PromptTemplate.new("").prepare.empty?.should.be.true
|
|
123
|
+
Brute::PromptTemplate.new("hi").prepare.empty?.should.be.false
|
|
124
|
+
end
|
|
125
|
+
end
|