From df5582db37e1483139b5d847c1ad1d1b7f25449f Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 23 Jun 2026 21:43:11 +0100 Subject: [PATCH] fix: remote embeds --- application/vectorstore/base.py | 39 ++++++++++++++++++++++++++------- tests/vectorstore/test_base.py | 25 +++++++++++++++++++++ 2 files changed, 56 insertions(+), 8 deletions(-) diff --git a/application/vectorstore/base.py b/application/vectorstore/base.py index c3fca698..cd7a0e6b 100644 --- a/application/vectorstore/base.py +++ b/application/vectorstore/base.py @@ -93,8 +93,38 @@ def _get_embeddings_wrapper(): class EmbeddingsSingleton: _instances = {} + @staticmethod + def _remote_instance(embeddings_name, embeddings_key=None): + """Return a cached ``RemoteEmbeddings`` for the configured remote API. + + Centralizes the ``EMBEDDINGS_BASE_URL`` dispatch so every caller — + including code that calls :meth:`get_instance` directly (GraphRAG, + semantic chunking) rather than via + :meth:`BaseVectorStore._get_embeddings` — routes to the remote + embeddings server instead of attempting a local model download. + + Args: + embeddings_name: Model name forwarded to the remote API. + embeddings_key: Optional API key; falls back to + ``settings.EMBEDDINGS_KEY`` when not provided. + + Returns: + RemoteEmbeddings: Shared instance keyed by base URL and model name. + """ + api_key = embeddings_key if embeddings_key is not None else settings.EMBEDDINGS_KEY + cache_key = f"remote_{settings.EMBEDDINGS_BASE_URL}_{embeddings_name}" + if cache_key not in EmbeddingsSingleton._instances: + EmbeddingsSingleton._instances[cache_key] = RemoteEmbeddings( + api_url=settings.EMBEDDINGS_BASE_URL, + model_name=embeddings_name, + api_key=api_key, + ) + return EmbeddingsSingleton._instances[cache_key] + @staticmethod def get_instance(embeddings_name, *args, **kwargs): + if settings.EMBEDDINGS_BASE_URL: + return EmbeddingsSingleton._remote_instance(embeddings_name) if embeddings_name not in EmbeddingsSingleton._instances: EmbeddingsSingleton._instances[embeddings_name] = ( EmbeddingsSingleton._create_instance(embeddings_name, *args, **kwargs) @@ -197,14 +227,7 @@ class BaseVectorStore(ABC): logging.info( f"Using remote embeddings API at: {settings.EMBEDDINGS_BASE_URL}" ) - cache_key = f"remote_{settings.EMBEDDINGS_BASE_URL}_{embeddings_name}" - if cache_key not in EmbeddingsSingleton._instances: - EmbeddingsSingleton._instances[cache_key] = RemoteEmbeddings( - api_url=settings.EMBEDDINGS_BASE_URL, - model_name=embeddings_name, - api_key=embeddings_key, - ) - return EmbeddingsSingleton._instances[cache_key] + return EmbeddingsSingleton._remote_instance(embeddings_name, embeddings_key) if embeddings_name == "openai_text-embedding-ada-002": if self.is_azure_configured(): diff --git a/tests/vectorstore/test_base.py b/tests/vectorstore/test_base.py index 4e8aa637..7e0346aa 100644 --- a/tests/vectorstore/test_base.py +++ b/tests/vectorstore/test_base.py @@ -228,6 +228,31 @@ class TestEmbeddingsSingleton: mock_wrapper_cls.assert_called_once_with("custom_model_name") assert result is mock_instance + @patch("application.vectorstore.base.settings") + def test_get_instance_uses_remote_when_base_url_set(self, mock_settings): + """Direct callers (GraphRAG, semantic chunking) must route to the + remote embeddings API instead of loading a local model.""" + mock_settings.EMBEDDINGS_BASE_URL = "http://remote:8080" + mock_settings.EMBEDDINGS_KEY = "sk-remote" + + result = EmbeddingsSingleton.get_instance("embeddinggemma", "sk-remote") + + assert isinstance(result, RemoteEmbeddings) + assert result.api_url == "http://remote:8080" + assert result.model_name == "embeddinggemma" + assert result.headers["Authorization"] == "Bearer sk-remote" + + @patch("application.vectorstore.base.settings") + def test_get_instance_remote_falls_back_to_settings_key(self, mock_settings): + """When no key is passed, the remote dispatch uses EMBEDDINGS_KEY.""" + mock_settings.EMBEDDINGS_BASE_URL = "http://remote:8080" + mock_settings.EMBEDDINGS_KEY = "sk-from-settings" + + result = EmbeddingsSingleton.get_instance("embeddinggemma") + + assert isinstance(result, RemoteEmbeddings) + assert result.headers["Authorization"] == "Bearer sk-from-settings" + # --- BaseVectorStore ---