Skip to content

Commit f9daca2

Browse files
ochanismochanism
committed
feat: pass IndexConfig to repositories with scalability options
1 parent c5fc05f commit f9daca2

14 files changed

Lines changed: 673 additions & 69 deletions

File tree

‎internal/application/repository/retriever/elasticsearch/v7/repository.go‎

Lines changed: 72 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@ import (
1212
"encoding/json"
1313
"errors"
1414
"fmt"
15-
"os"
1615
"strings"
1716

1817
elasticsearchRetriever "github.com/Tencent/WeKnora/internal/application/repository/retriever/elasticsearch"
@@ -29,22 +28,31 @@ type elasticsearchRepository struct {
2928
client *elasticsearch.Client
3029
index string
3130
useKeywordSuffix bool // Whether to append .keyword suffix to ID field names in queries
31+
numberOfShards int // Shard count for index creation (0 = ES default)
32+
numberOfReplicas int // Replica count for index creation (-1 = unset, use ES default)
3233
}
3334

35+
// NewElasticsearchEngineRepository creates and initializes a new Elasticsearch v7 repository.
36+
// indexCfg is optional — pass nil to use env var / default values (env path).
3437
func NewElasticsearchEngineRepository(client *elasticsearch.Client,
3538
config *config.Config,
39+
indexCfg *typesLocal.IndexConfig,
3640
) interfaces.RetrieveEngineRepository {
3741
log := logger.GetLogger(context.Background())
3842
log.Info("[ElasticsearchV7] Initializing Elasticsearch v7 retriever engine repository")
3943

40-
indexName := os.Getenv("ELASTICSEARCH_INDEX")
41-
if indexName == "" {
42-
log.Warn("[ElasticsearchV7] ELASTICSEARCH_INDEX environment variable not set, using default index name")
43-
indexName = "xwrag_default"
44-
}
44+
indexName := typesLocal.ResolveIndexName(indexCfg, "ELASTICSEARCH_INDEX", "xwrag_default")
4545

4646
log.Infof("[ElasticsearchV7] Using index: %s", indexName)
47-
res := &elasticsearchRepository{client: client, index: indexName}
47+
res := &elasticsearchRepository{
48+
client: client,
49+
index: indexName,
50+
numberOfShards: indexCfg.GetNumberOfShards(0),
51+
numberOfReplicas: indexCfg.GetNumberOfReplicas(-1),
52+
}
53+
if err := res.createIndexIfNotExists(context.Background()); err != nil {
54+
log.Errorf("[ElasticsearchV7] Failed to create index: %v", err)
55+
}
4856
res.detectFieldTypes(context.Background())
4957
return res
5058
}
@@ -118,6 +126,63 @@ func (e *elasticsearchRepository) detectFieldTypes(ctx context.Context) {
118126
}
119127
}
120128

129+
// createIndexIfNotExists checks if the specified index exists and creates it if not.
130+
// Uses esapi low-level client since v7 SDK does not have typed API.
131+
func (e *elasticsearchRepository) createIndexIfNotExists(ctx context.Context) error {
132+
log := logger.GetLogger(ctx)
133+
log.Debugf("[ElasticsearchV7] Checking if index exists: %s", e.index)
134+
135+
res, err := e.client.Indices.Exists([]string{e.index}, e.client.Indices.Exists.WithContext(ctx))
136+
if err != nil {
137+
return fmt.Errorf("check index existence: %w", err)
138+
}
139+
defer res.Body.Close()
140+
141+
if !res.IsError() {
142+
log.Debugf("[ElasticsearchV7] Index already exists: %s", e.index)
143+
return nil
144+
}
145+
146+
// Build settings body with optional shards/replicas
147+
var body string
148+
if e.numberOfShards > 0 || e.numberOfReplicas >= 0 {
149+
settings := make(map[string]interface{})
150+
if e.numberOfShards > 0 {
151+
settings["number_of_shards"] = e.numberOfShards
152+
}
153+
if e.numberOfReplicas >= 0 {
154+
settings["number_of_replicas"] = e.numberOfReplicas
155+
}
156+
bodyBytes, err := json.Marshal(map[string]interface{}{"settings": settings})
157+
if err != nil {
158+
return fmt.Errorf("marshal index settings: %w", err)
159+
}
160+
body = string(bodyBytes)
161+
}
162+
163+
log.Infof("[ElasticsearchV7] Creating index: %s", e.index)
164+
var opts []func(*esapi.IndicesCreateRequest)
165+
opts = append(opts, e.client.Indices.Create.WithContext(ctx))
166+
if body != "" {
167+
opts = append(opts, e.client.Indices.Create.WithBody(strings.NewReader(body)))
168+
}
169+
170+
createRes, err := e.client.Indices.Create(e.index, opts...)
171+
if err != nil {
172+
return fmt.Errorf("create index: %w", err)
173+
}
174+
defer createRes.Body.Close()
175+
176+
if createRes.IsError() {
177+
// Log detailed response server-side; return generic message to avoid leaking cluster info
178+
log.Errorf("[ElasticsearchV7] Create index response: %s", createRes.String())
179+
return fmt.Errorf("failed to create index %s", e.index)
180+
}
181+
182+
log.Infof("[ElasticsearchV7] Index created successfully: %s", e.index)
183+
return nil
184+
}
185+
121186
func (e *elasticsearchRepository) EngineType() typesLocal.RetrieverEngineType {
122187
return typesLocal.ElasticsearchRetrieverEngineType
123188
}

‎internal/application/repository/retriever/elasticsearch/v8/repository.go‎

Lines changed: 25 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import (
44
"context"
55
"encoding/json"
66
"fmt"
7-
"os"
87
"strings"
98

109
elasticsearchRetriever "github.com/Tencent/WeKnora/internal/application/repository/retriever/elasticsearch"
@@ -24,25 +23,28 @@ type elasticsearchRepository struct {
2423
client *elasticsearch.TypedClient // Elasticsearch client instance
2524
index string // Name of the Elasticsearch index to use
2625
useKeywordSuffix bool // Whether to append .keyword suffix to ID field names in queries
26+
numberOfShards int // Shard count for index creation (0 = ES default)
27+
numberOfReplicas int // Replica count for index creation (-1 = unset, use ES default)
2728
}
2829

29-
// NewElasticsearchEngineRepository creates and initializes a new Elasticsearch v8 repository
30-
// It sets up the index and returns a repository instance ready for use
30+
// NewElasticsearchEngineRepository creates and initializes a new Elasticsearch v8 repository.
31+
// indexCfg is optional — pass nil to use env var / default values (env path).
3132
func NewElasticsearchEngineRepository(client *elasticsearch.TypedClient,
3233
config *config.Config,
34+
indexCfg *typesLocal.IndexConfig,
3335
) interfaces.RetrieveEngineRepository {
3436
log := logger.GetLogger(context.Background())
3537
log.Info("[Elasticsearch] Initializing Elasticsearch v8 retriever engine repository")
3638

37-
// Get index name from environment variable or use default
38-
indexName := os.Getenv("ELASTICSEARCH_INDEX")
39-
if indexName == "" {
40-
log.Warn("[Elasticsearch] ELASTICSEARCH_INDEX environment variable not set, using default index name")
41-
indexName = "xwrag_default"
42-
}
39+
indexName := typesLocal.ResolveIndexName(indexCfg, "ELASTICSEARCH_INDEX", "xwrag_default")
4340

4441
// Create repository instance and ensure index exists
45-
res := &elasticsearchRepository{client: client, index: indexName}
42+
res := &elasticsearchRepository{
43+
client: client,
44+
index: indexName,
45+
numberOfShards: indexCfg.GetNumberOfShards(0),
46+
numberOfReplicas: indexCfg.GetNumberOfReplicas(-1),
47+
}
4648
if err := res.createIndexIfNotExists(context.Background()); err != nil {
4749
log.Errorf("[Elasticsearch] Failed to create index: %v", err)
4850
} else {
@@ -355,9 +357,20 @@ func (e *elasticsearchRepository) createIndexIfNotExists(ctx context.Context) er
355357
return nil
356358
}
357359

358-
// Create index if it doesn't exist
360+
// Create index if it doesn't exist, with optional shards/replicas settings
359361
log.Infof("[Elasticsearch] Creating index: %s", e.index)
360-
_, err = e.client.Indices.Create(e.index).Do(ctx)
362+
createReq := e.client.Indices.Create(e.index)
363+
if e.numberOfShards > 0 || e.numberOfReplicas >= 0 {
364+
settings := &types.IndexSettings{}
365+
if e.numberOfShards > 0 {
366+
settings.NumberOfShards = fmt.Sprintf("%d", e.numberOfShards)
367+
}
368+
if e.numberOfReplicas >= 0 {
369+
settings.NumberOfReplicas = fmt.Sprintf("%d", e.numberOfReplicas)
370+
}
371+
createReq = createReq.Settings(settings)
372+
}
373+
_, err = createReq.Do(ctx)
361374
if err != nil {
362375
log.Errorf("[Elasticsearch] Failed to create index: %v", err)
363376
return err

‎internal/application/repository/retriever/milvus/repository.go‎

Lines changed: 16 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -41,16 +41,13 @@ var (
4141
fieldKnowledgeID, fieldKnowledgeBaseID, fieldTagID, fieldIsEnabled, fieldEmbedding}
4242
)
4343

44-
// NewMilvusRetrieveEngineRepository creates and initializes a new Milvus repository
45-
func NewMilvusRetrieveEngineRepository(client *client.Client) interfaces.RetrieveEngineRepository {
44+
// NewMilvusRetrieveEngineRepository creates and initializes a new Milvus repository.
45+
// indexCfg is optional — pass nil to use env var / default values (env path).
46+
func NewMilvusRetrieveEngineRepository(client *client.Client, indexCfg *types.IndexConfig) interfaces.RetrieveEngineRepository {
4647
log := logger.GetLogger(context.Background())
4748
log.Info("[Milvus] Initializing Milvus retriever engine repository")
4849

49-
collectionBaseName := os.Getenv(envMilvusCollection)
50-
if collectionBaseName == "" {
51-
log.Warn("[Milvus] MILVUS_COLLECTION environment variable not set, using default collection name")
52-
collectionBaseName = defaultCollectionName
53-
}
50+
collectionBaseName := types.ResolveCollectionName(indexCfg, envMilvusCollection, defaultCollectionName)
5451

5552
metricType := entity.IP
5653
if mt := os.Getenv(envMilvusMetricType); mt != "" {
@@ -72,6 +69,8 @@ func NewMilvusRetrieveEngineRepository(client *client.Client) interfaces.Retriev
7269
client: client,
7370
collectionBaseName: collectionBaseName,
7471
metricType: metricType,
72+
shardsNum: indexCfg.GetShardsNum(0),
73+
replicaNumber: indexCfg.GetReplicaNumber(0),
7574
}
7675

7776
log.Info("[Milvus] Successfully initialized repository")
@@ -176,7 +175,11 @@ func (m *milvusRepository) ensureCollection(ctx context.Context, dimension int)
176175
}
177176

178177
// Create collection
179-
err = m.client.CreateCollection(ctx, client.NewCreateCollectionOption(collectionName, schema).WithIndexOptions(indexOpts...))
178+
createOpt := client.NewCreateCollectionOption(collectionName, schema).WithIndexOptions(indexOpts...)
179+
if m.shardsNum > 0 {
180+
createOpt = createOpt.WithShardNum(int32(m.shardsNum))
181+
}
182+
err = m.client.CreateCollection(ctx, createOpt)
180183
if err != nil {
181184
log.Errorf("[Milvus] Failed to create collection: %v", err)
182185
return fmt.Errorf("failed to create collection: %w", err)
@@ -185,7 +188,11 @@ func (m *milvusRepository) ensureCollection(ctx context.Context, dimension int)
185188
log.Infof("[Milvus] Successfully created collection %s", collectionName)
186189
}
187190

188-
loadTask, err := m.client.LoadCollection(ctx, client.NewLoadCollectionOption(collectionName))
191+
loadOpt := client.NewLoadCollectionOption(collectionName)
192+
if m.replicaNumber > 0 {
193+
loadOpt = loadOpt.WithReplica(m.replicaNumber)
194+
}
195+
loadTask, err := m.client.LoadCollection(ctx, loadOpt)
189196
if err != nil {
190197
log.Errorf("[Milvus] Failed to load collection: %v", err)
191198
return fmt.Errorf("failed to load collection: %w", err)

‎internal/application/repository/retriever/milvus/structs.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ type milvusRepository struct {
1212
client *client.Client
1313
collectionBaseName string
1414
metricType entity.MetricType
15+
shardsNum int // 0 = use Milvus default (1)
16+
replicaNumber int // 0 = use Milvus default (1); set at LoadCollection time
1517
// Cache for initialized collections (dimension -> true)
1618
initializedCollections sync.Map
1719
}

‎internal/application/repository/retriever/qdrant/repository.go‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import (
44
"context"
55
"fmt"
66
"maps"
7-
"os"
87
"slices"
98
"strings"
109
"unicode/utf8"
@@ -30,20 +29,19 @@ const (
3029
fieldIsEnabled = "is_enabled"
3130
)
3231

33-
// NewQdrantRetrieveEngineRepository creates and initializes a new Qdrant repository
34-
func NewQdrantRetrieveEngineRepository(client *qdrant.Client) interfaces.RetrieveEngineRepository {
32+
// NewQdrantRetrieveEngineRepository creates and initializes a new Qdrant repository.
33+
// indexCfg is optional — pass nil to use env var / default values (env path).
34+
func NewQdrantRetrieveEngineRepository(client *qdrant.Client, indexCfg *types.IndexConfig) interfaces.RetrieveEngineRepository {
3535
log := logger.GetLogger(context.Background())
3636
log.Info("[Qdrant] Initializing Qdrant retriever engine repository")
3737

38-
collectionBaseName := os.Getenv(envQdrantCollection)
39-
if collectionBaseName == "" {
40-
log.Warn("[Qdrant] QDRANT_COLLECTION environment variable not set, using default collection name")
41-
collectionBaseName = defaultCollectionName
42-
}
38+
collectionBaseName := types.ResolveCollectionName(indexCfg, envQdrantCollection, defaultCollectionName)
4339

4440
res := &qdrantRepository{
4541
client: client,
4642
collectionBaseName: collectionBaseName,
43+
shardNumber: indexCfg.GetShardNumber(0),
44+
replicationFactor: indexCfg.GetReplicationFactor(0),
4745
}
4846

4947
log.Info("[Qdrant] Successfully initialized repository")
@@ -82,6 +80,8 @@ func (q *qdrantRepository) ensureCollection(ctx context.Context, dimension int)
8280
Size: uint64(dimension),
8381
Distance: qdrant.Distance_Cosine,
8482
}),
83+
ShardNumber: types.OptionalUint32(q.shardNumber),
84+
ReplicationFactor: types.OptionalUint32(q.replicationFactor),
8585
})
8686
if err != nil {
8787
log.Errorf("[Qdrant] Failed to create collection: %v", err)

‎internal/application/repository/retriever/qdrant/structs.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ import (
99
type qdrantRepository struct {
1010
client *qdrant.Client
1111
collectionBaseName string
12+
shardNumber int // 0 = use Qdrant server default
13+
replicationFactor int // 0 = use Qdrant server default
1214
// Cache for initialized collections (dimension -> true)
1315
initializedCollections sync.Map
1416
}

‎internal/application/repository/retriever/weaviate/repository.go‎

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import (
44
"context"
55
"fmt"
66
"maps"
7-
"os"
87
"slices"
98
"strings"
109
"unicode/utf8"
@@ -35,19 +34,19 @@ const (
3534
fieldID = "id"
3635
)
3736

38-
func NewWeaviateRetrieveEngineRepository(client *weaviate.Client) interfaces.RetrieveEngineRepository {
37+
// NewWeaviateRetrieveEngineRepository creates and initializes a new Weaviate repository.
38+
// indexCfg is optional — pass nil to use env var / default values (env path).
39+
func NewWeaviateRetrieveEngineRepository(client *weaviate.Client, indexCfg *types.IndexConfig) interfaces.RetrieveEngineRepository {
3940
log := logger.GetLogger(context.Background())
4041
log.Info("[Weaviate] Initializing Weaviate retriever engine repository")
4142

42-
collectionBaseName := os.Getenv(envWeaviateCollection)
43-
if collectionBaseName == "" {
44-
log.Warn("[Weaviate] WEAVIATE_COLLECTION environment variable not set, using default collection name")
45-
collectionBaseName = defaultCollectionName
46-
}
43+
collectionBaseName := types.ResolveCollectionName(indexCfg, envWeaviateCollection, defaultCollectionName)
4744

4845
res := &weaviateRepository{
4946
client: client,
5047
collectionBaseName: collectionBaseName,
48+
replicationFactor: indexCfg.GetReplicationFactor(0),
49+
desiredShardCount: indexCfg.GetDesiredShardCount(0),
5150
}
5251

5352
log.Info("[Weaviate] Successfully initialized repository")
@@ -137,6 +136,18 @@ func (w *weaviateRepository) ensureCollection(ctx context.Context, dimension int
137136
},
138137
},
139138
}
139+
// Set replication factor if explicitly configured (> 0)
140+
if w.replicationFactor > 0 {
141+
classObj.ReplicationConfig = &models.ReplicationConfig{
142+
Factor: int64(w.replicationFactor),
143+
}
144+
}
145+
// Set shard count if explicitly configured (> 0)
146+
if w.desiredShardCount > 0 {
147+
classObj.ShardingConfig = map[string]interface{}{
148+
"desiredCount": w.desiredShardCount,
149+
}
150+
}
140151
//创建collection
141152
if err = w.client.Schema().ClassCreator().WithClass(&classObj).Do(ctx); err != nil {
142153
log.Errorf("[Weaviate] Failed to create collection: %v", err)

‎internal/application/repository/retriever/weaviate/structs.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ import (
99
type weaviateRepository struct {
1010
client *weaviate.Client
1111
collectionBaseName string
12+
replicationFactor int // 0 = use Weaviate server default
13+
desiredShardCount int // 0 = use Weaviate server default
1214
// Cache for initialized collections (dimension -> true)
1315
initializedCollections sync.Map
1416
}

‎internal/application/service/vectorstore.go‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package service
22

33
import (
44
"context"
5+
"fmt"
56
"os"
67
"time"
78

@@ -44,6 +45,11 @@ func (s *vectorStoreService) CreateStore(ctx context.Context, store *types.Vecto
4445
return err
4546
}
4647

48+
// 2.5. Index config validation (bounds, name characters)
49+
if err := types.ValidateIndexConfig(store.IndexConfig); err != nil {
50+
return err
51+
}
52+
4753
// 3. Duplicate check — DB stores
4854
endpoint := store.ConnectionConfig.GetEndpoint()
4955
indexName := store.IndexConfig.GetIndexNameOrDefault(store.EngineType)
@@ -66,7 +72,19 @@ func (s *vectorStoreService) CreateStore(ctx context.Context, store *types.Vecto
6672
}
6773
}
6874

69-
// 5. Persist
75+
// 5. Auto-detect server version via connection test.
76+
// This is required for engines where the version determines the SDK (e.g., ES v7 vs v8).
77+
// Without it, the wrong SDK may be used causing protocol errors (406, etc.).
78+
version, err := s.TestConnection(ctx, store.EngineType, store.ConnectionConfig)
79+
if err != nil {
80+
return errors.NewBadRequestError(
81+
fmt.Sprintf("connection test failed: %s. Ensure the server is reachable before saving.", err.Error()))
82+
}
83+
if version != "" {
84+
store.ConnectionConfig.Version = version
85+
}
86+
87+
// 6. Persist
7088
logger.Infof(ctx, "Creating vector store: tenant=%d, name=%s, engine=%s",
7189
store.TenantID, secutils.SanitizeForLog(store.Name), store.EngineType)
7290
if err := s.repo.Create(ctx, store); err != nil {

0 commit comments

Comments
 (0)