fix return type for GoogleGenerativeAIEmbeddings (langchain-ai#248)

* fix return type for GoogleGenerativeAIEmbeddings
eliasecchig · May 30, 2024 · d043e4f · d043e4f
1 parent ae72db3
commit d043e4f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/libs/genai/langchain_google_genai/embeddings.py b/libs/genai/langchain_google_genai/embeddings.py
@@ -148,7 +148,7 @@ def embed_documents(
             )
         except Exception as e:
             raise GoogleGenerativeAIError(f"Error embedding content: {e}") from e
-        return [e.values for e in result.embeddings]
+        return [list(e.values) for e in result.embeddings]
 
     def embed_query(
         self,

diff --git a/libs/genai/tests/integration_tests/test_embeddings.py b/libs/genai/tests/integration_tests/test_embeddings.py
@@ -20,6 +20,7 @@ def test_embed_query_different_lengths(query: str) -> None:
     model = GoogleGenerativeAIEmbeddings(model=_MODEL)
     result = model.embed_query(query)
     assert len(result) == 768
+    assert isinstance(result, list)
 
 
 @pytest.mark.parametrize(
@@ -35,6 +36,7 @@ async def test_aembed_query_different_lengths(query: str) -> None:
     model = GoogleGenerativeAIEmbeddings(model=_MODEL)
     result = await model.aembed_query(query)
     assert len(result) == 768
+    assert isinstance(result, list)
 
 
 def test_embed_documents() -> None:
@@ -46,6 +48,8 @@ def test_embed_documents() -> None:
     assert len(result) == 2
     assert len(result[0]) == 768
     assert len(result[1]) == 768
+    assert isinstance(result, list)
+    assert isinstance(result[0], list)
 
 
 async def test_aembed_documents() -> None:
@@ -57,6 +61,8 @@ async def test_aembed_documents() -> None:
     assert len(result) == 2
     assert len(result[0]) == 768
     assert len(result[1]) == 768
+    assert isinstance(result, list)
+    assert isinstance(result[0], list)
 
 
 def test_invalid_model_error_handling() -> None: