From 38ab9ce56adb8061b99e195ea2defde7ab880a6a Mon Sep 17 00:00:00 2001
From: Anshika Pathak <pathakanshika144@gmail.com>
Date: Sat, 11 Jul 2026 21:13:19 +0800
Subject: [PATCH] feat: support Gemini embedding models

---
 README.md               |  7 ++++
 openevolve/embedding.py | 16 +++++++--
 tests/test_embedding.py | 75 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 96 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_embedding.py

diff --git a/README.md b/README.md
index 740909f5e5..a726e1168e 100644
--- a/README.md
+++ b/README.md
@@ -464,6 +464,10 @@ database:
   migration_interval: 20
   feature_dimensions: ["complexity", "diversity", "performance"]
 
+  # Optional novelty filtering with Gemini embeddings
+  embedding_model: "gemini-embedding-001"
+  similarity_threshold: 0.99
+
 evaluator:
   enable_artifacts: true      # Error feedback to LLM
   cascade_evaluation: true    # Multi-stage testing
@@ -480,6 +484,9 @@ prompt:
   use_template_stochasticity: true  # Randomized prompts
 ```
 
+For Gemini embeddings, set `GEMINI_API_KEY`. `GOOGLE_API_KEY` is also supported
+as a fallback. OpenEvolve uses Google's OpenAI-compatible endpoint automatically.
+
 <details>
 <summary><b>🎯 Feature Engineering</b></summary>
 
diff --git a/openevolve/embedding.py b/openevolve/embedding.py
index 302d4513f6..4a6bd24c90 100644
--- a/openevolve/embedding.py
+++ b/openevolve/embedding.py
@@ -3,10 +3,11 @@
 Original source: https://github.com/SakanaAI/ShinkaEvolve/blob/main/shinka/llm/embedding.py
 """
 
+import logging
 import os
+from typing import List, Union
+
 import openai
-from typing import Union, List
-import logging
 
 logger = logging.getLogger(__name__)
 
@@ -22,6 +23,10 @@
     "azure-text-embedding-3-large",
 ]
 
+GEMINI_EMBEDDING_MODELS = [
+    "gemini-embedding-001",
+]
+
 OPENAI_EMBEDDING_COSTS = {
     "text-embedding-3-small": 0.02 / M,
     "text-embedding-3-large": 0.13 / M,
@@ -53,6 +58,13 @@ def _get_client_model(self, model_name: str) -> tuple[openai.OpenAI, str]:
                 api_version=os.getenv("AZURE_API_VERSION"),
                 azure_endpoint=os.getenv("AZURE_API_ENDPOINT"),
             )
+        elif model_name in GEMINI_EMBEDDING_MODELS:
+            gemini_api_key = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
+            client = openai.OpenAI(
+                api_key=gemini_api_key,
+                base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
+            )
+            model_to_use = model_name
         else:
             raise ValueError(f"Invalid embedding model: {model_name}")
 
diff --git a/tests/test_embedding.py b/tests/test_embedding.py
new file mode 100644
index 0000000000..857ed20ff1
--- /dev/null
+++ b/tests/test_embedding.py
@@ -0,0 +1,75 @@
+import os
+import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+from openevolve.embedding import EmbeddingClient
+
+
+class TestEmbeddingClient(unittest.TestCase):
+    @patch("openevolve.embedding.openai.OpenAI")
+    @patch.dict(
+        os.environ,
+        {"GEMINI_API_KEY": "gemini-key", "GOOGLE_API_KEY": "google-key"},
+        clear=True,
+    )
+    def test_uses_gemini_key_and_openai_compatible_endpoint(self, mock_openai):
+        client = EmbeddingClient("gemini-embedding-001")
+
+        mock_openai.assert_called_once_with(
+            api_key="gemini-key",
+            base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
+        )
+        self.assertEqual(client.model, "gemini-embedding-001")
+
+    @patch("openevolve.embedding.openai.OpenAI")
+    @patch.dict(os.environ, {"GOOGLE_API_KEY": "google-key"}, clear=True)
+    def test_falls_back_to_google_api_key(self, mock_openai):
+        EmbeddingClient("gemini-embedding-001")
+
+        self.assertEqual(mock_openai.call_args.kwargs["api_key"], "google-key")
+
+    def test_rejects_unknown_embedding_model(self):
+        with self.assertRaisesRegex(ValueError, "Invalid embedding model: unknown-model"):
+            EmbeddingClient("unknown-model")
+
+    def test_get_embedding_returns_a_single_embedding(self):
+        client = EmbeddingClient.__new__(EmbeddingClient)
+        client.model = "gemini-embedding-001"
+        client.client = MagicMock()
+        client.client.embeddings.create.return_value = SimpleNamespace(
+            data=[SimpleNamespace(embedding=[0.1, 0.2])]
+        )
+
+        result = client.get_embedding("def example(): pass")
+
+        self.assertEqual(result, [0.1, 0.2])
+        client.client.embeddings.create.assert_called_once_with(
+            model="gemini-embedding-001",
+            input=["def example(): pass"],
+            encoding_format="float",
+        )
+
+    def test_get_embedding_returns_batch_embeddings(self):
+        client = EmbeddingClient.__new__(EmbeddingClient)
+        client.model = "gemini-embedding-001"
+        client.client = MagicMock()
+        client.client.embeddings.create.return_value = SimpleNamespace(
+            data=[
+                SimpleNamespace(embedding=[0.1, 0.2]),
+                SimpleNamespace(embedding=[0.3, 0.4]),
+            ]
+        )
+
+        result = client.get_embedding(["first", "second"])
+
+        self.assertEqual(result, [[0.1, 0.2], [0.3, 0.4]])
+        client.client.embeddings.create.assert_called_once_with(
+            model="gemini-embedding-001",
+            input=["first", "second"],
+            encoding_format="float",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()