refactor: improve code formatting for consistency and readability in document embedder and tests

Author: garybadwal

integrations/google_genai/src/haystack_integrations/components/embedders/google_genai/document_embedder.py

@@ -85,8 +85,7 @@ def __init__(
         self._meta_fields_to_embed = meta_fields_to_embed or []
         self._embedding_separator = embedding_separator
         self._client = genai.Client(api_key=api_key.resolve_value())
-        self._config = config if config is not None else {
-            "task_type": "SEMANTIC_SIMILARITY"}
+        self._config = config if config is not None else {"task_type": "SEMANTIC_SIMILARITY"}
 
     def to_dict(self) -> Dict[str, Any]:
         """
@@ -134,8 +133,7 @@ def _prepare_texts_to_embed(self, documents: List[Document]) -> Dict[str, str]:
             ]
 
             text_to_embed = (
-                self._prefix + self._embedding_separator.join(
-                    [*meta_values_to_embed, doc.content or ""]) + self._suffix
+                self._prefix + self._embedding_separator.join([*meta_values_to_embed, doc.content or ""]) + self._suffix
             )
             texts_to_embed.append(text_to_embed)
 
@@ -151,11 +149,9 @@ def _embed_batch(self, texts_to_embed: List[str], batch_size: int) -> Tuple[List
         for batch in tqdm(
             batched(texts_to_embed, batch_size), disable=not self._progress_bar, desc="Calculating embeddings"
         ):
-            args: Dict[str, Any] = {"model": self._model,
-                                    "contents": [b[1] for b in batch]}
+            args: Dict[str, Any] = {"model": self._model, "contents": [b[1] for b in batch]}
             if self._config:
-                args["config"] = types.EmbedContentConfig(
-                    **self._config) if self._config else None
+                args["config"] = types.EmbedContentConfig(**self._config) if self._config else None
 
             response = self._client.models.embed_content(**args)
 
@@ -189,8 +185,7 @@ def run(self, documents: List[Document]) -> Dict[str, Union[List[Document], Dict
 
         texts_to_embed = self._prepare_texts_to_embed(documents=documents)
 
-        embeddings, meta = self._embed_batch(
-            texts_to_embed=texts_to_embed, batch_size=self._batch_size)
+        embeddings, meta = self._embed_batch(texts_to_embed=texts_to_embed, batch_size=self._batch_size)
 
         for doc, emb in zip(documents, embeddings):
             doc.embedding = emb

integrations/google_genai/tests/test_document_embedder.py

@@ -92,8 +92,7 @@ def test_to_dict(self, monkeypatch):
         data = component.to_dict()
         assert data == {
             "type": (
-                "haystack_integrations.components.embedders"
-                ".google_genai.document_embedder.GoogleGenAIDocumentEmbedder"
+                "haystack_integrations.components.embedders.google_genai.document_embedder.GoogleGenAIDocumentEmbedder"
             ),
             "init_parameters": {
                 "model": "text-embedding-004",
@@ -124,8 +123,7 @@ def test_to_dict_with_custom_init_parameters(self, monkeypatch):
         data = component.to_dict()
         assert data == {
             "type": (
-                "haystack_integrations.components.embedders"
-                ".google_genai.document_embedder.GoogleGenAIDocumentEmbedder"
+                "haystack_integrations.components.embedders.google_genai.document_embedder.GoogleGenAIDocumentEmbedder"
             ),
             "init_parameters": {
                 "model": "model",
@@ -142,8 +140,7 @@ def test_to_dict_with_custom_init_parameters(self, monkeypatch):
 
     def test_prepare_texts_to_embed_w_metadata(self):
         documents = [
-            Document(id=f"{i}", content=f"document number {i}:\ncontent", meta={
-                     "meta_field": f"meta_value {i}"})
+            Document(id=f"{i}", content=f"document number {i}:\ncontent", meta={"meta_field": f"meta_value {i}"})
             for i in range(5)
         ]
 
@@ -157,12 +154,11 @@ def test_prepare_texts_to_embed_w_metadata(self):
             "meta_value 1 | document number 1:\ncontent",
             "meta_value 2 | document number 2:\ncontent",
             "meta_value 3 | document number 3:\ncontent",
-            "meta_value 4 | document number 4:\ncontent"
+            "meta_value 4 | document number 4:\ncontent",
         ]
 
     def test_run_wrong_input_format(self):
-        embedder = GoogleGenAIDocumentEmbedder(
-            api_key=Secret.from_token("fake-api-key"))
+        embedder = GoogleGenAIDocumentEmbedder(api_key=Secret.from_token("fake-api-key"))
 
         # wrong formats
         string_input = "text"
@@ -175,8 +171,7 @@ def test_run_wrong_input_format(self):
             embedder.run(documents=list_integers_input)
 
     def test_run_on_empty_list(self):
-        embedder = GoogleGenAIDocumentEmbedder(
-            api_key=Secret.from_token("fake-api-key"))
+        embedder = GoogleGenAIDocumentEmbedder(api_key=Secret.from_token("fake-api-key"))
 
         empty_list_input = []
         result = embedder.run(documents=empty_list_input)
@@ -192,14 +187,12 @@ def test_run_on_empty_list(self):
     def test_run(self):
         docs = [
             Document(content="I love cheese", meta={"topic": "Cuisine"}),
-            Document(content="A transformer is a deep learning architecture", meta={
-                     "topic": "ML"}),
+            Document(content="A transformer is a deep learning architecture", meta={"topic": "ML"}),
         ]
 
         model = "text-embedding-004"
 
-        embedder = GoogleGenAIDocumentEmbedder(model=model, meta_fields_to_embed=[
-                                               "topic"], embedding_separator=" | ")
+        embedder = GoogleGenAIDocumentEmbedder(model=model, meta_fields_to_embed=["topic"], embedding_separator=" | ")
 
         result = embedder.run(documents=docs)
         documents_with_embeddings = result["documents"]
@@ -211,6 +204,6 @@ def test_run(self):
             assert len(doc.embedding) == 768
             assert all(isinstance(x, float) for x in doc.embedding)
 
-        assert (
-            "text" in result["meta"]["model"] and "004" in result["meta"]["model"]
-        ), "The model name does not contain 'text' and '004'"
+        assert "text" in result["meta"]["model"] and "004" in result["meta"]["model"], (
+            "The model name does not contain 'text' and '004'"
+        )

integrations/google_genai/tests/test_text_embedder.py

@@ -160,6 +160,6 @@ def test_run(self):
         assert len(result["embedding"]) == 768
         assert all(isinstance(x, float) for x in result["embedding"])
 
-        assert (
-            "text" in result["meta"]["model"] and "004" in result["meta"]["model"]
-        ), "The model name does not contain 'text' and '004'"
+        assert "text" in result["meta"]["model"] and "004" in result["meta"]["model"], (
+            "The model name does not contain 'text' and '004'"
+        )