Skip to content

Commit 614c4b2

Browse files
committed
Refactor optimizer to interface-based architecture with consolidated packages
This refactoring improves code organization and testability by: **Interface-Based Design:** - Introduce clean `Optimizer` interface with 5 methods (FindTool, CallTool, Close, etc.) - Replace concrete `OptimizerIntegration` with interface-based approach - Implement `EmbeddingOptimizer` as production implementation - Use Factory pattern for dependency injection at server startup **Package Consolidation:** - Move implementation from `cmd/thv-operator/pkg/optimizer/` to `pkg/vmcp/optimizer/internal/` - Encapsulate implementation details in internal/ subdirectories: - internal/embeddings/ (Ollama, OpenAI-compatible, vLLM backends) - internal/db/ (chromem-go vectors + SQLite FTS5) - internal/ingestion/ (tool ingestion pipeline) - internal/models/ (data structures) - internal/tokens/ (token counting) - Public API now consists of only: optimizer.go, config.go, README.md **Server Integration:** - Replace `OptimizerIntegration` field with `Optimizer` interface - Add `OptimizerFactory` for clean dependency injection - Factory creates optimizer at startup with all required dependencies - Maintain backward compatibility with existing configs **Benefits:** - Better testability: Easy to mock Optimizer interface for unit tests - Cleaner separation: Public API vs internal implementation - Package design: Internal packages prevent external coupling - Extensibility: Easy to add new optimizer implementations **Migration:** - Old: `OptimizerIntegration optimizer.Integration` - New: `Optimizer optimizer.Optimizer` + `OptimizerFactory optimizer.Factory` - All existing functionality preserved - No breaking changes to CRD or YAML configs This addresses code review feedback to consolidate optimizer packages and implement a clean interface-based architecture.
1 parent 679ce3f commit 614c4b2

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

47 files changed

+932
-766
lines changed

cmd/thv-operator/pkg/optimizer/doc.go

Lines changed: 0 additions & 88 deletions
This file was deleted.

cmd/vmcp/app/commands.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -450,6 +450,7 @@ func runServe(cmd *cobra.Command, _ []string) error {
450450
if cfg.Optimizer != nil && cfg.Optimizer.Enabled {
451451
logger.Info("🔬 Optimizer enabled via configuration (chromem-go)")
452452
optimizerCfg := vmcpoptimizer.ConfigFromVMCPConfig(cfg.Optimizer)
453+
serverCfg.OptimizerFactory = vmcpoptimizer.NewEmbeddingOptimizer
453454
serverCfg.OptimizerConfig = optimizerCfg
454455
persistInfo := "in-memory"
455456
if cfg.Optimizer.PersistPath != "" {

pkg/vmcp/optimizer/README.md

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
# VMCPOptimizer Package
2+
3+
This package provides semantic tool discovery for Virtual MCP Server, reducing token usage by allowing LLMs to discover relevant tools on-demand instead of receiving all tool definitions upfront.
4+
5+
## Architecture
6+
7+
The optimizer exposes a clean interface-based architecture:
8+
9+
```
10+
pkg/vmcp/optimizer/
11+
├── optimizer.go # Public Optimizer interface and EmbeddingOptimizer implementation
12+
├── config.go # Configuration types
13+
├── README.md # This file
14+
└── internal/ # Implementation details (not part of public API)
15+
├── embeddings/ # Embedding backends (Ollama, OpenAI-compatible, vLLM)
16+
├── db/ # Database operations (chromem-go vectors, SQLite FTS5)
17+
├── ingestion/ # Tool ingestion service
18+
├── models/ # Internal data models
19+
└── tokens/ # Token counting utilities
20+
```
21+
22+
## Public API
23+
24+
### Optimizer Interface
25+
26+
```go
27+
type Optimizer interface {
28+
// FindTool searches for tools matching the description and keywords
29+
FindTool(ctx context.Context, input FindToolInput) (*FindToolOutput, error)
30+
31+
// CallTool invokes a tool by name with parameters
32+
CallTool(ctx context.Context, input CallToolInput) (*mcp.CallToolResult, error)
33+
34+
// Close cleans up optimizer resources
35+
Close() error
36+
37+
// HandleSessionRegistration handles session setup for optimizer mode
38+
HandleSessionRegistration(...) (bool, error)
39+
40+
// OptimizerHandlerProvider provides tool handlers for MCP integration
41+
adapter.OptimizerHandlerProvider
42+
}
43+
```
44+
45+
### Factory Pattern
46+
47+
```go
48+
// Factory creates an Optimizer instance
49+
type Factory func(
50+
ctx context.Context,
51+
cfg *Config,
52+
mcpServer *server.MCPServer,
53+
backendClient vmcp.BackendClient,
54+
sessionManager *transportsession.Manager,
55+
) (Optimizer, error)
56+
57+
// NewEmbeddingOptimizer is the production implementation
58+
func NewEmbeddingOptimizer(...) (Optimizer, error)
59+
```
60+
61+
## Usage
62+
63+
### In vMCP Server
64+
65+
```go
66+
import "github.com/stacklok/toolhive/pkg/vmcp/optimizer"
67+
68+
// Configure server with optimizer
69+
serverCfg := &vmcpserver.Config{
70+
OptimizerFactory: optimizer.NewEmbeddingOptimizer,
71+
OptimizerConfig: &optimizer.Config{
72+
Enabled: true,
73+
PersistPath: "/data/optimizer",
74+
HybridSearchRatio: 70, // 70% semantic, 30% keyword
75+
EmbeddingConfig: &embeddings.Config{
76+
BackendType: "ollama",
77+
BaseURL: "http://localhost:11434",
78+
Model: "nomic-embed-text",
79+
Dimension: 768,
80+
},
81+
},
82+
}
83+
```
84+
85+
### MCP Tools Exposed
86+
87+
When the optimizer is enabled, vMCP exposes two tools instead of all backend tools:
88+
89+
1. **`optim_find_tool`**: Semantic search for tools
90+
- Input: `tool_description` (natural language), optional `tool_keywords`, `limit`
91+
- Output: Ranked tools with similarity scores and token metrics
92+
93+
2. **`optim_call_tool`**: Dynamic tool invocation
94+
- Input: `backend_id`, `tool_name`, `parameters`
95+
- Output: Tool execution result
96+
97+
## Benefits
98+
99+
- **Token Savings**: Only relevant tools are sent to the LLM (typically 80-95% reduction)
100+
- **Hybrid Search**: Combines semantic embeddings (70%) with BM25 keyword matching (30%)
101+
- **Startup Ingestion**: Tools are indexed once at startup, not per-session
102+
- **Clean Architecture**: Interface-based design allows easy testing and alternative implementations
103+
104+
## Implementation Details
105+
106+
The `internal/` directory contains implementation details that are not part of the public API:
107+
108+
- **embeddings/**: Pluggable embedding backends (Ollama, vLLM, OpenAI-compatible)
109+
- **db/**: Hybrid search using chromem-go (vector DB) + SQLite FTS5 (BM25)
110+
- **ingestion/**: Tool ingestion pipeline with background embedding generation
111+
- **models/**: Internal data structures for backend tools and metadata
112+
- **tokens/**: Token counting for metrics calculation
113+
114+
These internal packages use internal import paths and cannot be imported from outside the optimizer package.
115+
116+
## Testing
117+
118+
The interface-based design enables easy testing:
119+
120+
```go
121+
// Mock the interface for unit tests
122+
mockOpt := mocks.NewMockOptimizer(ctrl)
123+
mockOpt.EXPECT().FindTool(...).Return(...)
124+
mockOpt.EXPECT().Close()
125+
126+
// Use in server configuration
127+
cfg.Optimizer = mockOpt
128+
```
129+
130+
## Migration from Integration Pattern
131+
132+
Previous versions used an `Integration` interface. The current `Optimizer` interface provides the same functionality with cleaner separation of concerns:
133+
134+
**Before (Integration):**
135+
- `OptimizerIntegration optimizer.Integration`
136+
- `optimizer.NewIntegration(...)`
137+
138+
**After (Optimizer):**
139+
- `Optimizer optimizer.Optimizer`
140+
- `OptimizerFactory optimizer.Factory`
141+
- `optimizer.NewEmbeddingOptimizer(...)`
142+
143+
The factory pattern allows the server to create the optimizer at startup with all necessary dependencies.

0 commit comments

Comments
 (0)