Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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>

Expand Down
16 changes: 14 additions & 2 deletions openevolve/embedding.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)

Expand All @@ -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,
Expand Down Expand Up @@ -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}")

Expand Down
75 changes: 75 additions & 0 deletions tests/test_embedding.py
Original file line number Diff line number Diff line change
@@ -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()
Loading