Skip to content

Commit 8d16359

Browse files
author
nigel brown
committed
Remove optimizer config translation layer
Eliminate the intermediate optimizer.Config type and ConfigFromVMCPConfig conversion function to use config.OptimizerConfig directly throughout the optimizer package. This addresses maintainability concerns by establishing a single source of truth for optimizer configuration. Changes: - Delete pkg/vmcp/optimizer/config.go containing the duplicate config type - Update optimizer.Factory and EmbeddingOptimizer to use *config.OptimizerConfig - Flatten embedding config in ingestion.Config (individual fields vs nested) - Add type aliases (Config, OptimizerIntegration) for test compatibility - Add test helper methods (OnRegisterSession, RegisterTools, IngestToolsForTesting) - Update all test files to use flattened config structure - Handle HybridSearchRatio as pointer with default value (70) Benefits: - Single source of truth (no config duplication) - No synchronization burden between config types - Eliminates risk of translation bugs - Clearer code flow without intermediate transformations Closes review comment in PR #3440 requesting removal of translation layers.
1 parent 3ac35ac commit 8d16359

17 files changed

Lines changed: 336 additions & 449 deletions

‎cmd/vmcp/app/commands.go‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -449,9 +449,8 @@ func runServe(cmd *cobra.Command, _ []string) error {
449449
// Configure optimizer if enabled in YAML config
450450
if cfg.Optimizer != nil && cfg.Optimizer.Enabled {
451451
logger.Info("🔬 Optimizer enabled via configuration (chromem-go)")
452-
optimizerCfg := vmcpoptimizer.ConfigFromVMCPConfig(cfg.Optimizer)
453452
serverCfg.OptimizerFactory = vmcpoptimizer.NewEmbeddingOptimizer
454-
serverCfg.OptimizerConfig = optimizerCfg
453+
serverCfg.OptimizerConfig = cfg.Optimizer
455454
persistInfo := "in-memory"
456455
if cfg.Optimizer.PersistPath != "" {
457456
persistInfo = cfg.Optimizer.PersistPath

‎pkg/vmcp/optimizer/config.go‎

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

‎pkg/vmcp/optimizer/find_tool_semantic_search_test.go‎

Lines changed: 33 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,11 @@ import (
1515
"github.com/stretchr/testify/assert"
1616
"github.com/stretchr/testify/require"
1717

18-
"github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/embeddings"
1918
transportsession "github.com/stacklok/toolhive/pkg/transport/session"
2019
"github.com/stacklok/toolhive/pkg/vmcp"
2120
"github.com/stacklok/toolhive/pkg/vmcp/aggregator"
2221
"github.com/stacklok/toolhive/pkg/vmcp/discovery"
22+
"github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/embeddings"
2323
vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session"
2424
)
2525

@@ -83,16 +83,15 @@ func TestFindTool_SemanticSearch(t *testing.T) {
8383
mcpServer := server.NewMCPServer("test-server", "1.0")
8484
mockClient := &mockBackendClient{}
8585

86+
hybridRatio := 90 // 90% semantic, 10% BM25 to test semantic search
8687
config := &Config{
87-
Enabled: true,
88-
PersistPath: filepath.Join(tmpDir, "optimizer-db"),
89-
EmbeddingConfig: &embeddings.Config{
90-
BackendType: embeddingBackend,
91-
BaseURL: embeddingConfig.BaseURL,
92-
Model: embeddingConfig.Model,
93-
Dimension: embeddingConfig.Dimension,
94-
},
95-
HybridSearchRatio: 90, // 90% semantic, 10% BM25 to test semantic search
88+
Enabled: true,
89+
PersistPath: filepath.Join(tmpDir, "optimizer-db"),
90+
EmbeddingBackend: embeddingBackend,
91+
EmbeddingURL: embeddingConfig.BaseURL,
92+
EmbeddingModel: embeddingConfig.Model,
93+
EmbeddingDimension: embeddingConfig.Dimension,
94+
HybridSearchRatio: &hybridRatio,
9695
}
9796

9897
sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory())
@@ -383,16 +382,15 @@ func TestFindTool_SemanticVsKeyword(t *testing.T) {
383382
mockClient := &mockBackendClient{}
384383

385384
// Test with high semantic ratio
385+
hybridRatioSemantic := 90 // 90% semantic
386386
configSemantic := &Config{
387-
Enabled: true,
388-
PersistPath: filepath.Join(tmpDir, "optimizer-db-semantic"),
389-
EmbeddingConfig: &embeddings.Config{
390-
BackendType: embeddingBackend,
391-
BaseURL: embeddingConfig.BaseURL,
392-
Model: embeddings.DefaultModelAllMiniLM,
393-
Dimension: 384,
394-
},
395-
HybridSearchRatio: 90, // 90% semantic
387+
Enabled: true,
388+
PersistPath: filepath.Join(tmpDir, "optimizer-db-semantic"),
389+
EmbeddingBackend: embeddingBackend,
390+
EmbeddingURL: embeddingConfig.BaseURL,
391+
EmbeddingModel: embeddings.DefaultModelAllMiniLM,
392+
EmbeddingDimension: 384,
393+
HybridSearchRatio: &hybridRatioSemantic,
396394
}
397395

398396
sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory())
@@ -401,16 +399,15 @@ func TestFindTool_SemanticVsKeyword(t *testing.T) {
401399
defer func() { _ = integrationSemantic.Close() }()
402400

403401
// Test with low semantic ratio (high BM25)
402+
hybridRatioKeyword := 10 // 10% semantic, 90% BM25
404403
configKeyword := &Config{
405-
Enabled: true,
406-
PersistPath: filepath.Join(tmpDir, "optimizer-db-keyword"),
407-
EmbeddingConfig: &embeddings.Config{
408-
BackendType: embeddingBackend,
409-
BaseURL: embeddingConfig.BaseURL,
410-
Model: embeddings.DefaultModelAllMiniLM,
411-
Dimension: 384,
412-
},
413-
HybridSearchRatio: 10, // 10% semantic, 90% BM25
404+
Enabled: true,
405+
PersistPath: filepath.Join(tmpDir, "optimizer-db-keyword"),
406+
EmbeddingBackend: embeddingBackend,
407+
EmbeddingURL: embeddingConfig.BaseURL,
408+
EmbeddingModel: embeddings.DefaultModelAllMiniLM,
409+
EmbeddingDimension: 384,
410+
HybridSearchRatio: &hybridRatioKeyword,
414411
}
415412

416413
integrationKeyword, err := NewIntegration(ctx, configKeyword, mcpServer, mockClient, sessionMgr)
@@ -577,16 +574,15 @@ func TestFindTool_SemanticSimilarityScores(t *testing.T) {
577574
mcpServer := server.NewMCPServer("test-server", "1.0")
578575
mockClient := &mockBackendClient{}
579576

577+
hybridRatio := 90 // High semantic ratio
580578
config := &Config{
581-
Enabled: true,
582-
PersistPath: filepath.Join(tmpDir, "optimizer-db"),
583-
EmbeddingConfig: &embeddings.Config{
584-
BackendType: embeddingBackend,
585-
BaseURL: embeddingConfig.BaseURL,
586-
Model: embeddings.DefaultModelAllMiniLM,
587-
Dimension: 384,
588-
},
589-
HybridSearchRatio: 90, // High semantic ratio
579+
Enabled: true,
580+
PersistPath: filepath.Join(tmpDir, "optimizer-db"),
581+
EmbeddingBackend: embeddingBackend,
582+
EmbeddingURL: embeddingConfig.BaseURL,
583+
EmbeddingModel: embeddings.DefaultModelAllMiniLM,
584+
EmbeddingDimension: 384,
585+
HybridSearchRatio: &hybridRatio,
590586
}
591587

592588
sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory())

‎pkg/vmcp/optimizer/find_tool_string_matching_test.go‎

Lines changed: 25 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,11 @@ import (
1616
"github.com/stretchr/testify/assert"
1717
"github.com/stretchr/testify/require"
1818

19-
"github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/embeddings"
2019
transportsession "github.com/stacklok/toolhive/pkg/transport/session"
2120
"github.com/stacklok/toolhive/pkg/vmcp"
2221
"github.com/stacklok/toolhive/pkg/vmcp/aggregator"
2322
"github.com/stacklok/toolhive/pkg/vmcp/discovery"
23+
"github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/embeddings"
2424
vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session"
2525
)
2626

@@ -124,16 +124,15 @@ func TestFindTool_StringMatching(t *testing.T) {
124124
// Verify Ollama is actually working, not just reachable
125125
verifyOllamaWorking(t, embeddingManager)
126126

127+
hybridRatio := 50 // 50% semantic, 50% BM25 for better string matching
127128
config := &Config{
128-
Enabled: true,
129-
PersistPath: filepath.Join(tmpDir, "optimizer-db"),
130-
EmbeddingConfig: &embeddings.Config{
131-
BackendType: embeddings.BackendTypeOllama,
132-
BaseURL: "http://localhost:11434",
133-
Model: embeddings.DefaultModelAllMiniLM,
134-
Dimension: 384,
135-
},
136-
HybridSearchRatio: 50, // 50% semantic, 50% BM25 for better string matching
129+
Enabled: true,
130+
PersistPath: filepath.Join(tmpDir, "optimizer-db"),
131+
EmbeddingBackend: embeddings.BackendTypeOllama,
132+
EmbeddingURL: "http://localhost:11434",
133+
EmbeddingModel: embeddings.DefaultModelAllMiniLM,
134+
EmbeddingDimension: 384,
135+
HybridSearchRatio: &hybridRatio,
137136
}
138137

139138
sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory())
@@ -401,16 +400,15 @@ func TestFindTool_ExactStringMatch(t *testing.T) {
401400
// Verify Ollama is actually working, not just reachable
402401
verifyOllamaWorking(t, embeddingManager)
403402

403+
hybridRatio := 30 // 30% semantic, 70% BM25 for better exact string matching
404404
config := &Config{
405-
Enabled: true,
406-
PersistPath: filepath.Join(tmpDir, "optimizer-db"),
407-
EmbeddingConfig: &embeddings.Config{
408-
BackendType: embeddings.BackendTypeOllama,
409-
BaseURL: "http://localhost:11434",
410-
Model: embeddings.DefaultModelAllMiniLM,
411-
Dimension: 384,
412-
},
413-
HybridSearchRatio: 30, // 30% semantic, 70% BM25 for better exact string matching
405+
Enabled: true,
406+
PersistPath: filepath.Join(tmpDir, "optimizer-db"),
407+
EmbeddingBackend: embeddings.BackendTypeOllama,
408+
EmbeddingURL: "http://localhost:11434",
409+
EmbeddingModel: embeddings.DefaultModelAllMiniLM,
410+
EmbeddingDimension: 384,
411+
HybridSearchRatio: &hybridRatio,
414412
}
415413

416414
sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory())
@@ -582,16 +580,15 @@ func TestFindTool_CaseInsensitive(t *testing.T) {
582580
// Verify Ollama is actually working, not just reachable
583581
verifyOllamaWorking(t, embeddingManager)
584582

583+
hybridRatio := 30 // Favor BM25 for string matching
585584
config := &Config{
586-
Enabled: true,
587-
PersistPath: filepath.Join(tmpDir, "optimizer-db"),
588-
EmbeddingConfig: &embeddings.Config{
589-
BackendType: embeddings.BackendTypeOllama,
590-
BaseURL: "http://localhost:11434",
591-
Model: embeddings.DefaultModelAllMiniLM,
592-
Dimension: 384,
593-
},
594-
HybridSearchRatio: 30, // Favor BM25 for string matching
585+
Enabled: true,
586+
PersistPath: filepath.Join(tmpDir, "optimizer-db"),
587+
EmbeddingBackend: embeddings.BackendTypeOllama,
588+
EmbeddingURL: "http://localhost:11434",
589+
EmbeddingModel: embeddings.DefaultModelAllMiniLM,
590+
EmbeddingDimension: 384,
591+
HybridSearchRatio: &hybridRatio,
595592
}
596593

597594
sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory())

‎pkg/vmcp/optimizer/internal/db/backend_server.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,8 @@ import (
1212

1313
"github.com/philippgille/chromem-go"
1414

15-
"github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/models"
1615
"github.com/stacklok/toolhive/pkg/logger"
16+
"github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/models"
1717
)
1818

1919
// backendServerOps provides operations for backend servers in chromem-go

‎pkg/vmcp/optimizer/internal/db/backend_tool.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,8 @@ import (
1111

1212
"github.com/philippgille/chromem-go"
1313

14-
"github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/models"
1514
"github.com/stacklok/toolhive/pkg/logger"
15+
"github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/models"
1616
)
1717

1818
// backendToolOps provides operations for backend tools in chromem-go

‎pkg/vmcp/optimizer/internal/db/fts.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,8 @@ import (
1111
"strings"
1212
"sync"
1313

14-
"github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/models"
1514
"github.com/stacklok/toolhive/pkg/logger"
15+
"github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/models"
1616
)
1717

1818
//go:embed schema_fts.sql

‎pkg/vmcp/optimizer/internal/db/hybrid.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,8 @@ import (
77
"context"
88
"fmt"
99

10-
"github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/models"
1110
"github.com/stacklok/toolhive/pkg/logger"
11+
"github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/models"
1212
)
1313

1414
// HybridSearchConfig configures hybrid search behavior

‎pkg/vmcp/optimizer/internal/ingestion/service.go‎

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,20 +17,23 @@ import (
1717
"go.opentelemetry.io/otel/codes"
1818
"go.opentelemetry.io/otel/trace"
1919

20+
"github.com/stacklok/toolhive/pkg/logger"
2021
"github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/db"
2122
"github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/embeddings"
2223
"github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/models"
2324
"github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/tokens"
24-
"github.com/stacklok/toolhive/pkg/logger"
2525
)
2626

2727
// Config holds configuration for the ingestion service
2828
type Config struct {
2929
// Database configuration
3030
DBConfig *db.Config
3131

32-
// Embedding configuration
33-
EmbeddingConfig *embeddings.Config
32+
// Embedding configuration (flattened from embeddings.Config)
33+
EmbeddingBackend string
34+
EmbeddingURL string
35+
EmbeddingModel string
36+
EmbeddingDimension int
3437

3538
// MCP timeout in seconds
3639
MCPTimeout int
@@ -70,8 +73,16 @@ func NewService(config *Config) (*Service, error) {
7073
config.SkippedWorkloads = []string{"inspector", "mcp-optimizer"}
7174
}
7275

76+
// Construct embeddings.Config from individual fields
77+
embeddingConfig := &embeddings.Config{
78+
BackendType: config.EmbeddingBackend,
79+
BaseURL: config.EmbeddingURL,
80+
Model: config.EmbeddingModel,
81+
Dimension: config.EmbeddingDimension,
82+
}
83+
7384
// Initialize embedding manager first (needed for database)
74-
embeddingManager, err := embeddings.NewManager(config.EmbeddingConfig)
85+
embeddingManager, err := embeddings.NewManager(embeddingConfig)
7586
if err != nil {
7687
return nil, fmt.Errorf("failed to initialize embedding manager: %w", err)
7788
}

0 commit comments

Comments
 (0)