refactor: format code for better readability and consistency in document embedder

Author: garybadwal

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

@@ -85,7 +85,8 @@ 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]:
         """
@@ -127,14 +128,14 @@ def _prepare_texts_to_embed(self, documents: List[Document]) -> Dict[str, str]:
         texts_to_embed: List[str] = []
         for doc in documents:
             meta_values_to_embed = [
-                str(doc.meta[key]) 
-                for key in 
-                self._meta_fields_to_embed 
+                str(doc.meta[key])
+                for key in self._meta_fields_to_embed
                 if key in doc.meta and doc.meta[key] is not None
             ]
 
             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)
 
@@ -150,9 +151,11 @@ 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)
 
@@ -186,7 +189,8 @@ 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

@@ -142,7 +142,8 @@ 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)
         ]
 
@@ -152,15 +153,16 @@ def test_prepare_texts_to_embed_w_metadata(self):
 
         prepared_texts = embedder._prepare_texts_to_embed(documents)
         assert prepared_texts == [
-            'meta_value 0 | document number 0:\ncontent', 
-            '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 0 | document number 0:\ncontent",
+            "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"
         ]
 
     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"
@@ -173,7 +175,8 @@ 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)
@@ -189,12 +192,14 @@ 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"]