feat: add universal api embedder (#81)

* feat: add api embedder

* test: add unittest for universal api embedder

* Update src/memos/embedders/universal_api.py

Co-authored-by: Copilot <[email protected]>

* Update src/memos/configs/embedder.py

Co-authored-by: Copilot <[email protected]>

* style: reformat

* style: reformat

---------

Co-authored-by: chunyu li <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: 3 people

evaluation/scripts/locomo/openai_memory_locomo_eval_guide.md

@@ -81,7 +81,7 @@ Can you please extract relevant information from this conversation and create me
     * Click on **Manage** in the memory confirmation to view the newly generated memories.
     * Create a new local `.txt` file with the same name as the input file (e.g., `0-D1.txt`).
     * Copy each memory entry from ChatGPT and paste it into the new file, with each memory on a new line.
-5.  **Reset Memories for the Next Conversation:** 
+5.  **Reset Memories for the Next Conversation:**
     * Once all sessions for a conversation are complete, it is essential to **delete all memories to ensure a clean state for the next conversation**. Navigate to Settings -> Personalization -> Manage and click Delete all.
 
 **Example Memory Output (`0-D9.txt`):**

examples/basic_modules/embedder.py

@@ -46,3 +46,23 @@
 text_hf = "This is a sample text for Hugging Face embedding generation."
 embedding_hf = embedder_hf.embed([text_hf])
 print("Scenario 3 HF embedding shape:", len(embedding_hf[0]))
+print("==" * 20)
+
+# === Scenario 4: Using UniversalAPIEmbedder ===
+
+config_api = EmbedderConfigFactory.model_validate(
+    {
+        "backend": "universal_api",
+        "config": {
+            "provider": "openai",
+            "api_key": "<YOUR_KEY>",
+            "model_name_or_path": "text-embedding-3-large",
+            "base_url": "https://api.myproxy.com/v1",
+        },
+    }
+)
+embedder_api = EmbedderFactory.from_config(config_api)
+text_api = "This is a sample text for embedding generation using OpenAI API."
+embedding_api = embedder_api.embed([text_api])
+print("Scenario 4: OpenAI API embedding vector length:", len(embedding_api[0]))
+print("Embedding preview:", embedding_api[0][:10])

src/memos/configs/embedder.py

@@ -35,6 +35,19 @@ class SenTranEmbedderConfig(BaseEmbedderConfig):
     )
 
 
+class UniversalAPIEmbedderConfig(BaseEmbedderConfig):
+    """
+    Configuration class for universal API embedding providers, e.g.,
+    OpenAI, etc.
+    """
+
+    provider: str = Field(..., description="Provider name, e.g., 'openai'")
+    api_key: str = Field(..., description="API key for the embedding provider")
+    base_url: str | None = Field(
+        default=None, description="Optional base URL for custom or proxied endpoint"
+    )
+
+
 class EmbedderConfigFactory(BaseConfig):
     """Factory class for creating embedder configurations."""
 
@@ -45,6 +58,7 @@ class EmbedderConfigFactory(BaseConfig):
         "ollama": OllamaEmbedderConfig,
         "sentence_transformer": SenTranEmbedderConfig,
         "ark": ArkEmbedderConfig,
+        "universal_api": UniversalAPIEmbedderConfig,
     }
 
     @field_validator("backend")

src/memos/embedders/factory.py

@@ -5,6 +5,7 @@
 from memos.embedders.base import BaseEmbedder
 from memos.embedders.ollama import OllamaEmbedder
 from memos.embedders.sentence_transformer import SenTranEmbedder
+from memos.embedders.universal_api import UniversalAPIEmbedder
 
 
 class EmbedderFactory(BaseEmbedder):
@@ -14,6 +15,7 @@ class EmbedderFactory(BaseEmbedder):
         "ollama": OllamaEmbedder,
         "sentence_transformer": SenTranEmbedder,
         "ark": ArkEmbedder,
+        "universal_api": UniversalAPIEmbedder,
     }
 
     @classmethod

src/memos/embedders/universal_api.py

@@ -0,0 +1,25 @@
+from openai import OpenAI as OpenAIClient
+
+from memos.configs.embedder import UniversalAPIEmbedderConfig
+from memos.embedders.base import BaseEmbedder
+
+
+class UniversalAPIEmbedder(BaseEmbedder):
+    def __init__(self, config: UniversalAPIEmbedderConfig):
+        self.provider = config.provider
+        self.config = config
+
+        if self.provider == "openai":
+            self.client = OpenAIClient(api_key=config.api_key, base_url=config.base_url)
+        else:
+            raise ValueError(f"Unsupported provider: {self.provider}")
+
+    def embed(self, texts: list[str]) -> list[list[float]]:
+        if self.provider == "openai":
+            response = self.client.embeddings.create(
+                model=getattr(self.config, "model_name_or_path", "text-embedding-3-large"),
+                input=texts,
+            )
+            return [r.embedding for r in response.data]
+        else:
+            raise ValueError(f"Unsupported provider: {self.provider}")

tests/embedders/test_universal_api.py

@@ -0,0 +1,76 @@
+import unittest
+
+from unittest.mock import MagicMock, patch
+
+from memos.configs.embedder import UniversalAPIEmbedderConfig
+from memos.embedders.universal_api import UniversalAPIEmbedder
+
+
+class TestUniversalAPIEmbedder(unittest.TestCase):
+    @patch("memos.embedders.universal_api.OpenAIClient")
+    def test_embed_single_text(self, mock_openai_client):
+        """Test embedding a single text with OpenAI provider."""
+        # Mock the embeddings.create return value
+        mock_response = MagicMock()
+        mock_response.data = [MagicMock(embedding=[0.1, 0.2, 0.3, 0.4])]
+        mock_openai_client.return_value.embeddings.create.return_value = mock_response
+
+        config = UniversalAPIEmbedderConfig(
+            provider="openai",
+            api_key="fake-api-key",
+            base_url="https://api.openai.com/v1",
+            model_name_or_path="text-embedding-3-large",
+        )
+
+        embedder = UniversalAPIEmbedder(config)
+        text = ["Test input for embedding."]
+        result = embedder.embed(text)
+
+        # Assert OpenAIClient was created with proper args
+        mock_openai_client.assert_called_once_with(
+            api_key="fake-api-key",
+            base_url="https://api.openai.com/v1",
+        )
+
+        # Assert embeddings.create called with correct params
+        embedder.client.embeddings.create.assert_called_once_with(
+            model="text-embedding-3-large",
+            input=text,
+        )
+
+        self.assertEqual(len(result[0]), 4)
+
+    @patch("memos.embedders.universal_api.OpenAIClient")
+    def test_embed_batch_text(self, mock_openai_client):
+        """Test embedding multiple texts at once with OpenAI provider."""
+        # Mock response for multiple texts
+        mock_response = MagicMock()
+        mock_response.data = [
+            MagicMock(embedding=[0.1, 0.2]),
+            MagicMock(embedding=[0.3, 0.4]),
+            MagicMock(embedding=[0.5, 0.6]),
+        ]
+        mock_openai_client.return_value.embeddings.create.return_value = mock_response
+
+        config = UniversalAPIEmbedderConfig(
+            provider="openai",
+            api_key="fake-api-key",
+            base_url="https://api.openai.com/v1",
+            model_name_or_path="text-embedding-3-large",
+        )
+
+        embedder = UniversalAPIEmbedder(config)
+        texts = ["First text.", "Second text.", "Third text."]
+        result = embedder.embed(texts)
+
+        embedder.client.embeddings.create.assert_called_once_with(
+            model="text-embedding-3-large",
+            input=texts,
+        )
+
+        self.assertEqual(len(result), 3)
+        self.assertEqual(result[0], [0.1, 0.2])
+
+
+if __name__ == "__main__":
+    unittest.main()