From 149aaa5b22bb237bbf695de23507db6916cd8bc2 Mon Sep 17 00:00:00 2001 From: Enrico Zimuel Date: Mon, 10 Nov 2025 12:31:03 +0100 Subject: [PATCH] Updated Elasticsearch configurations --- src/llama_stack/distributions/starter-gpu/build.yaml | 1 + .../starter-gpu/run-with-postgres-store.yaml | 8 ++++++++ src/llama_stack/distributions/starter-gpu/run.yaml | 8 ++++++++ src/llama_stack/distributions/starter/build.yaml | 1 + .../starter/run-with-postgres-store.yaml | 8 ++++++++ src/llama_stack/distributions/starter/starter.py | 11 +++++++++++ .../remote/vector_io/elasticsearch/config.py | 8 ++++---- .../remote/vector_io/elasticsearch/elasticsearch.py | 6 ++++-- 8 files changed, 45 insertions(+), 6 deletions(-) diff --git a/src/llama_stack/distributions/starter-gpu/build.yaml b/src/llama_stack/distributions/starter-gpu/build.yaml index 10cbb1389..7b4465206 100644 --- a/src/llama_stack/distributions/starter-gpu/build.yaml +++ b/src/llama_stack/distributions/starter-gpu/build.yaml @@ -28,6 +28,7 @@ distribution_spec: - provider_type: remote::pgvector - provider_type: remote::qdrant - provider_type: remote::weaviate + - provider_type: remote::elasticsearch files: - provider_type: inline::localfs safety: diff --git a/src/llama_stack/distributions/starter-gpu/run-with-postgres-store.yaml b/src/llama_stack/distributions/starter-gpu/run-with-postgres-store.yaml index 1920ebd9d..3f06f6c85 100644 --- a/src/llama_stack/distributions/starter-gpu/run-with-postgres-store.yaml +++ b/src/llama_stack/distributions/starter-gpu/run-with-postgres-store.yaml @@ -146,6 +146,14 @@ providers: persistence: namespace: vector_io::weaviate backend: kv_default + - provider_id: ${env.ELASTICSEARCH_URL:+elasticsearch} + provider_type: remote::elasticsearch + config: + elasticsearch_url: ${env.ELASTICSEARCH_URL:=localhost:9200} + elasticsearch_api_key: ${env.ELASTICSEARCH_API_KEY:=} + persistence: + namespace: vector_io::elasticsearch + backend: kv_default files: - provider_id: meta-reference-files provider_type: inline::localfs diff --git a/src/llama_stack/distributions/starter-gpu/run.yaml b/src/llama_stack/distributions/starter-gpu/run.yaml index 7149b8659..41394e1b8 100644 --- a/src/llama_stack/distributions/starter-gpu/run.yaml +++ b/src/llama_stack/distributions/starter-gpu/run.yaml @@ -146,6 +146,14 @@ providers: persistence: namespace: vector_io::weaviate backend: kv_default + - provider_id: ${env.ELASTICSEARCH_URL:+elasticsearch} + provider_type: remote::elasticsearch + config: + elasticsearch_url: ${env.ELASTICSEARCH_URL:=localhost:9200} + elasticsearch_api_key: ${env.ELASTICSEARCH_API_KEY:=} + persistence: + namespace: vector_io::elasticsearch + backend: kv_default files: - provider_id: meta-reference-files provider_type: inline::localfs diff --git a/src/llama_stack/distributions/starter/build.yaml b/src/llama_stack/distributions/starter/build.yaml index acd51f773..35566be15 100644 --- a/src/llama_stack/distributions/starter/build.yaml +++ b/src/llama_stack/distributions/starter/build.yaml @@ -28,6 +28,7 @@ distribution_spec: - provider_type: remote::pgvector - provider_type: remote::qdrant - provider_type: remote::weaviate + - provider_type: remote::elasticsearch files: - provider_type: inline::localfs safety: diff --git a/src/llama_stack/distributions/starter/run-with-postgres-store.yaml b/src/llama_stack/distributions/starter/run-with-postgres-store.yaml index 702f95381..6283061c6 100644 --- a/src/llama_stack/distributions/starter/run-with-postgres-store.yaml +++ b/src/llama_stack/distributions/starter/run-with-postgres-store.yaml @@ -146,6 +146,14 @@ providers: persistence: namespace: vector_io::weaviate backend: kv_default + - provider_id: ${env.ELASTICSEARCH_URL:+elasticsearch} + provider_type: remote::elasticsearch + config: + elasticsearch_api_key: ${env.ELASTICSEARCH_API_KEY:=} + elasticsearch_url: ${env.ELASTICSEARCH_URL:=localhost:9200} + persistence: + namespace: vector_io::elasticsearch + backend: kv_default files: - provider_id: meta-reference-files provider_type: inline::localfs diff --git a/src/llama_stack/distributions/starter/starter.py b/src/llama_stack/distributions/starter/starter.py index 88cd3a4fe..7f47c47c9 100644 --- a/src/llama_stack/distributions/starter/starter.py +++ b/src/llama_stack/distributions/starter/starter.py @@ -41,6 +41,7 @@ from llama_stack.providers.remote.vector_io.pgvector.config import ( ) from llama_stack.providers.remote.vector_io.qdrant.config import QdrantVectorIOConfig from llama_stack.providers.remote.vector_io.weaviate.config import WeaviateVectorIOConfig +from llama_stack.providers.remote.vector_io.elasticsearch.config import ElasticsearchVectorIOConfig from llama_stack.providers.utils.kvstore.config import PostgresKVStoreConfig from llama_stack.providers.utils.sqlstore.sqlstore import PostgresSqlStoreConfig @@ -126,6 +127,7 @@ def get_distribution_template(name: str = "starter") -> DistributionTemplate: BuildProvider(provider_type="remote::pgvector"), BuildProvider(provider_type="remote::qdrant"), BuildProvider(provider_type="remote::weaviate"), + BuildProvider(provider_type="remote::elasticsearch"), ], "files": [BuildProvider(provider_type="inline::localfs")], "safety": [ @@ -240,6 +242,15 @@ def get_distribution_template(name: str = "starter") -> DistributionTemplate: cluster_url="${env.WEAVIATE_CLUSTER_URL:=}", ), ), + Provider( + provider_id="${env.ELASTICSEARCH_URL:+elasticsearch}", + provider_type="remote::elasticsearch", + config=ElasticsearchVectorIOConfig.sample_run_config( + f"~/.llama/distributions/{name}", + elasticsearch_url="${env.ELASTICSEARCH_URL:=localhost:9200}", + elasticsearch_api_key="${env.ELASTICSEARCH_API_KEY:=}", + ), + ), ], "files": [files_provider], } diff --git a/src/llama_stack/providers/remote/vector_io/elasticsearch/config.py b/src/llama_stack/providers/remote/vector_io/elasticsearch/config.py index 4f2237d79..87624b265 100644 --- a/src/llama_stack/providers/remote/vector_io/elasticsearch/config.py +++ b/src/llama_stack/providers/remote/vector_io/elasticsearch/config.py @@ -14,8 +14,8 @@ from llama_stack.schema_utils import json_schema_type @json_schema_type class ElasticsearchVectorIOConfig(BaseModel): - api_key: str | None = Field(description="The API key for the Elasticsearch instance", default=None) - hosts: str | None = Field(description="The URL of the Elasticsearch instance", default="localhost:9200") + elasticsearch_api_key: str | None = Field(description="The API key for the Elasticsearch instance", default=None) + elasticsearch_hosts: str | None = Field(description="The URL of the Elasticsearch instance", default="localhost:9200") persistence: KVStoreReference | None = Field( description="Config for KV store backend (SQLite only for now)", default=None ) @@ -23,8 +23,8 @@ class ElasticsearchVectorIOConfig(BaseModel): @classmethod def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]: return { - "hosts": "${env.ELASTICSEARCH_API_KEY:=None}", - "api_key": "${env.ELASTICSEARCH_URL:=localhost:9200}", + "elasticsearch_url": "${env.ELASTICSEARCH_URL:=localhost:9200}", + "elasticsearch_api_key": "${env.ELASTICSEARCH_API_KEY:=}", "persistence": KVStoreReference( backend="kv_default", namespace="vector_io::elasticsearch", diff --git a/src/llama_stack/providers/remote/vector_io/elasticsearch/elasticsearch.py b/src/llama_stack/providers/remote/vector_io/elasticsearch/elasticsearch.py index 2bbb167a0..8fc876b23 100644 --- a/src/llama_stack/providers/remote/vector_io/elasticsearch/elasticsearch.py +++ b/src/llama_stack/providers/remote/vector_io/elasticsearch/elasticsearch.py @@ -244,8 +244,10 @@ class ElasticsearchVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStore self.vector_store_table = None async def initialize(self) -> None: - client_config = self.config.model_dump(exclude_none=True) - self.client = AsyncElasticsearch(**client_config) + self.client = AsyncElasticsearch( + hosts=self.config.elasticsearch_url, + api_key=self.config.elasticsearch_api_key + ) self.kvstore = await kvstore_impl(self.config.persistence) start_key = VECTOR_DBS_PREFIX