aiagents4pharma 1.43.0__py3-none-any.whl → 1.45.0__py3-none-any.whl

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_pubchem_utils.py

@@ -4,6 +4,7 @@ Test cases for utils/pubchem_utils.py
 
 from ..utils import pubchem_utils
 
+
 def test_cas_rn2pubchem_cid():
     """
     Test the casRN2pubchem_cid function.
@@ -15,6 +16,7 @@ def test_cas_rn2pubchem_cid():
     pubchem_cid = pubchem_utils.cas_rn2pubchem_cid(casrn)
     assert pubchem_cid == 7766
 
+
 def test_external_id2pubchem_cid():
     """
     Test the external_id2pubchem_cid function.
@@ -26,24 +28,26 @@ def test_external_id2pubchem_cid():
     The PubChem CID for Butylated Hydroxyanisole is 24667.
     """
     drugbank_id = "DB00240"
-    pubchem_cid = pubchem_utils.external_id2pubchem_cid('drugbank', drugbank_id)
+    pubchem_cid = pubchem_utils.external_id2pubchem_cid("drugbank", drugbank_id)
     assert pubchem_cid == 5311000
 
     ctd_id = "D002083"
     pubchem_cid = pubchem_utils.external_id2pubchem_cid(
-                    'comparative toxicogenomics database',
-                    ctd_id)
-    assert pubchem_cid == 24667
+        "comparative toxicogenomics database", ctd_id
+    )
+    assert pubchem_cid == 8456
+
 
 def test_pubchem_cid_description():
     """
     Test the pubchem_cid_description function.
 
     The PubChem CID for Alclometasone is 5311000.
-    The description for Alclometasone starts with 
+    The description for Alclometasone starts with
         "Alclometasone is a prednisolone compound having an alpha-chloro substituent".
     """
     pubchem_cid = 5311000
     description = pubchem_utils.pubchem_cid_description(pubchem_cid)
     assert description.startswith(
-        "Alclometasone is a prednisolone compound having an alpha-chloro substituent")
+        "Alclometasone is a prednisolone compound having an alpha-chloro substituent"
+    )

aiagents4pharma/talk2knowledgegraphs/tools/__init__.py

@@ -1,8 +1,11 @@
-'''
+"""
 This file is used to import all the models in the package.
-'''
-from . import subgraph_extraction
-from . import multimodal_subgraph_extraction
-from . import milvus_multimodal_subgraph_extraction
-from . import subgraph_summarization
-from . import graphrag_reasoning
+"""
+
+from . import (
+    graphrag_reasoning,
+    milvus_multimodal_subgraph_extraction,
+    multimodal_subgraph_extraction,
+    subgraph_extraction,
+    subgraph_summarization,
+)

aiagents4pharma/talk2knowledgegraphs/tools/graphrag_reasoning.py

@@ -3,20 +3,21 @@ Tool for performing Graph RAG reasoning.
 """
 
 import logging
-from typing import Type, Annotated
-from pydantic import BaseModel, Field
-from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.messages import ToolMessage
-from langchain_core.tools.base import InjectedToolCallId
-from langchain_core.tools import BaseTool
-from langchain_core.vectorstores import InMemoryVectorStore
-from langchain.chains.retrieval import create_retrieval_chain
+from typing import Annotated
+
+import hydra
 from langchain.chains.combine_documents import create_stuff_documents_chain
+from langchain.chains.retrieval import create_retrieval_chain
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.document_loaders import PyPDFLoader
-from langgraph.types import Command
+from langchain_core.messages import ToolMessage
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.tools import BaseTool
+from langchain_core.tools.base import InjectedToolCallId
+from langchain_core.vectorstores import InMemoryVectorStore
 from langgraph.prebuilt import InjectedState
-import hydra
+from langgraph.types import Command
+from pydantic import BaseModel, Field
 
 # Initialize logger
 logging.basicConfig(level=logging.INFO)
@@ -33,9 +34,7 @@ class GraphRAGReasoningInput(BaseModel):
         extraction_name: Name assigned to the subgraph extraction process
     """
 
-    tool_call_id: Annotated[str, InjectedToolCallId] = Field(
-        description="Tool call ID."
-    )
+    tool_call_id: Annotated[str, InjectedToolCallId] = Field(description="Tool call ID.")
     state: Annotated[dict, InjectedState] = Field(description="Injected state.")
     prompt: str = Field(description="Prompt to interact with the backend.")
     extraction_name: str = Field(
@@ -53,7 +52,7 @@ class GraphRAGReasoningTool(BaseTool):
     name: str = "graphrag_reasoning"
     description: str = """A tool to perform reasoning using a Graph RAG approach
                         by considering textualized subgraph context and document context."""
-    args_schema: Type[BaseModel] = GraphRAGReasoningInput
+    args_schema: type[BaseModel] = GraphRAGReasoningInput
 
     def _run(
         self,
@@ -71,9 +70,7 @@ class GraphRAGReasoningTool(BaseTool):
             prompt: The prompt to interact with the backend.
             extraction_name: The name assigned to the subgraph extraction process.
         """
-        logger.log(
-            logging.INFO, "Invoking graphrag_reasoning tool for %s", extraction_name
-        )
+        logger.log(logging.INFO, "Invoking graphrag_reasoning tool for %s", extraction_name)
 
         # Load Hydra configuration
         with hydra.initialize(version_base=None, config_path="../configs"):
@@ -88,9 +85,7 @@ class GraphRAGReasoningTool(BaseTool):
             for uploaded_file in state["uploaded_files"]:
                 if uploaded_file["file_type"] == "drug_data":
                     # Load documents
-                    raw_documents = PyPDFLoader(
-                        file_path=uploaded_file["file_path"]
-                    ).load()
+                    raw_documents = PyPDFLoader(file_path=uploaded_file["file_path"]).load()
 
                     # Split documents
                     # May need to find an optimal chunk size and overlap configuration