feat(google): Add embedding support. (#68)

Author: daavoo

src/any_llm/providers/google/google.py

@@ -14,12 +14,18 @@
 from openai._streaming import Stream
 from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
 from openai.types.chat.chat_completion import ChatCompletion
+from openai.types import CreateEmbeddingResponse
 from any_llm.provider import Provider, ApiConfig
 from any_llm.exceptions import MissingApiKeyError, UnsupportedParameterError
 from any_llm.providers.helpers import (
     create_completion_from_response,
 )
-from any_llm.providers.google.utils import _convert_tool_spec, _convert_messages, _create_openai_chunk_from_google_chunk
+from any_llm.providers.google.utils import (
+    _convert_tool_spec,
+    _convert_messages,
+    _create_openai_chunk_from_google_chunk,
+    _create_openai_embedding_response_from_google,
+)
 
 
 class GoogleProvider(Provider):
@@ -29,7 +35,7 @@ class GoogleProvider(Provider):
     PROVIDER_DOCUMENTATION_URL = "https://cloud.google.com/vertex-ai/docs"
 
     SUPPORTS_STREAMING = True
-    SUPPORTS_EMBEDDING = False
+    SUPPORTS_EMBEDDING = True
 
     def __init__(self, config: ApiConfig) -> None:
         """Initialize Google GenAI provider."""
@@ -51,6 +57,20 @@ def __init__(self, config: ApiConfig) -> None:
 
             self.client = genai.Client(api_key=api_key)
 
+    def embedding(
+        self,
+        model: str,
+        inputs: str | list[str],
+        **kwargs: Any,
+    ) -> CreateEmbeddingResponse:
+        result = self.client.models.embed_content(
+            model=model,
+            contents=inputs,  # type: ignore[arg-type]
+            **kwargs,
+        )
+
+        return _create_openai_embedding_response_from_google(model, result)
+
     def verify_kwargs(self, kwargs: dict[str, Any]) -> None:
         """Verify the kwargs for the Google provider."""
         if kwargs.get("stream", False) and kwargs.get("response_format", None) is not None:

src/any_llm/providers/google/utils.py

@@ -2,6 +2,9 @@
 from time import time
 from typing import Any
 
+from openai.types import CreateEmbeddingResponse
+from openai.types.embedding import Embedding
+from openai.types.create_embedding_response import Usage
 from openai.types.chat.chat_completion_chunk import (
     ChatCompletionChunk,
     Choice,
@@ -92,6 +95,32 @@ def _convert_messages(messages: list[dict[str, Any]]) -> list[types.Content]:
     return formatted_messages
 
 
+def _create_openai_embedding_response_from_google(
+    model: str, result: types.EmbedContentResponse
+) -> CreateEmbeddingResponse:
+    """Convert a Google embedding response to an OpenAI-compatible format."""
+
+    data = [
+        Embedding(
+            embedding=embedding.values,
+            index=i,
+            object="embedding",
+        )
+        for i, embedding in enumerate(result.embeddings or [])
+        if embedding.values
+    ]
+
+    # Google does not provide usage data in the embedding response
+    usage = Usage(prompt_tokens=0, total_tokens=0)
+
+    return CreateEmbeddingResponse(
+        data=data,
+        model=model,
+        object="list",
+        usage=usage,
+    )
+
+
 def _create_openai_chunk_from_google_chunk(
     response: types.GenerateContentResponse,
 ) -> ChatCompletionChunk:

tests/conftest.py

@@ -46,6 +46,7 @@ def embedding_provider_model_map() -> dict[ProviderName, str]:
         ProviderName.AWS: "amazon.titan-embed-text-v2:0",
         ProviderName.OLLAMA: "llama3.2:1b",
         ProviderName.LMSTUDIO: "text-embedding-nomic-embed-text-v1.5",
+        ProviderName.GOOGLE: "gemini-embedding-001",
     }
 
 

tests/integration/test_embedding.py

@@ -27,11 +27,11 @@ def test_embedding_providers(provider: ProviderName, embedding_provider_model_ma
         if "model" in str(e).lower() or "embedding" in str(e).lower():
             pytest.skip(f"{provider.value} embedding model not available: {e}")
         raise
-    # Verify result is a list of floats
     assert isinstance(result, CreateEmbeddingResponse)
     assert len(result.data) > 0
-    assert all(isinstance(x.embedding, list) for x in result.data)
-    # LM Studio follows OpenAI Spec but doesn't output token use
-    if provider not in ProviderName.LMSTUDIO:
+    for entry in result.data:
+        assert all(isinstance(v, float) for v in entry.embedding)
+    # These providers don't output token use
+    if provider not in (ProviderName.GOOGLE, ProviderName.LMSTUDIO):
         assert result.usage.prompt_tokens > 0
         assert result.usage.total_tokens > 0