Skip to content

Commit ad34abf

Browse files
RaezilRaezil
authored
feat: caching tool selection (#150)
1 parent 1a652c0 commit ad34abf

File tree

5 files changed

+1026
-1
lines changed

5 files changed

+1026
-1
lines changed

‎src/plugins/codemode/cache_benchmark_test.go‎

Lines changed: 228 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,228 @@
1+
package codemode
2+
3+
import (
4+
"context"
5+
"testing"
6+
7+
"github.com/universal-tool-calling-protocol/go-utcp/src/tools"
8+
)
9+
10+
// Benchmark tool specs without cache
11+
func BenchmarkToolSpecs_NoCache(b *testing.B) {
12+
mock := &mockUTCP{
13+
searchToolsFn: func(query string, limit int) ([]tools.Tool, error) {
14+
// Simulate SearchTools returning a large list
15+
result := make([]tools.Tool, 50)
16+
for i := 0; i < 50; i++ {
17+
result[i] = tools.Tool{
18+
Name: "test.tool" + string(rune(i)),
19+
Description: "Test tool " + string(rune(i)),
20+
}
21+
}
22+
return result, nil
23+
},
24+
}
25+
26+
mockModel := &mockModel{
27+
GenerateFunc: func(ctx context.Context, prompt string) (any, error) {
28+
return "{}", nil
29+
},
30+
}
31+
32+
cm := NewCodeModeUTCP(mock, mockModel)
33+
// Disable cache to benchmark without caching
34+
cm.cache = nil
35+
36+
b.ResetTimer()
37+
for i := 0; i < b.N; i++ {
38+
_ = cm.ToolSpecs()
39+
}
40+
}
41+
42+
// Benchmark tool specs with cache (first call - miss)
43+
func BenchmarkToolSpecs_WithCache_Miss(b *testing.B) {
44+
mock := &mockUTCP{
45+
searchToolsFn: func(query string, limit int) ([]tools.Tool, error) {
46+
result := make([]tools.Tool, 50)
47+
for i := 0; i < 50; i++ {
48+
result[i] = tools.Tool{
49+
Name: "test.tool" + string(rune(i)),
50+
Description: "Test tool " + string(rune(i)),
51+
}
52+
}
53+
return result, nil
54+
},
55+
}
56+
57+
mockModel := &mockModel{
58+
GenerateFunc: func(ctx context.Context, prompt string) (any, error) {
59+
return "{}", nil
60+
},
61+
}
62+
63+
b.ResetTimer()
64+
for i := 0; i < b.N; i++ {
65+
b.StopTimer()
66+
cm := NewCodeModeUTCP(mock, mockModel)
67+
b.StartTimer()
68+
_ = cm.ToolSpecs()
69+
}
70+
}
71+
72+
// Benchmark tool specs with cache (subsequent calls - hits)
73+
func BenchmarkToolSpecs_WithCache_Hit(b *testing.B) {
74+
mock := &mockUTCP{
75+
searchToolsFn: func(query string, limit int) ([]tools.Tool, error) {
76+
result := make([]tools.Tool, 50)
77+
for i := 0; i < 50; i++ {
78+
result[i] = tools.Tool{
79+
Name: "test.tool" + string(rune(i)),
80+
Description: "Test tool " + string(rune(i)),
81+
}
82+
}
83+
return result, nil
84+
},
85+
}
86+
87+
mockModel := &mockModel{
88+
GenerateFunc: func(ctx context.Context, prompt string) (any, error) {
89+
return "{}", nil
90+
},
91+
}
92+
93+
cm := NewCodeModeUTCP(mock, mockModel)
94+
// Warm up cache
95+
_ = cm.ToolSpecs()
96+
97+
b.ResetTimer()
98+
for i := 0; i < b.N; i++ {
99+
_ = cm.ToolSpecs()
100+
}
101+
}
102+
103+
// Benchmark selectTools without cache
104+
func BenchmarkSelectTools_NoCache(b *testing.B) {
105+
mock := &mockUTCP{}
106+
mockModel := &mockModel{
107+
GenerateFunc: func(ctx context.Context, prompt string) (any, error) {
108+
return `{"tools": ["test.tool1", "test.tool2"]}`, nil
109+
},
110+
}
111+
112+
cm := NewCodeModeUTCP(mock, mockModel)
113+
// Disable cache
114+
cm.cache = nil
115+
116+
ctx := context.Background()
117+
query := "find memory tools"
118+
toolsStr := "test.tool1, test.tool2, test.tool3"
119+
120+
b.ResetTimer()
121+
for i := 0; i < b.N; i++ {
122+
_, _ = cm.selectTools(ctx, query, toolsStr)
123+
}
124+
}
125+
126+
// Benchmark selectTools with cache (first call - miss)
127+
func BenchmarkSelectTools_WithCache_Miss(b *testing.B) {
128+
mock := &mockUTCP{}
129+
mockModel := &mockModel{
130+
GenerateFunc: func(ctx context.Context, prompt string) (any, error) {
131+
return `{"tools": ["test.tool1", "test.tool2"]}`, nil
132+
},
133+
}
134+
135+
ctx := context.Background()
136+
query := "find memory tools"
137+
toolsStr := "test.tool1, test.tool2, test.tool3"
138+
139+
b.ResetTimer()
140+
for i := 0; i < b.N; i++ {
141+
b.StopTimer()
142+
cm := NewCodeModeUTCP(mock, mockModel)
143+
b.StartTimer()
144+
_, _ = cm.selectTools(ctx, query, toolsStr)
145+
}
146+
}
147+
148+
// Benchmark selectTools with cache (subsequent calls - hits)
149+
func BenchmarkSelectTools_WithCache_Hit(b *testing.B) {
150+
mock := &mockUTCP{}
151+
mockModel := &mockModel{
152+
GenerateFunc: func(ctx context.Context, prompt string) (any, error) {
153+
return `{"tools": ["test.tool1", "test.tool2"]}`, nil
154+
},
155+
}
156+
157+
cm := NewCodeModeUTCP(mock, mockModel)
158+
159+
ctx := context.Background()
160+
query := "find memory tools"
161+
toolsStr := "test.tool1, test.tool2, test.tool3"
162+
163+
// Warm up cache
164+
_, _ = cm.selectTools(ctx, query, toolsStr)
165+
166+
b.ResetTimer()
167+
for i := 0; i < b.N; i++ {
168+
_, _ = cm.selectTools(ctx, query, toolsStr)
169+
}
170+
}
171+
172+
// Benchmark full CallTool workflow with cache
173+
func BenchmarkCallTool_WithCache(b *testing.B) {
174+
mock := &mockUTCP{
175+
searchToolsFn: func(query string, limit int) ([]tools.Tool, error) {
176+
return []tools.Tool{
177+
{Name: "test.tool1", Description: "Test tool 1"},
178+
{Name: "test.tool2", Description: "Test tool 2"},
179+
}, nil
180+
},
181+
}
182+
183+
mockModel := &mockModel{
184+
GenerateFunc: func(ctx context.Context, prompt string) (any, error) {
185+
// Simulate different responses for different stages
186+
if stringContains(prompt, "Decide if") {
187+
return `{"needs": true}`, nil
188+
}
189+
if stringContains(prompt, "Select ALL UTCP tools") {
190+
return `{"tools": ["test.tool1"]}`, nil
191+
}
192+
if stringContains(prompt, "Generate a Go snippet") {
193+
return `{"code": "__out = map[string]any{\"result\": \"test\"}", "stream": false}`, nil
194+
}
195+
return "{}", nil
196+
},
197+
}
198+
199+
cm := NewCodeModeUTCP(mock, mockModel)
200+
// Mock execute to avoid actual code execution
201+
cm.executeFunc = func(ctx context.Context, args CodeModeArgs) (CodeModeResult, error) {
202+
return CodeModeResult{Value: "mocked result"}, nil
203+
}
204+
205+
ctx := context.Background()
206+
query := "test query"
207+
208+
// First call to warm up cache
209+
_, _, _ = cm.CallTool(ctx, query)
210+
211+
b.ResetTimer()
212+
for i := 0; i < b.N; i++ {
213+
_, _, _ = cm.CallTool(ctx, query)
214+
}
215+
}
216+
217+
// Helper function
218+
func stringContains(s, substr string) bool {
219+
if len(s) < len(substr) {
220+
return false
221+
}
222+
for i := 0; i <= len(s)-len(substr); i++ {
223+
if s[i:i+len(substr)] == substr {
224+
return true
225+
}
226+
}
227+
return false
228+
}

‎src/plugins/codemode/codemode.go‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,12 +72,18 @@ type CodeModeUTCP struct {
7272
}
7373
// For testing purposes, to mock the Execute method.
7474
executeFunc func(ctx context.Context, args CodeModeArgs) (CodeModeResult, error)
75+
// Cache for tool specs and selection results
76+
cache *ToolCache
7577
}
7678

7779
func NewCodeModeUTCP(client utcp.UtcpClientInterface, model interface {
7880
Generate(ctx context.Context, prompt string) (any, error)
7981
}) *CodeModeUTCP {
80-
return &CodeModeUTCP{client: client, model: model}
82+
return &CodeModeUTCP{
83+
client: client,
84+
model: model,
85+
cache: NewToolCache(),
86+
}
8187
}
8288

8389
//

‎src/plugins/codemode/orchestrator.go‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88
"os"
99
"strconv"
1010
"strings"
11+
"time"
1112

1213
"github.com/universal-tool-calling-protocol/go-utcp/src/tools"
1314
)
@@ -287,6 +288,13 @@ func renderUtcpToolsForPrompt(specs []tools.Tool) string {
287288
}
288289

289290
func (a *CodeModeUTCP) ToolSpecs() []tools.Tool {
291+
// Check cache first
292+
if a.cache != nil {
293+
if cached := a.cache.GetToolSpecs(); cached != nil {
294+
return cached
295+
}
296+
}
297+
290298
var allSpecs []tools.Tool
291299
seen := make(map[string]bool)
292300

@@ -319,6 +327,12 @@ func (a *CodeModeUTCP) ToolSpecs() []tools.Tool {
319327
}
320328
}
321329
}
330+
331+
// Store in cache
332+
if a.cache != nil {
333+
a.cache.SetToolSpecs(allSpecs)
334+
}
335+
322336
return allSpecs
323337
}
324338

@@ -463,6 +477,13 @@ func (cm *CodeModeUTCP) selectTools(
463477
tools string,
464478
) ([]string, error) {
465479

480+
// Check cache first
481+
if cm.cache != nil {
482+
if cached := cm.cache.GetSelectedTools(query, tools); cached != nil {
483+
return cached, nil
484+
}
485+
}
486+
466487
prompt := fmt.Sprintf(`
467488
Select ALL UTCP tools that match the user's intent.
468489
@@ -498,5 +519,52 @@ Rules:
498519
}
499520

500521
_ = json.Unmarshal([]byte(jsonStr), &resp)
522+
523+
// Store in cache
524+
if cm.cache != nil && resp.Tools != nil {
525+
cm.cache.SetSelectedTools(query, tools, resp.Tools)
526+
}
527+
501528
return resp.Tools, nil
502529
}
530+
531+
// ───────────────────────────────────────────────────────────
532+
// Cache Management Methods
533+
// ───────────────────────────────────────────────────────────
534+
535+
// InvalidateToolSpecsCache clears the cached tool specifications
536+
func (cm *CodeModeUTCP) InvalidateToolSpecsCache() {
537+
if cm.cache != nil {
538+
cm.cache.InvalidateToolSpecs()
539+
}
540+
}
541+
542+
// InvalidateSelectionsCache clears all cached tool selection results
543+
func (cm *CodeModeUTCP) InvalidateSelectionsCache() {
544+
if cm.cache != nil {
545+
cm.cache.InvalidateSelections()
546+
}
547+
}
548+
549+
// InvalidateAllCaches clears all caches (tool specs and selections)
550+
func (cm *CodeModeUTCP) InvalidateAllCaches() {
551+
if cm.cache != nil {
552+
cm.cache.InvalidateAll()
553+
}
554+
}
555+
556+
// CacheStats returns performance statistics for the tool cache
557+
func (cm *CodeModeUTCP) CacheStats() CacheStats {
558+
if cm.cache == nil {
559+
return CacheStats{}
560+
}
561+
return cm.cache.Stats()
562+
}
563+
564+
// StartCacheCleanup starts a background routine to clean expired cache entries
565+
// Call this with a context to control the cleanup lifecycle
566+
func (cm *CodeModeUTCP) StartCacheCleanup(ctx context.Context, interval time.Duration) {
567+
if cm.cache != nil {
568+
cm.cache.StartCleanupRoutine(ctx, interval)
569+
}
570+
}

0 commit comments

Comments
 (0)