55 lines
1.8 KiB
Go
55 lines
1.8 KiB
Go
|
package rag
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"path/filepath"
|
||
|
|
"strings"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"code.tczkiot.com/wlw/ai-agent/contract"
|
||
|
|
"code.tczkiot.com/wlw/ai-agent/internal/ai"
|
||
|
|
"code.tczkiot.com/wlw/ai-agent/internal/ai/rag/vectordb"
|
||
|
|
"code.tczkiot.com/wlw/ai-agent/internal/pkg/config"
|
||
|
|
)
|
||
|
|
|
||
|
|
type dimensionTestPlatformProvider struct{}
|
||
|
|
|
||
|
|
func (dimensionTestPlatformProvider) ModelSource(context.Context) (string, error) {
|
||
|
|
return contract.ModelSourcePlatform, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (dimensionTestPlatformProvider) Config(context.Context) (*contract.PlatformAIConfig, error) {
|
||
|
|
return &contract.PlatformAIConfig{
|
||
|
|
APIKey: "license-signed",
|
||
|
|
BaseURL: "https://platform.example/v1",
|
||
|
|
EmbeddingModel: "qwen3.7-text-embedding",
|
||
|
|
EmbeddingDimension: 4,
|
||
|
|
}, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (dimensionTestPlatformProvider) Status(context.Context) (*contract.PlatformAIStatus, error) {
|
||
|
|
return &contract.PlatformAIStatus{Enabled: true, EmbeddingEnabled: true}, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestEnsureCollectionRejectsChangedEmbeddingDimension(t *testing.T) {
|
||
|
|
if err := vectordb.Init(&config.VectorDBConfig{Path: filepath.Join(t.TempDir(), "vectors.db")}); err != nil {
|
||
|
|
t.Fatalf("vectordb.Init() error = %v", err)
|
||
|
|
}
|
||
|
|
t.Cleanup(func() { _ = vectordb.Close() })
|
||
|
|
provider := vectordb.GetProvider()
|
||
|
|
if err := provider.CreateCollection(context.Background(), knowledgeCollectionName, 3); err != nil {
|
||
|
|
t.Fatalf("CreateCollection() error = %v", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
ai.SetPlatformAIProvider(dimensionTestPlatformProvider{})
|
||
|
|
t.Cleanup(func() { ai.SetPlatformAIProvider(nil) })
|
||
|
|
|
||
|
|
err := Index.EnsureCollection(context.Background())
|
||
|
|
if err == nil {
|
||
|
|
t.Fatal("expected dimension mismatch error")
|
||
|
|
}
|
||
|
|
if message := err.Error(); !strings.Contains(message, "dimension is 3") || !strings.Contains(message, "uses 4") || !strings.Contains(message, "rebuild") {
|
||
|
|
t.Fatalf("unexpected dimension mismatch error: %v", err)
|
||
|
|
}
|
||
|
|
}
|