Merge pull request #2559 from arc53/fix-embeds-remote

fix: remote embeds

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():

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 ---