aiagents4pharma 1.44.0__py3-none-any.whl → 1.45.1__py3-none-any.whl

aiagents4pharma/talk2knowledgegraphs/tests/test_tools_multimodal_subgraph_extraction.py

@@ -3,6 +3,7 @@ Test cases for tools/subgraph_extraction.py
 """
 
 import pytest
+
 # from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 from ..tools.multimodal_subgraph_extraction import MultimodalSubgraphExtractionTool
 
@@ -24,7 +25,7 @@ def agent_state_fixture():
             "cellular_component": [],
             "biological_process": [],
             "drug": [],
-            "disease": []
+            "disease": [],
         },
         "uploaded_files": [],
         "topk_nodes": 3,
@@ -59,13 +60,16 @@ def test_extract_multimodal_subgraph_wo_doc(agent_state):
 
     # Invoking the subgraph_extraction_tool
     response = subgraph_extraction_tool.invoke(
-        input={"prompt": prompt,
-               "tool_call_id": "subgraph_extraction_tool",
-               "state": agent_state,
-               "arg_data": {"extraction_name": "subkg_12345"}})
+        input={
+            "prompt": prompt,
+            "tool_call_id": "subgraph_extraction_tool",
+            "state": agent_state,
+            "arg_data": {"extraction_name": "subkg_12345"},
+        }
+    )
 
     # Check tool message
-    assert response.update["messages"][-1].tool_call_id  == "subgraph_extraction_tool"
+    assert response.update["messages"][-1].tool_call_id == "subgraph_extraction_tool"
 
     # Check extracted subgraph dictionary
     dic_extracted_graph = response.update["dic_extracted_graph"][0]
@@ -80,13 +84,13 @@ def test_extract_multimodal_subgraph_wo_doc(agent_state):
     assert isinstance(dic_extracted_graph["graph_text"], str)
     # Check if the nodes are in the graph_text
     assert all(
-        n[0] in dic_extracted_graph["graph_text"].replace('"', '')
+        n[0] in dic_extracted_graph["graph_text"].replace('"', "")
         for n in dic_extracted_graph["graph_dict"]["nodes"]
     )
     # Check if the edges are in the graph_text
     assert all(
         ",".join([e[0], str(tuple(e[2]["relation"])), e[1]])
-        in dic_extracted_graph["graph_text"].replace('"', '')
+        in dic_extracted_graph["graph_text"].replace('"', "")
         for e in dic_extracted_graph["graph_dict"]["edges"]
     )
 
@@ -120,13 +124,16 @@ def test_extract_multimodal_subgraph_w_doc(agent_state):
 
     # Invoking the subgraph_extraction_tool
     response = subgraph_extraction_tool.invoke(
-        input={"prompt": prompt,
-               "tool_call_id": "subgraph_extraction_tool",
-               "state": agent_state,
-               "arg_data": {"extraction_name": "subkg_12345"}})
+        input={
+            "prompt": prompt,
+            "tool_call_id": "subgraph_extraction_tool",
+            "state": agent_state,
+            "arg_data": {"extraction_name": "subkg_12345"},
+        }
+    )
 
     # Check tool message
-    assert response.update["messages"][-1].tool_call_id  == "subgraph_extraction_tool"
+    assert response.update["messages"][-1].tool_call_id == "subgraph_extraction_tool"
 
     # Check extracted subgraph dictionary
     dic_extracted_graph = response.update["dic_extracted_graph"][0]
@@ -141,12 +148,12 @@ def test_extract_multimodal_subgraph_w_doc(agent_state):
     assert isinstance(dic_extracted_graph["graph_text"], str)
     # Check if the nodes are in the graph_text
     assert all(
-        n[0] in dic_extracted_graph["graph_text"].replace('"', '')
+        n[0] in dic_extracted_graph["graph_text"].replace('"', "")
         for n in dic_extracted_graph["graph_dict"]["nodes"]
     )
     # Check if the edges are in the graph_text
     assert all(
         ",".join([e[0], str(tuple(e[2]["relation"])), e[1]])
-        in dic_extracted_graph["graph_text"].replace('"', '')
+        in dic_extracted_graph["graph_text"].replace('"', "")
         for e in dic_extracted_graph["graph_dict"]["edges"]
     )

aiagents4pharma/talk2knowledgegraphs/tests/test_tools_subgraph_extraction.py

@@ -4,6 +4,7 @@ Test cases for tools/subgraph_extraction.py
 
 import pytest
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+
 from ..tools.subgraph_extraction import SubgraphExtractionTool
 
 # Define the data path
@@ -41,7 +42,7 @@ def test_extract_subgraph_wo_docs(agent_state):
         agent_state: Agent state in the form of a dictionary.
     """
     prompt = """
-    Extract all relevant information related to nodes of genes related to inflammatory bowel disease 
+    Extract all relevant information related to nodes of genes related to inflammatory bowel disease
     (IBD) that existed in the knowledge graph.
     Please set the extraction name for this process as `subkg_12345`.
     """
@@ -51,13 +52,16 @@ def test_extract_subgraph_wo_docs(agent_state):
 
     # Invoking the subgraph_extraction_tool
     response = subgraph_extraction_tool.invoke(
-        input={"prompt": prompt,
-               "tool_call_id": "subgraph_extraction_tool",
-               "state": agent_state,
-               "arg_data": {"extraction_name": "subkg_12345"}})
+        input={
+            "prompt": prompt,
+            "tool_call_id": "subgraph_extraction_tool",
+            "state": agent_state,
+            "arg_data": {"extraction_name": "subkg_12345"},
+        }
+    )
 
     # Check tool message
-    assert response.update["messages"][-1].tool_call_id  == "subgraph_extraction_tool"
+    assert response.update["messages"][-1].tool_call_id == "subgraph_extraction_tool"
 
     # Check extracted subgraph dictionary
     dic_extracted_graph = response.update["dic_extracted_graph"][0]
@@ -103,7 +107,7 @@ def test_extract_subgraph_w_docs(agent_state):
     ]
 
     prompt = """
-    Extract all relevant information related to nodes of genes related to inflammatory bowel disease 
+    Extract all relevant information related to nodes of genes related to inflammatory bowel disease
     (IBD) that existed in the knowledge graph.
     Please set the extraction name for this process as `subkg_12345`.
     """
@@ -113,13 +117,16 @@ def test_extract_subgraph_w_docs(agent_state):
 
     # Invoking the subgraph_extraction_tool
     response = subgraph_extraction_tool.invoke(
-        input={"prompt": prompt,
-               "tool_call_id": "subgraph_extraction_tool",
-               "state": agent_state,
-               "arg_data": {"extraction_name": "subkg_12345"}})
+        input={
+            "prompt": prompt,
+            "tool_call_id": "subgraph_extraction_tool",
+            "state": agent_state,
+            "arg_data": {"extraction_name": "subkg_12345"},
+        }
+    )
 
     # Check tool message
-    assert response.update["messages"][-1].tool_call_id  == "subgraph_extraction_tool"
+    assert response.update["messages"][-1].tool_call_id == "subgraph_extraction_tool"
 
     # Check extracted subgraph dictionary
     dic_extracted_graph = response.update["dic_extracted_graph"][0]

aiagents4pharma/talk2knowledgegraphs/tests/test_tools_subgraph_summarization.py

@@ -5,6 +5,7 @@ Test cases for tools/subgraph_summarization.py
 import pytest
 from langchain_core.messages import HumanMessage
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+
 from ..agents.t2kg_agent import get_app
 
 # Define the data path
@@ -37,58 +38,104 @@ def input_dict_fixture():
                 "topk_nodes": 3,
                 "topk_edges": 3,
                 "graph_dict": {
-                    'nodes': [('IFNG_(3495)', {}), 
-                              ('IKBKG_(3672)', {}),
-                              ('ATG16L1_(6661)', {}),
-                              ('inflammatory bowel disease_(28158)', {}),
-                              ('Crohn ileitis and jejunitis_(35814)', {}),
-                              ("Crohn's colitis_(83770)", {})],
-                    'edges': [('IFNG_(3495)', 'inflammatory bowel disease_(28158)', 
-                               {'relation': ['gene/protein', 'associated with', 'disease'],
-                                'label': ['gene/protein', 'associated with', 'disease']}), 
-                              ('IFNG_(3495)', "Crohn's colitis_(83770)",
-                               {'relation': ['gene/protein', 'associated with', 'disease'],
-                                'label': ['gene/protein', 'associated with', 'disease']}), 
-                              ('IFNG_(3495)', 'Crohn ileitis and jejunitis_(35814)',
-                               {'relation': ['gene/protein', 'associated with', 'disease'],
-                                'label': ['gene/protein', 'associated with', 'disease']}), 
-                              ('ATG16L1_(6661)', 'IKBKG_(3672)',
-                               {'relation': ['gene/protein', 'ppi', 'gene/protein'],
-                                'label': ['gene/protein', 'ppi', 'gene/protein']}), 
-                              ("Crohn's colitis_(83770)", 'ATG16L1_(6661)',
-                               {'relation': ['disease', 'associated with', 'gene/protein'],
-                                'label': ['disease', 'associated with', 'gene/protein']})]},
+                    "nodes": [
+                        ("IFNG_(3495)", {}),
+                        ("IKBKG_(3672)", {}),
+                        ("ATG16L1_(6661)", {}),
+                        ("inflammatory bowel disease_(28158)", {}),
+                        ("Crohn ileitis and jejunitis_(35814)", {}),
+                        ("Crohn's colitis_(83770)", {}),
+                    ],
+                    "edges": [
+                        (
+                            "IFNG_(3495)",
+                            "inflammatory bowel disease_(28158)",
+                            {
+                                "relation": [
+                                    "gene/protein",
+                                    "associated with",
+                                    "disease",
+                                ],
+                                "label": ["gene/protein", "associated with", "disease"],
+                            },
+                        ),
+                        (
+                            "IFNG_(3495)",
+                            "Crohn's colitis_(83770)",
+                            {
+                                "relation": [
+                                    "gene/protein",
+                                    "associated with",
+                                    "disease",
+                                ],
+                                "label": ["gene/protein", "associated with", "disease"],
+                            },
+                        ),
+                        (
+                            "IFNG_(3495)",
+                            "Crohn ileitis and jejunitis_(35814)",
+                            {
+                                "relation": [
+                                    "gene/protein",
+                                    "associated with",
+                                    "disease",
+                                ],
+                                "label": ["gene/protein", "associated with", "disease"],
+                            },
+                        ),
+                        (
+                            "ATG16L1_(6661)",
+                            "IKBKG_(3672)",
+                            {
+                                "relation": ["gene/protein", "ppi", "gene/protein"],
+                                "label": ["gene/protein", "ppi", "gene/protein"],
+                            },
+                        ),
+                        (
+                            "Crohn's colitis_(83770)",
+                            "ATG16L1_(6661)",
+                            {
+                                "relation": [
+                                    "disease",
+                                    "associated with",
+                                    "gene/protein",
+                                ],
+                                "label": ["disease", "associated with", "gene/protein"],
+                            },
+                        ),
+                    ],
+                },
                 "graph_text": """
             node_id,node_attr
-            IFNG_(3495),"IFNG belongs to gene/protein category. 
-            This gene encodes a soluble cytokine that is a member of the type II interferon class. 
-            The encoded protein is secreted by cells of both the innate and adaptive immune systems. 
-            The active protein is a homodimer that binds to the interferon gamma receptor 
-            which triggers a cellular response to viral and microbial infections. 
-            Mutations in this gene are associated with an increased susceptibility to viral, 
-            bacterial and parasitic infections and to several autoimmune diseases. 
+            IFNG_(3495),"IFNG belongs to gene/protein category.
+            This gene encodes a soluble cytokine that is a member of the type II interferon class.
+            The encoded protein is secreted by cells of both the innate and adaptive immune systems.
+            The active protein is a homodimer that binds to the interferon gamma receptor
+            which triggers a cellular response to viral and microbial infections.
+            Mutations in this gene are associated with an increased susceptibility to viral,
+            bacterial and parasitic infections and to several autoimmune diseases.
             [provided by RefSeq, Dec 2015]."
-            IKBKG_(3672),"IKBKG belongs to gene/protein category. This gene encodes the regulatory 
-            subunit of the inhibitor of kappaB kinase (IKK) complex, which activates NF-kappaB 
-            resulting in activation of genes involved in inflammation, immunity, cell survival, 
-            and other pathways. Mutations in this gene result in incontinentia pigmenti, 
-            hypohidrotic ectodermal dysplasia, and several other types of immunodeficiencies. 
-            A pseudogene highly similar to this locus is located in an adjacent region of the 
+            IKBKG_(3672),"IKBKG belongs to gene/protein category. This gene encodes the regulatory
+            subunit of the inhibitor of kappaB kinase (IKK) complex, which activates NF-kappaB
+            resulting in activation of genes involved in inflammation, immunity, cell survival,
+            and other pathways. Mutations in this gene result in incontinentia pigmenti,
+            hypohidrotic ectodermal dysplasia, and several other types of immunodeficiencies.
+            A pseudogene highly similar to this locus is located in an adjacent region of the
             X chromosome. [provided by RefSeq, Mar 2016]."
-            ATG16L1_(6661),"ATG16L1 belongs to gene/protein category. The protein encoded 
-            by this gene is part of a large protein complex that is necessary for autophagy, 
-            the major process by which intracellular components are targeted to lysosomes 
-            for degradation. Defects in this gene are a cause of susceptibility to inflammatory 
-            bowel disease type 10 (IBD10). Several transcript variants encoding different 
+            ATG16L1_(6661),"ATG16L1 belongs to gene/protein category. The protein encoded
+            by this gene is part of a large protein complex that is necessary for autophagy,
+            the major process by which intracellular components are targeted to lysosomes
+            for degradation. Defects in this gene are a cause of susceptibility to inflammatory
+            bowel disease type 10 (IBD10). Several transcript variants encoding different
             isoforms have been found for this gene.[provided by RefSeq, Jun 2010]."
-            inflammatory bowel disease_(28158),inflammatory bowel disease belongs to disease 
-            category. Any inflammatory bowel disease in which the cause of the disease 
-            is a mutation in the NOD2 gene.  
-            Crohn ileitis and jejunitis_(35814),Crohn ileitis and jejunitis belongs to 
-            disease category. An Crohn disease involving a pathogenic inflammatory 
-            response in the ileum.  
-            Crohn's colitis_(83770),Crohn's colitis belongs to disease category. 
-            Crohn's disease affecting the colon.  
+            inflammatory bowel disease_(28158),inflammatory bowel disease belongs to disease
+            category. Any inflammatory bowel disease in which the cause of the disease
+            is a mutation in the NOD2 gene.
+            Crohn ileitis and jejunitis_(35814),Crohn ileitis and jejunitis belongs to
+            disease category. An Crohn disease involving a pathogenic inflammatory
+            response in the ileum.
+            Crohn's colitis_(83770),Crohn's colitis belongs to disease category.
+            Crohn's disease affecting the colon.
 
             head_id,edge_type,tail_id
             Crohn's colitis_(83770),"('disease', 'associated with', 'gene/protein')",
@@ -129,7 +176,7 @@ def test_summarize_subgraph(input_dict):
         input_dict,
     )
     prompt = """
-    Please directly invoke `subgraph_summarization` tool without calling any other tools 
+    Please directly invoke `subgraph_summarization` tool without calling any other tools
     to respond to the following prompt:
 
     You are given a subgraph in the forms of textualized subgraph representing

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_embeddings.py

@@ -3,8 +3,10 @@ Test cases for utils/embeddings/embeddings.py
 """
 
 import pytest
+
 from ..utils.embeddings.embeddings import Embeddings
 
+
 class TestEmbeddings(Embeddings):
     """Test implementation of the Embeddings interface for testing purposes."""
 
@@ -14,6 +16,7 @@ class TestEmbeddings(Embeddings):
     def embed_query(self, text: str) -> list[float]:
         return [0.1, 0.2, 0.3]
 
+
 def test_embed_documents():
     """Test embedding documents using the Embeddings interface."""
     embeddings = TestEmbeddings()
@@ -29,6 +32,7 @@ def test_embed_query():
     result = embeddings.embed_query(text)
     assert result == [0.1, 0.2, 0.3]
 
+
 @pytest.mark.asyncio
 async def test_aembed_documents():
     """Test asynchronous embedding of documents using the Embeddings interface."""

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_huggingface.py

@@ -3,8 +3,10 @@ Test cases for utils/embeddings/huggingface.py
 """
 
 import pytest
+
 from ..utils.embeddings.huggingface import EmbeddingWithHuggingFace
 
+
 @pytest.fixture(name="embedding_model")
 def embedding_model_fixture():
     """Return the configuration object for the HuggingFace embedding model and model object"""
@@ -14,6 +16,7 @@ def embedding_model_fixture():
         truncation=True,
     )
 
+
 def test_embedding_with_huggingface_embed_documents(embedding_model):
     """Test embedding documents using the EmbeddingWithHuggingFace class."""
     # Perform embedding
@@ -23,6 +26,7 @@ def test_embedding_with_huggingface_embed_documents(embedding_model):
     assert len(result) == 3
     assert len(result[0]) == 768
 
+
 def test_embedding_with_huggingface_embed_query(embedding_model):
     """Test embedding a query using the EmbeddingWithHuggingFace class."""
     # Perform embedding
@@ -31,6 +35,7 @@ def test_embedding_with_huggingface_embed_query(embedding_model):
     # Check the result
     assert len(result) == 768
 
+
 def test_embedding_with_huggingface_failed():
     """Test embedding documents using the EmbeddingWithHuggingFace class."""
     # Check if the model is available on HuggingFace Hub

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_nim_molmim.py

@@ -5,27 +5,28 @@ Test cases for utils/embeddings/nim_molmim.py
 """
 
 import unittest
-from unittest.mock import patch, MagicMock
+from unittest.mock import MagicMock, patch
+
 from ..utils.embeddings.nim_molmim import EmbeddingWithMOLMIM
 
+
 class TestEmbeddingWithMOLMIM(unittest.TestCase):
     """
     Test cases for EmbeddingWithMOLMIM class.
     """
+
     def setUp(self):
         self.base_url = "https://fake-nim-api.com/embeddings"
         self.embeddings_model = EmbeddingWithMOLMIM(self.base_url)
         self.test_texts = ["CCO", "CCC", "C=O"]
         self.test_query = "CCO"
-        self.mock_response = {
-            "embeddings": [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]]
-        }
+        self.mock_response = {"embeddings": [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]]}
 
     @patch("requests.post")
     def test_embed_documents(self, mock_post):
-        '''
+        """
         Test the embed_documents method.
-        '''
+        """
         # Mock the response from requests.post
         mock_post.return_value = MagicMock()
         mock_post.return_value.json.return_value = self.mock_response
@@ -34,19 +35,16 @@ class TestEmbeddingWithMOLMIM(unittest.TestCase):
         self.assertEqual(embeddings, self.mock_response["embeddings"])
         mock_post.assert_called_once_with(
             self.base_url,
-            headers={
-                'accept': 'application/json',
-                'Content-Type': 'application/json'
-            },
+            headers={"accept": "application/json", "Content-Type": "application/json"},
             data='{"sequences": ["CCO", "CCC", "C=O"]}',
-            timeout=60
+            timeout=60,
         )
 
     @patch("requests.post")
     def test_embed_query(self, mock_post):
-        '''
+        """
         Test the embed_query method.
-        '''
+        """
         # Mock the response from requests.post
         mock_post.return_value = MagicMock()
         mock_post.return_value.json.return_value = {"embeddings": [[0.1, 0.2, 0.3]]}
@@ -55,10 +53,7 @@ class TestEmbeddingWithMOLMIM(unittest.TestCase):
         self.assertEqual(embedding, [[0.1, 0.2, 0.3]])
         mock_post.assert_called_once_with(
             self.base_url,
-            headers={
-                'accept': 'application/json',
-                'Content-Type': 'application/json'
-            },
+            headers={"accept": "application/json", "Content-Type": "application/json"},
             data='{"sequences": ["CCO"]}',
-            timeout=60
+            timeout=60,
         )

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_ollama.py

@@ -2,17 +2,20 @@
 Test cases for utils/embeddings/ollama.py
 """
 
-import pytest
 import ollama
+import pytest
+
 from ..utils.embeddings.ollama import EmbeddingWithOllama
 
+
 @pytest.fixture(name="ollama_config")
 def fixture_ollama_config():
     """Return a dictionary with Ollama configuration."""
     return {
-        "model_name": "all-minilm", # Choose a small model
+        "model_name": "all-minilm",  # Choose a small model
     }
 
+
 def test_no_model_ollama(ollama_config):
     """Test the case when the Ollama model is not available."""
     cfg = ollama_config
@@ -25,16 +28,19 @@ def test_no_model_ollama(ollama_config):
 
     # Check if the model is available
     with pytest.raises(
-        ValueError, match=f"Error: Pulled {cfg["model_name"]} model and restarted Ollama server."
+        ValueError,
+        match=f"Error: Pulled {cfg['model_name']} model and restarted Ollama server.",
     ):
         EmbeddingWithOllama(model_name=cfg["model_name"])
 
+
 @pytest.fixture(name="embedding_model")
 def embedding_model_fixture(ollama_config):
     """Return the configuration object for the Ollama embedding model and model object"""
     cfg = ollama_config
     return EmbeddingWithOllama(model_name=cfg["model_name"])
 
+
 def test_embedding_with_ollama_embed_documents(embedding_model):
     """Test embedding documents using the EmbeddingWithOllama class."""
     # Perform embedding
@@ -44,6 +50,7 @@ def test_embedding_with_ollama_embed_documents(embedding_model):
     assert len(result) == 3
     assert len(result[0]) == 384
 
+
 def test_embedding_with_ollama_embed_query(embedding_model):
     """Test embedding a query using the EmbeddingWithOllama class."""
     # Perform embedding

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_enrichments_enrichments.py

@@ -4,19 +4,19 @@ Test cases for utils/enrichments/enrichments.py
 
 from ..utils.enrichments.enrichments import Enrichments
 
+
 class TestEnrichments(Enrichments):
     """Test implementation of the Enrichments interface for testing purposes."""
 
     def enrich_documents(self, texts: list[str]) -> list[list[float]]:
-        return [
-            f"Additional text description of {text} as the input." for text in texts
-        ]
+        return [f"Additional text description of {text} as the input." for text in texts]
 
     def enrich_documents_with_rag(self, texts, docs):
         # Currently we don't have a RAG model to test this method.
         # Thus, we will just call the enrich_documents method instead.
         return self.enrich_documents(texts)
 
+
 def test_enrich_documents():
     """Test enriching documents using the Enrichments interface."""
     enrichments = TestEnrichments()
@@ -27,6 +27,7 @@ def test_enrich_documents():
         "Additional text description of text2 as the input.",
     ]
 
+
 def test_enrich_documents_with_rag():
     """Test enriching documents with RAG using the Enrichments interface."""
     enrichments = TestEnrichments()

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_enrichments_ollama.py

@@ -2,8 +2,9 @@
 Test cases for utils/enrichments/ollama.py
 """
 
-import pytest
 import ollama
+import pytest
+
 from ..utils.enrichments.ollama import EnrichmentWithOllama
 
 
@@ -13,7 +14,7 @@ def fixture_ollama_config():
     return {
         "model_name": "llama3.2:1b",
         "prompt_enrichment": """
-            Given the input as a list of strings, please return the list of addditional information 
+            Given the input as a list of strings, please return the list of addditional information
             of each input terms using your prior knowledge.
 
             Example:

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_enrichments_ols.py

@@ -5,6 +5,7 @@ Test cases for utils/enrichments/ols_terms.py
 """
 
 import pytest
+
 from ..utils.enrichments.ols_terms import EnrichmentWithOLS
 
 # In this test, we will consider 5 examples:

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_enrichments_pubchem.py

@@ -5,23 +5,26 @@ Test cases for utils/enrichments/pubchem_strings.py
 """
 
 import pytest
+
 from ..utils.enrichments.pubchem_strings import EnrichmentWithPubChem
 
 # In this test, we will consider 2 examples:
 # 1. PubChem ID: 5311000 (Alclometasone)
 # 2. PubChem ID: 1X (Fake ID)
 # The expected SMILES representation for the first PubChem ID is:
-SMILES_FIRST = 'C[C@@H]1C[C@H]2[C@@H]3[C@@H](CC4=CC(=O)C=C[C@@]'
-SMILES_FIRST += '4([C@H]3[C@H](C[C@@]2([C@]1(C(=O)CO)O)C)O)C)Cl'
+SMILES_FIRST = "C[C@@H]1C[C@H]2[C@@H]3[C@@H](CC4=CC(=O)C=C[C@@]"
+SMILES_FIRST += "4([C@H]3[C@H](C[C@@]2([C@]1(C(=O)CO)O)C)O)C)Cl"
 # The expected description for the first PubChem ID starts with:
 DESCRIPTION_FIRST = "Alclometasone is a prednisolone compound having an alpha-chloro substituent"
 # The expected SMILES representation and description for the second PubChem ID is None.
 
+
 @pytest.fixture(name="enrich_obj")
 def fixture_pubchem_config():
     """Return a dictionary with the configuration for the PubChem enrichment."""
     return EnrichmentWithPubChem()
 
+
 def test_enrich_documents(enrich_obj):
     """Test the enrich_documents method."""
     pubchem_ids = ["5311000", "1X"]
@@ -30,11 +33,13 @@ def test_enrich_documents(enrich_obj):
     assert enriched_descriptions[0].startswith(DESCRIPTION_FIRST)
     assert enriched_descriptions[1] is None
 
+
 def test_enrich_documents_with_rag(enrich_obj):
     """Test the enrich_documents_with_rag method."""
     pubchem_ids = ["5311000", "1X"]
-    enriched_descriptions, enriched_strings = enrich_obj.enrich_documents_with_rag(pubchem_ids,
-                                                                                   None)
+    enriched_descriptions, enriched_strings = enrich_obj.enrich_documents_with_rag(
+        pubchem_ids, None
+    )
     assert enriched_strings == [SMILES_FIRST, None]
     assert enriched_descriptions[0].startswith(DESCRIPTION_FIRST)
     assert enriched_descriptions[1] is None