feat: configure vector-io provider with an embedding model

Signed-off-by: Mustafa Elbehery <melbeher@redhat.com>
This commit is contained in:
Mustafa Elbehery 2025-07-31 13:07:03 +02:00
parent 1f0766308d
commit d8f013b35a
29 changed files with 228 additions and 24 deletions

View file

@ -9,6 +9,7 @@ from typing import Any
from pydantic import BaseModel, Field
from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig
from llama_stack.providers.utils.vector_io.embedding_config import EmbeddingConfig
from llama_stack.schema_utils import json_schema_type
@ -16,6 +17,10 @@ from llama_stack.schema_utils import json_schema_type
class ChromaVectorIOConfig(BaseModel):
db_path: str
kvstore: KVStoreConfig = Field(description="Config for KV store backend")
embedding: EmbeddingConfig | None = Field(
default=None,
description="Default embedding configuration for this provider. When specified, vector databases created with this provider will use these embedding settings as defaults.",
)
@classmethod
def sample_run_config(
@ -27,4 +32,9 @@ class ChromaVectorIOConfig(BaseModel):
__distro_dir__=__distro_dir__,
db_name="chroma_inline_registry.db",
),
# Optional: Configure default embedding model for this provider
# "embedding": {
# "model": "${env.CHROMA_EMBEDDING_MODEL:=all-MiniLM-L6-v2}",
# "dimensions": 384
# },
}

View file

@ -6,18 +6,23 @@
from typing import Any
from pydantic import BaseModel
from pydantic import BaseModel, Field
from llama_stack.providers.utils.kvstore.config import (
KVStoreConfig,
SqliteKVStoreConfig,
)
from llama_stack.providers.utils.vector_io.embedding_config import EmbeddingConfig
from llama_stack.schema_utils import json_schema_type
@json_schema_type
class FaissVectorIOConfig(BaseModel):
kvstore: KVStoreConfig
embedding: EmbeddingConfig | None = Field(
default=None,
description="Default embedding configuration for this provider. When specified, vector databases created with this provider will use these embedding settings as defaults.",
)
@classmethod
def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]:
@ -25,5 +30,10 @@ class FaissVectorIOConfig(BaseModel):
"kvstore": SqliteKVStoreConfig.sample_run_config(
__distro_dir__=__distro_dir__,
db_name="faiss_store.db",
)
),
# Optional: Configure default embedding model for this provider
# "embedding": {
# "model": "${env.FAISS_EMBEDDING_MODEL:=all-MiniLM-L6-v2}",
# "dimensions": 384
# },
}

View file

@ -35,6 +35,7 @@ from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIV
from llama_stack.providers.utils.memory.vector_store import (
EmbeddingIndex,
VectorDBWithIndex,
apply_provider_embedding_defaults,
)
from .config import FaissVectorIOConfig
@ -237,6 +238,9 @@ class FaissVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolPr
self,
vector_db: VectorDB,
) -> None:
# Apply provider-level embedding defaults if configured
vector_db = apply_provider_embedding_defaults(vector_db, self.config.embedding)
assert self.kvstore is not None
key = f"{VECTOR_DBS_PREFIX}{vector_db.identifier}"

View file

@ -12,6 +12,7 @@ from llama_stack.providers.utils.kvstore.config import (
KVStoreConfig,
SqliteKVStoreConfig,
)
from llama_stack.providers.utils.vector_io.embedding_config import EmbeddingConfig
from llama_stack.schema_utils import json_schema_type
@ -20,6 +21,10 @@ class MilvusVectorIOConfig(BaseModel):
db_path: str
kvstore: KVStoreConfig = Field(description="Config for KV store backend (SQLite only for now)")
consistency_level: str = Field(description="The consistency level of the Milvus server", default="Strong")
embedding: EmbeddingConfig | None = Field(
default=None,
description="Default embedding configuration for this provider. When specified, vector databases created with this provider will use these embedding settings as defaults.",
)
@classmethod
def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]:
@ -29,4 +34,9 @@ class MilvusVectorIOConfig(BaseModel):
__distro_dir__=__distro_dir__,
db_name="milvus_registry.db",
),
# Optional: Configure default embedding model for this provider
# "embedding": {
# "model": "${env.MILVUS_EMBEDDING_MODEL:=all-MiniLM-L6-v2}",
# "dimensions": 384
# },
}

View file

@ -12,11 +12,18 @@ from llama_stack.providers.utils.kvstore.config import (
KVStoreConfig,
SqliteKVStoreConfig,
)
from llama_stack.providers.utils.vector_io.embedding_config import EmbeddingConfig
from llama_stack.schema_utils import json_schema_type
@json_schema_type
class SQLiteVectorIOConfig(BaseModel):
db_path: str = Field(description="Path to the SQLite database file")
kvstore: KVStoreConfig = Field(description="Config for KV store backend (SQLite only for now)")
embedding: EmbeddingConfig | None = Field(
default=None,
description="Default embedding configuration for this provider. When specified, vector databases created with this provider will use these embedding settings as defaults.",
)
@classmethod
def sample_run_config(cls, __distro_dir__: str) -> dict[str, Any]:
@ -26,4 +33,9 @@ class SQLiteVectorIOConfig(BaseModel):
__distro_dir__=__distro_dir__,
db_name="sqlite_vec_registry.db",
),
# Optional: Configure default embedding model for this provider
# "embedding": {
# "model": "${env.SQLITE_VEC_EMBEDDING_MODEL:=all-MiniLM-L6-v2}",
# "dimensions": 384
# },
}