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

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_pubchem_utils.py

@@ -1,49 +1,123 @@
 """
 Test cases for utils/pubchem_utils.py
+
+These tests mock PubChem HTTP calls to avoid network dependency and verify
+parsing logic without modifying business logic.
 """
 
+from types import SimpleNamespace
+
+import pytest
+
 from ..utils import pubchem_utils
 
+
+def _resp(json_data):
+    """Create a minimal response-like object exposing .json()."""
+    return SimpleNamespace(json=lambda: json_data)
+
+
+@pytest.fixture(autouse=True)
+def mock_requests_get(monkeypatch):
+    """Mock requests.get used by pubchem_utils to return deterministic JSON."""
+
+    def _mock_get(url, _timeout=60, **_kwargs):
+        # CAS RN lookup: ethyl carbonate 105-58-8 -> CID 7766
+        if "substance/xref/RN/105-58-8/record/JSON" in url:
+            return _resp(
+                {"PC_Substances": [{"compound": [{"id": {"type": 1, "id": {"cid": 7766}}}]}]}
+            )
+
+        # DrugBank: DB00240 (Alclometasone) -> CID 5311000
+        if (
+            "substance/sourceid//drugbank/DB00240/JSON" in url
+            or "substance/sourceid/drugbank/DB00240/JSON" in url
+        ):
+            return _resp(
+                {"PC_Substances": [{"compound": [{"id": {"type": 1, "id": {"cid": 5311000}}}]}]}
+            )
+
+        # CTD: D002083 (Butylated Hydroxyanisole) -> CID 24667
+        if (
+            "substance/sourceid//comparative toxicogenomics database/D002083/JSON" in url
+            or "substance/sourceid/comparative toxicogenomics database/D002083/JSON" in url
+        ):
+            return _resp(
+                {"PC_Substances": [{"compound": [{"id": {"type": 1, "id": {"cid": 24667}}}]}]}
+            )
+
+        # CID description for 5311000
+        if (
+            "compound/cid//5311000/description/JSON" in url
+            or "compound/cid/5311000/description/JSON" in url
+        ):
+            return _resp(
+                {
+                    "InformationList": {
+                        "Information": [
+                            {
+                                "Description": (
+                                    "Alclometasone is a prednisolone compound having an "
+                                    "alpha-chloro substituent at the 9alpha-position."
+                                ),
+                            }
+                        ]
+                    }
+                }
+            )
+
+        # Default empty response
+        return _resp({})
+
+    monkeypatch.setattr(pubchem_utils.requests, "get", _mock_get)
+    return _mock_get
+
+
 def test_cas_rn2pubchem_cid():
     """
     Test the casRN2pubchem_cid function.
 
-    The CAS RN for ethyl cabonate is 105-58-8.
-    The PubChem CID for ethyl cabonate is 7766.
+    The CAS RN for ethyl carbonate is 105-58-8.
+    The PubChem CID for ethyl carbonate is 7766.
     """
     casrn = "105-58-8"
     pubchem_cid = pubchem_utils.cas_rn2pubchem_cid(casrn)
     assert pubchem_cid == 7766
 
+
 def test_external_id2pubchem_cid():
     """
     Test the external_id2pubchem_cid function.
 
-    The DrugBank ID for Alclometasone is DB00240.
-    The PubChem CID for Alclometasone is 5311000.
-
-    The CTD ID for Butylated Hydroxyanisole is D002083
-    The PubChem CID for Butylated Hydroxyanisole is 24667.
+    The DrugBank ID for Alclometasone is DB00240 -> CID 5311000.
+    The CTD ID for Butylated Hydroxyanisole is D002083 -> CID 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)
+        "comparative toxicogenomics database", ctd_id
+    )
     assert pubchem_cid == 24667
 
+
 def test_pubchem_cid_description():
     """
     Test the pubchem_cid_description function.
 
     The PubChem CID for Alclometasone is 5311000.
-    The description for Alclometasone starts with 
-        "Alclometasone is a prednisolone compound having an alpha-chloro substituent".
+    The description starts with the expected prefix.
     """
     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"
+    )
+
+
+def test_mock_fallback_path_coverage():
+    """Exercise the default branch of the mocked requests.get for coverage."""
+    resp = pubchem_utils.requests.get("unknown://unmatched/url", timeout=1)
+    assert resp.json() == {}

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