Skip to content

Commit 6a0fbf3

Browse files
jharan1coderabbitai[bot]
jharan1
and
coderabbitai[bot]
authored
feat: RHOAIENG-51773 add external vector stores support in lsd install (gen-ai bff) (#6812)
* add external vector stores support in lsd install (gen-ai bff) * Update packages/gen-ai/bff/internal/integrations/kubernetes/k8smocks/token_k8s_client_mock.go Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> * update InstallLlamaStackDistribution to fail fast if credential secret doesn't exist * add check that provider id doesn't collide with default provided id * update generateLlamaStackConfig to ensure persistence is present across all providers * use models.ModelTypeEmbedding * remove default adding of token key for milvus (instead docs will require it added in configmap) * remove use of credentialEnvVarField * update env var naming for vs credentials to include segment with sanitized vector-io ProviderID * remove "milvus without credential gets empty token field" test * add buildEmbeddingModelLookup and use to derive EmbeddingModel set in llamastack registered vector stores * update comment * add vector_stores to LlamaStackDistributionInstallRequest in gen-ai.yaml --------- Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
1 parent 8eb809e commit 6a0fbf3

19 files changed

Lines changed: 1380 additions & 124 deletions

‎packages/gen-ai/bff/internal/api/aaa_vectorstores_handler_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ var _ = Describe("VectorStoresAAHandler", func() {
7171
err = json.Unmarshal(body, &response)
7272
require.NoError(t, err)
7373

74-
assert.Len(t, response.Data, 3, "Should return 3 vector stores from the ConfigMap")
74+
assert.Len(t, response.Data, 4, "Should return 4 vector stores from the ConfigMap")
7575

7676
for _, store := range response.Data {
7777
assert.NotEmpty(t, store.VectorStoreID, "Store should have a vector_store_id")
@@ -108,7 +108,7 @@ var _ = Describe("VectorStoresAAHandler", func() {
108108
err = json.Unmarshal(body, &response)
109109
require.NoError(t, err)
110110

111-
assert.Len(t, response.Data, 3, "mock-test-namespace-1 should have the same 3 vector stores from its ConfigMap")
111+
assert.Len(t, response.Data, 4, "mock-test-namespace-1 should have the same 4 vector stores from its ConfigMap")
112112
})
113113

114114
It("should not include ConfigMap metadata in the response", func() {

‎packages/gen-ai/bff/internal/api/external_vectorstores_handler_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -72,8 +72,8 @@ var _ = Describe("ExternalVectorStoresListHandler", func() {
7272
err = json.Unmarshal(body, &response)
7373
require.NoError(t, err)
7474

75-
assert.Equal(t, 3, response.Data.TotalCount, "Should return 3 vector stores")
76-
assert.Equal(t, 3, len(response.Data.VectorStores))
75+
assert.Equal(t, 4, response.Data.TotalCount, "Should return 4 vector stores")
76+
assert.Equal(t, 4, len(response.Data.VectorStores))
7777

7878
for _, store := range response.Data.VectorStores {
7979
assert.NotEmpty(t, store.VectorStoreID, "Store should have a vector_store_id")

‎packages/gen-ai/bff/internal/api/lsd_handler_test.go‎

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -671,6 +671,167 @@ var _ = Describe("LlamaStackDistributionInstallHandlerWithMaaSModels", func() {
671671
})
672672
})
673673

674+
var _ = Describe("LlamaStackDistributionInstallHandlerWithVectorStores", func() {
675+
var app App
676+
677+
BeforeEach(func() {
678+
k8sFactory, err := k8smocks.NewTokenClientFactory(testK8sClient, testCfg, slog.Default())
679+
require.NoError(GinkgoT(), err)
680+
681+
app = App{
682+
config: config.EnvConfig{
683+
Port: 4000,
684+
},
685+
logger: slog.Default(),
686+
kubernetesClientFactory: k8sFactory,
687+
repositories: repositories.NewRepositories(),
688+
}
689+
})
690+
691+
It("should install LSD successfully with empty vector_stores", func() {
692+
t := GinkgoT()
693+
namespace := fmt.Sprintf("vs-empty-test-%d", time.Now().UnixNano())
694+
ctx := context.Background()
695+
696+
requestBody := map[string]interface{}{
697+
"models": []map[string]interface{}{
698+
{"model_name": "llama-2-7b-chat", "model_source_type": "maas"},
699+
},
700+
"vector_stores": []interface{}{},
701+
}
702+
jsonBody, err := json.Marshal(requestBody)
703+
require.NoError(t, err)
704+
705+
req, err := http.NewRequest(http.MethodPost, "/gen-ai/api/v1/llamastack-distribution/install", bytes.NewReader(jsonBody))
706+
require.NoError(t, err)
707+
req.Header.Set("Content-Type", "application/json")
708+
709+
ctx = context.WithValue(ctx, constants.NamespaceQueryParameterKey, namespace)
710+
ctx = context.WithValue(ctx, constants.RequestIdentityKey, &integrations.RequestIdentity{Token: "FAKE_BEARER_TOKEN"})
711+
ctx = context.WithValue(ctx, constants.MaaSClientKey, maasmocks.NewMockMaaSClient())
712+
req = req.WithContext(ctx)
713+
714+
rr := httptest.NewRecorder()
715+
app.LlamaStackDistributionInstallHandler(rr, req, nil)
716+
717+
assert.Equal(t, http.StatusOK, rr.Code)
718+
var response map[string]interface{}
719+
err = json.Unmarshal(rr.Body.Bytes(), &response)
720+
assert.NoError(t, err)
721+
_, hasData := response["data"]
722+
assert.True(t, hasData)
723+
})
724+
725+
It("should install LSD with a valid vector store", func() {
726+
t := GinkgoT()
727+
// mock-test-namespace-1 has gen-ai-aa-vector-stores ConfigMap pre-seeded.
728+
namespace := "mock-test-namespace-1"
729+
ctx := context.Background()
730+
cleanupTestNamespace(ctx, namespace)
731+
732+
requestBody := map[string]interface{}{
733+
"models": []map[string]interface{}{
734+
{"model_name": "llama-2-7b-chat", "model_source_type": "maas"},
735+
},
736+
"vector_stores": []map[string]interface{}{
737+
{"vector_store_id": "vs_4c4b74e3-30ac-4e46-9057-213154f83dba"},
738+
},
739+
}
740+
jsonBody, err := json.Marshal(requestBody)
741+
require.NoError(t, err)
742+
743+
req, err := http.NewRequest(http.MethodPost, "/gen-ai/api/v1/llamastack-distribution/install", bytes.NewReader(jsonBody))
744+
require.NoError(t, err)
745+
req.Header.Set("Content-Type", "application/json")
746+
747+
ctx = context.WithValue(ctx, constants.NamespaceQueryParameterKey, namespace)
748+
ctx = context.WithValue(ctx, constants.RequestIdentityKey, &integrations.RequestIdentity{Token: "FAKE_BEARER_TOKEN"})
749+
ctx = context.WithValue(ctx, constants.MaaSClientKey, maasmocks.NewMockMaaSClient())
750+
req = req.WithContext(ctx)
751+
752+
rr := httptest.NewRecorder()
753+
app.LlamaStackDistributionInstallHandler(rr, req, nil)
754+
755+
body, _ := io.ReadAll(rr.Result().Body)
756+
t.Logf("Response: %s", string(body))
757+
assert.Equal(t, http.StatusOK, rr.Code)
758+
})
759+
760+
It("should return 400 when vector_store_id is not found in ConfigMap", func() {
761+
t := GinkgoT()
762+
namespace := "mock-test-namespace-1"
763+
ctx := context.Background()
764+
cleanupTestNamespace(ctx, namespace)
765+
766+
requestBody := map[string]interface{}{
767+
"models": []map[string]interface{}{
768+
{"model_name": "llama-2-7b-chat", "model_source_type": "maas"},
769+
},
770+
"vector_stores": []map[string]interface{}{
771+
{"vector_store_id": "nonexistent-store"},
772+
},
773+
}
774+
jsonBody, err := json.Marshal(requestBody)
775+
require.NoError(t, err)
776+
777+
req, err := http.NewRequest(http.MethodPost, "/gen-ai/api/v1/llamastack-distribution/install", bytes.NewReader(jsonBody))
778+
require.NoError(t, err)
779+
req.Header.Set("Content-Type", "application/json")
780+
781+
ctx = context.WithValue(ctx, constants.NamespaceQueryParameterKey, namespace)
782+
ctx = context.WithValue(ctx, constants.RequestIdentityKey, &integrations.RequestIdentity{Token: "FAKE_BEARER_TOKEN"})
783+
ctx = context.WithValue(ctx, constants.MaaSClientKey, maasmocks.NewMockMaaSClient())
784+
req = req.WithContext(ctx)
785+
786+
rr := httptest.NewRecorder()
787+
app.LlamaStackDistributionInstallHandler(rr, req, nil)
788+
789+
assert.Equal(t, http.StatusBadRequest, rr.Code)
790+
var response map[string]interface{}
791+
err = json.Unmarshal(rr.Body.Bytes(), &response)
792+
assert.NoError(t, err)
793+
errorMap, _ := response["error"].(map[string]interface{})
794+
assert.Contains(t, errorMap["message"], "nonexistent-store")
795+
})
796+
797+
It("should return 400 when vector stores ConfigMap is absent from the namespace", func() {
798+
t := GinkgoT()
799+
// Use a fresh namespace that has no pre-seeded ConfigMap.
800+
namespace := fmt.Sprintf("vs-no-configmap-%d", time.Now().UnixNano())
801+
ctx := context.Background()
802+
803+
requestBody := map[string]interface{}{
804+
"models": []map[string]interface{}{
805+
{"model_name": "llama-2-7b-chat", "model_source_type": "maas"},
806+
},
807+
"vector_stores": []map[string]interface{}{
808+
{"vector_store_id": "vs_4c4b74e3-30ac-4e46-9057-213154f83dba"},
809+
},
810+
}
811+
jsonBody, err := json.Marshal(requestBody)
812+
require.NoError(t, err)
813+
814+
req, err := http.NewRequest(http.MethodPost, "/gen-ai/api/v1/llamastack-distribution/install", bytes.NewReader(jsonBody))
815+
require.NoError(t, err)
816+
req.Header.Set("Content-Type", "application/json")
817+
818+
ctx = context.WithValue(ctx, constants.NamespaceQueryParameterKey, namespace)
819+
ctx = context.WithValue(ctx, constants.RequestIdentityKey, &integrations.RequestIdentity{Token: "FAKE_BEARER_TOKEN"})
820+
ctx = context.WithValue(ctx, constants.MaaSClientKey, maasmocks.NewMockMaaSClient())
821+
req = req.WithContext(ctx)
822+
823+
rr := httptest.NewRecorder()
824+
app.LlamaStackDistributionInstallHandler(rr, req, nil)
825+
826+
assert.Equal(t, http.StatusBadRequest, rr.Code)
827+
var response map[string]interface{}
828+
err = json.Unmarshal(rr.Body.Bytes(), &response)
829+
assert.NoError(t, err)
830+
errorMap, _ := response["error"].(map[string]interface{})
831+
assert.Contains(t, errorMap["message"], "gen-ai-aa-vector-stores")
832+
})
833+
})
834+
674835
var _ = Describe("LlamaStackDistributionDeleteHandler", func() {
675836
var app App
676837

‎packages/gen-ai/bff/internal/api/lsd_install_handler.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@ func (app *App) LlamaStackDistributionInstallHandler(w http.ResponseWriter, r *h
8484

8585
// Pass the InstallModel structs directly to the repository
8686
// enableGuardrails - if true, safety providers with shields will be configured for all models
87-
response, err := app.repositories.LlamaStackDistribution.InstallLlamaStackDistribution(client, ctx, identity, namespace, installRequest.Models, installRequest.EnableGuardrails, maasClient)
87+
response, err := app.repositories.LlamaStackDistribution.InstallLlamaStackDistribution(client, ctx, identity, namespace, installRequest.Models, installRequest.EnableGuardrails, installRequest.VectorStores, maasClient)
8888
if err != nil {
8989
app.badRequestResponse(w, r, err)
9090
return

‎packages/gen-ai/bff/internal/integrations/kubernetes/client.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ type KubernetesClientInterface interface {
3333
// LlamaStack Distribution
3434
GetLlamaStackDistributions(ctx context.Context, identity *integrations.RequestIdentity, namespace string) (*lsdapi.LlamaStackDistributionList, error)
3535
CanListLlamaStackDistributions(ctx context.Context, identity *integrations.RequestIdentity, namespace string) (bool, error)
36-
InstallLlamaStackDistribution(ctx context.Context, identity *integrations.RequestIdentity, namespace string, models []models.InstallModel, enableGuardrails bool, maasClient maas.MaaSClientInterface) (*lsdapi.LlamaStackDistribution, error)
36+
InstallLlamaStackDistribution(ctx context.Context, identity *integrations.RequestIdentity, namespace string, installModels []models.InstallModel, enableGuardrails bool, vectorStores []models.InstallVectorStore, maasClient maas.MaaSClientInterface) (*lsdapi.LlamaStackDistribution, error)
3737
DeleteLlamaStackDistribution(ctx context.Context, identity *integrations.RequestIdentity, namespace string, name string) (*lsdapi.LlamaStackDistribution, error)
3838
GetModelProviderInfo(ctx context.Context, identity *integrations.RequestIdentity, namespace string, modelID string) (*types.ModelProviderInfo, error)
3939

‎packages/gen-ai/bff/internal/integrations/kubernetes/k8smocks/base_testenv.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -458,7 +458,7 @@ registered_resources:
458458
provider_id: vllm-inference-1
459459
model_type: llm
460460
shields: []
461-
vector_dbs: []
461+
vector_stores: []
462462
datasets: []
463463
scoring_fns: []
464464
benchmarks: []

‎packages/gen-ai/bff/internal/integrations/kubernetes/k8smocks/testdata/vector_stores.yaml‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,16 @@ providers:
3939
config:
4040
uri: http://milvus.vector.svc.cluster.local:19530
4141

42+
- provider_id: milvus-secure-provider
43+
provider_type: remote::milvus
44+
config:
45+
uri: http://milvus-secure.vector.svc.cluster.local:19530
46+
custom_gen_ai:
47+
credentials:
48+
secretRefs:
49+
- name: milvus-credentials
50+
key: token
51+
4252
registered_resources:
4353
vector_stores:
4454
- provider_id: pgvector-provider
@@ -64,3 +74,11 @@ registered_resources:
6474
vector_store_name: "Code Vector Store (Milvus)"
6575
metadata:
6676
description: "Code embeddings for repository search"
77+
78+
- provider_id: milvus-secure-provider
79+
vector_store_id: vs_c9e12f45-3b67-4d89-a012-3b4567890cde
80+
embedding_model: ibm-granite/granite-embedding-125m-english
81+
embedding_dimension: 768
82+
vector_store_name: "Secure Knowledge Base (Milvus)"
83+
metadata:
84+
description: "Secure internal knowledge base with token auth"

‎packages/gen-ai/bff/internal/integrations/kubernetes/k8smocks/token_k8s_client_mock.go‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -384,7 +384,13 @@ func (m *TokenKubernetesClientMock) GetLlamaStackDistributions(ctx context.Conte
384384
}, nil
385385
}
386386

387-
func (m *TokenKubernetesClientMock) InstallLlamaStackDistribution(ctx context.Context, identity *integrations.RequestIdentity, namespace string, installModels []models.InstallModel, enableGuardrails bool, maasClient maas.MaaSClientInterface) (*lsdapi.LlamaStackDistribution, error) {
387+
func (m *TokenKubernetesClientMock) InstallLlamaStackDistribution(ctx context.Context, identity *integrations.RequestIdentity, namespace string, installModels []models.InstallModel, enableGuardrails bool, vectorStores []models.InstallVectorStore, maasClient maas.MaaSClientInterface) (*lsdapi.LlamaStackDistribution, error) {
388+
if len(vectorStores) > 0 {
389+
if _, err := m.LoadAndValidateVectorStores(ctx, identity, namespace, vectorStores); err != nil {
390+
return nil, err
391+
}
392+
}
393+
388394
// Check if LSD already exists in the namespace
389395
existingLSDList, err := m.GetLlamaStackDistributions(ctx, identity, namespace)
390396
if err != nil {
@@ -604,7 +610,7 @@ registered_resources:
604610
provider_id: vllm-inference-1
605611
model_type: llm
606612
` + shieldsSection + `
607-
vector_dbs: []
613+
vector_stores: []
608614
datasets: []
609615
scoring_fns: []
610616
benchmarks: []

‎packages/gen-ai/bff/internal/integrations/kubernetes/llamastack_config.go‎

Lines changed: 35 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -62,13 +62,13 @@ type MetadataStore struct {
6262
}
6363

6464
type RegisteredResources struct {
65-
Models []Model `json:"models" yaml:"models"`
66-
Shields []Shield `json:"shields" yaml:"shields"`
67-
VectorDBs []VectorDB `json:"vector_dbs" yaml:"vector_dbs"`
68-
Datasets []Dataset `json:"datasets" yaml:"datasets"`
69-
ScoringFns []ScoringFn `json:"scoring_fns" yaml:"scoring_fns"`
70-
Benchmarks []Benchmark `json:"benchmarks" yaml:"benchmarks"`
71-
ToolGroups []ToolGroup `json:"tool_groups" yaml:"tool_groups"`
65+
Models []Model `json:"models" yaml:"models"`
66+
Shields []Shield `json:"shields" yaml:"shields"`
67+
VectorStores []VectorStore `json:"vector_stores" yaml:"vector_stores"`
68+
Datasets []Dataset `json:"datasets" yaml:"datasets"`
69+
ScoringFns []ScoringFn `json:"scoring_fns" yaml:"scoring_fns"`
70+
Benchmarks []Benchmark `json:"benchmarks" yaml:"benchmarks"`
71+
ToolGroups []ToolGroup `json:"tool_groups" yaml:"tool_groups"`
7272
}
7373

7474
type Storage struct {
@@ -229,12 +229,12 @@ func NewDefaultLlamaStackConfig() *LlamaStackConfig {
229229
},
230230
RegisteredResources: RegisteredResources{
231231
// Ensure these serialize as `[]` (not `null`) when no values exist.
232-
Models: []Model{},
233-
Shields: []Shield{},
234-
VectorDBs: []VectorDB{},
235-
Datasets: []Dataset{},
236-
ScoringFns: []ScoringFn{},
237-
Benchmarks: []Benchmark{},
232+
Models: []Model{},
233+
Shields: []Shield{},
234+
VectorStores: []VectorStore{},
235+
Datasets: []Dataset{},
236+
ScoringFns: []ScoringFn{},
237+
Benchmarks: []Benchmark{},
238238
ToolGroups: []ToolGroup{
239239
{
240240
ToolGroupID: "builtin::rag",
@@ -351,7 +351,7 @@ func NewEmbeddingModel(modelID, providerID, providerModelID string, embeddingDim
351351
ModelID: modelID,
352352
ProviderID: providerID,
353353
ProviderModelID: providerModelID,
354-
ModelType: "embedding",
354+
ModelType: string(models.ModelTypeEmbedding),
355355
Metadata: map[string]interface{}{
356356
"embedding_dimension": embeddingDimension,
357357
},
@@ -363,7 +363,7 @@ func NewLLMModel(modelID, providerID string, displayName string) Model {
363363
return Model{
364364
ModelID: modelID,
365365
ProviderID: providerID,
366-
ModelType: "llm",
366+
ModelType: string(models.ModelTypeLLM),
367367
Metadata: map[string]interface{}{
368368
"display_name": displayName,
369369
},
@@ -545,6 +545,11 @@ func (c *LlamaStackConfig) RegisterShield(shield Shield) {
545545
c.RegisteredResources.Shields = append(c.RegisteredResources.Shields, shield)
546546
}
547547

548+
// RegisterVectorStore adds a vector store to the registered resources.
549+
func (c *LlamaStackConfig) RegisterVectorStore(store VectorStore) {
550+
c.RegisteredResources.VectorStores = append(c.RegisteredResources.VectorStores, store)
551+
}
552+
548553
// GetModelProviderInfo extracts model provider information for a given model ID
549554
// This is a two-step process:
550555
// 1. Find the model in the Models section and get its provider_id
@@ -631,13 +636,15 @@ type Shield struct {
631636
Metadata map[string]interface{} `json:"metadata,omitempty" yaml:"metadata,omitempty"`
632637
}
633638

634-
// VectorDB represents a vector database configuration
635-
type VectorDB struct {
636-
DBID string `json:"db_id" yaml:"db_id"`
637-
Name string `json:"name" yaml:"name"`
638-
ProviderID string `json:"provider_id" yaml:"provider_id"`
639-
Config map[string]interface{} `json:"config" yaml:"config"`
640-
Metadata map[string]interface{} `json:"metadata,omitempty" yaml:"metadata,omitempty"`
639+
// VectorStore represents a vector store configuration in registered_resources
640+
type VectorStore struct {
641+
VectorStoreID string `json:"vector_store_id" yaml:"vector_store_id"`
642+
EmbeddingModel string `json:"embedding_model" yaml:"embedding_model"`
643+
EmbeddingDimension int `json:"embedding_dimension" yaml:"embedding_dimension"`
644+
ProviderID string `json:"provider_id,omitempty" yaml:"provider_id,omitempty"`
645+
ProviderVectorStoreID string `json:"provider_vector_store_id,omitempty" yaml:"provider_vector_store_id,omitempty"`
646+
VectorStoreName string `json:"vector_store_name,omitempty" yaml:"vector_store_name,omitempty"`
647+
Metadata map[string]interface{} `json:"metadata,omitempty" yaml:"metadata,omitempty"`
641648
}
642649

643650
// Dataset represents a dataset configuration
@@ -680,14 +687,12 @@ func NewShield(shieldID, shieldType, providerID string, config map[string]interf
680687
}
681688
}
682689

683-
// NewVectorDB creates a new VectorDB instance
684-
func NewVectorDB(dbID, name, providerID string, config map[string]interface{}) VectorDB {
685-
return VectorDB{
686-
DBID: dbID,
687-
Name: name,
688-
ProviderID: providerID,
689-
Config: config,
690-
Metadata: EmptyConfig(),
690+
// NewVectorStore creates a new VectorStore instance
691+
func NewVectorStore(vectorStoreID, embeddingModel string, embeddingDimension int) VectorStore {
692+
return VectorStore{
693+
VectorStoreID: vectorStoreID,
694+
EmbeddingModel: embeddingModel,
695+
EmbeddingDimension: embeddingDimension,
691696
}
692697
}
693698

0 commit comments

Comments
 (0)