aiagents4pharma 1.21.0__py3-none-any.whl → 1.22.0__py3-none-any.whl

aiagents4pharma/talk2knowledgegraphs/configs/__init__.py

@@ -5,3 +5,4 @@ Import all the modules in the package
 from . import agents
 from . import tools
 from . import app
+from . import utils

aiagents4pharma/talk2knowledgegraphs/configs/config.yaml

@@ -4,4 +4,5 @@ defaults:
   - tools/subgraph_extraction: default
   - tools/subgraph_summarization: default
   - tools/graphrag_reasoning: default
+  - utils/pubchem_utils: default
   - app/frontend: default

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_nim_molmim.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python3
+
+"""
+Test cases for utils/embeddings/nim_molmim.py
+"""
+
+import unittest
+from unittest.mock import patch, MagicMock
+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]]
+        }
+
+    @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
+        embeddings = self.embeddings_model.embed_documents(self.test_texts)
+        # Assertions
+        self.assertEqual(embeddings, self.mock_response["embeddings"])
+        mock_post.assert_called_once_with(
+            self.base_url,
+            headers={
+                'accept': 'application/json',
+                'Content-Type': 'application/json'
+            },
+            data='{"sequences": ["CCO", "CCC", "C=O"]}',
+            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]]}
+        embedding = self.embeddings_model.embed_query(self.test_query)
+        # Assertions
+        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'
+            },
+            data='{"sequences": ["CCO"]}',
+            timeout=60
+        )

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_enrichments_pubchem.py

@@ -0,0 +1,33 @@
+#!/usr/bin/env python3
+
+"""
+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'
+# The expected SMILES representation 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"]
+    enriched_strings = enrich_obj.enrich_documents(pubchem_ids)
+    assert enriched_strings == [SMILES_FIRST, None]
+
+def test_enrich_documents_with_rag(enrich_obj):
+    """Test the enrich_documents_with_rag method."""
+    pubchem_ids = ["5311000", "1X"]
+    enriched_strings = enrich_obj.enrich_documents_with_rag(pubchem_ids, None)
+    assert enriched_strings == [SMILES_FIRST, None]

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_pubchem_utils.py

@@ -0,0 +1,16 @@
+"""
+Test cases for utils/pubchem_utils.py
+"""
+
+from ..utils import pubchem_utils
+
+def test_drugbank_id2pubchem_cid():
+    """
+    Test the drugbank_id2pubchem_cid method.
+
+    The DrugBank ID for Alclometasone is DB00240.
+    The PubChem CID for Alclometasone is 5311000.
+    """
+    drugbank_id = "DB00240"
+    pubchem_cid = pubchem_utils.drugbank_id2pubchem_cid(drugbank_id)
+    assert pubchem_cid == 5311000

aiagents4pharma/talk2knowledgegraphs/utils/__init__.py

@@ -5,3 +5,4 @@ from . import embeddings
 from . import enrichments
 from . import extractions
 from . import kg_utils
+from . import pubchem_utils

aiagents4pharma/talk2knowledgegraphs/utils/embeddings/__init__.py

@@ -5,3 +5,4 @@ from . import embeddings
 from . import sentence_transformer
 from . import huggingface
 from . import ollama
+from . import nim_molmim

aiagents4pharma/talk2knowledgegraphs/utils/embeddings/nim_molmim.py

@@ -0,0 +1,54 @@
+"""
+Embedding class using MOLMIM model from NVIDIA NIM.
+"""
+
+import json
+from typing import List
+import requests
+from .embeddings import Embeddings
+
+class EmbeddingWithMOLMIM(Embeddings):
+    """
+    Embedding class using MOLMIM model from NVIDIA NIM
+    """
+    def __init__(self, base_url: str):
+        """
+        Initialize the EmbeddingWithMOLMIM class.
+
+        Args:
+            base_url: The base URL for the NIM/MOLMIM model.
+        """
+        # Set base URL
+        self.base_url = base_url
+
+    def embed_documents(self, texts: List[str]) -> List[float]:
+        """
+        Generate embedding for a list of SMILES strings using MOLMIM model.
+
+        Args:
+            texts: The list of SMILES strings to be embedded.
+
+        Returns:
+            The list of embeddings for the given SMILES strings.
+        """
+        headers = {
+            'accept': 'application/json',
+            'Content-Type': 'application/json'
+        }
+        data = json.dumps({"sequences": texts})
+        response = requests.post(self.base_url, headers=headers, data=data, timeout=60)
+        embeddings = response.json()["embeddings"]
+        return embeddings
+
+    def embed_query(self, text: str) -> List[float]:
+        """
+        Generate embeddings for an input query using MOLMIM model.
+
+        Args:
+            text: A query to be embedded.
+        Returns:
+            The embeddings for the given query.
+        """
+        # Generate the embedding
+        embeddings = self.embed_documents([text])
+        return embeddings

aiagents4pharma/talk2knowledgegraphs/utils/enrichments/__init__.py

@@ -3,3 +3,4 @@ This package contains modules to use the enrichment model
 """
 from . import enrichments
 from . import ollama
+from . import pubchem_strings

aiagents4pharma/talk2knowledgegraphs/utils/enrichments/pubchem_strings.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python3
+
+"""
+Enrichment class for enriching PubChem IDs with their STRINGS representation.
+"""
+
+from typing import List
+import pubchempy as pcp
+from .enrichments import Enrichments
+
+class EnrichmentWithPubChem(Enrichments):
+    """
+    Enrichment class using PubChem
+    """
+    def enrich_documents(self, texts: List[str]) -> List[str]:
+        """
+        Enrich a list of input PubChem IDs with their STRINGS representation.
+
+        Args:
+            texts: The list of pubchem IDs to be enriched.
+
+        Returns:
+            The list of enriched STRINGS
+        """
+
+        enriched_pubchem_ids = []
+        pubchem_cids = texts
+        for pubchem_cid in pubchem_cids:
+            try:
+                c = pcp.Compound.from_cid(pubchem_cid)
+            except pcp.BadRequestError:
+                enriched_pubchem_ids.append(None)
+                continue
+            enriched_pubchem_ids.append(c.isomeric_smiles)
+
+        return enriched_pubchem_ids
+
+    def enrich_documents_with_rag(self, texts, docs):
+        """
+        Enrich a list of input PubChem IDs with their STRINGS representation.
+
+        Args:
+            texts: The list of pubchem IDs to be enriched.
+            docs: None
+
+        Returns:
+            The list of enriched STRINGS
+        """
+        return self.enrich_documents(texts)

aiagents4pharma/talk2knowledgegraphs/utils/pubchem_utils.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+
+"""
+Enrichment class for enriching PubChem IDs with their STRINGS representation.
+"""
+
+import logging
+import requests
+import hydra
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def drugbank_id2pubchem_cid(drugbank_id):
+    """
+    Convert DrugBank ID to PubChem CID.
+
+    Args:
+        drugbank_id: The DrugBank ID of the drug.
+
+    Returns:
+        The PubChem CID of the drug.
+    """
+    logger.log(logging.INFO, "Load Hydra configuration for PubChem ID conversion.")
+    with hydra.initialize(version_base=None, config_path="../configs"):
+        cfg = hydra.compose(config_name='config',
+                            overrides=['utils/pubchem_utils=default'])
+        cfg = cfg.utils.pubchem_utils
+    # Prepare the URL
+    pubchem_url_for_drug = cfg.drugbank_id_to_pubchem_cid_url + drugbank_id + '/JSON'
+    # Get the data
+    response = requests.get(pubchem_url_for_drug, timeout=60)
+    data = response.json()
+    # Extract the PubChem CID
+    cid = None
+    for substance in data.get("PC_Substances", []):
+        for compound in substance.get("compound", []):
+            if "id" in compound and "type" in compound["id"] and compound["id"]["type"] == 1:
+                cid = compound["id"].get("id", {}).get("cid")
+                break
+    return cid

{aiagents4pharma-1.21.0.dist-info → aiagents4pharma-1.22.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: aiagents4pharma
-Version: 1.21.0
+Version: 1.22.0
 Summary: AI Agents for drug discovery, drug development, and other pharmaceutical R&D.
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License
@@ -30,6 +30,7 @@ Requires-Dist: ollama==0.4.6
 Requires-Dist: pandas==2.2.3
 Requires-Dist: pcst_fast==1.0.10
 Requires-Dist: plotly==5.24.1
+Requires-Dist: pubchempy==1.0.4
 Requires-Dist: pydantic==2.9.2
 Requires-Dist: pylint==3.3.1
 Requires-Dist: pypdf==5.2.0

{aiagents4pharma-1.21.0.dist-info → aiagents4pharma-1.22.0.dist-info}/RECORD

@@ -59,8 +59,8 @@ aiagents4pharma/talk2cells/tools/scp_agent/search_studies.py,sha256=MLe-twtFnOu-
 aiagents4pharma/talk2knowledgegraphs/__init__.py,sha256=Z0Eo7LTiKk0STsr8VI7wkCLq7PHrK1vYlH4I1hSNLiA,165
 aiagents4pharma/talk2knowledgegraphs/agents/__init__.py,sha256=iOAzuy_8A03tQDFtSBhC9dldUo62z5gfxcVtXAdLOJs,92
 aiagents4pharma/talk2knowledgegraphs/agents/t2kg_agent.py,sha256=j6MA1LB28mqpb6ZEmNLGcvDZvOnlGbJB9r7VXyEGask,3079
-aiagents4pharma/talk2knowledgegraphs/configs/__init__.py,sha256=Y49ucO22v9oe9EwFiXN6MU2wvyB3_ZBpmHwHbeh-ZVQ,106
-aiagents4pharma/talk2knowledgegraphs/configs/config.yaml,sha256=rwUIZ2t5j5hlFyre7VnV8zMsP0qpPTwvAFExgvQD6q0,196
+aiagents4pharma/talk2knowledgegraphs/configs/__init__.py,sha256=4_DVdpahaJ55yPl0aZotlFA_MYWLFF2cubWyKtBVI_Q,126
+aiagents4pharma/talk2knowledgegraphs/configs/config.yaml,sha256=bag4w3JCSqaojG37MTksy3ZehAPe3qoVzjIN2uh3nrc,229
 aiagents4pharma/talk2knowledgegraphs/configs/agents/t2kg_agent/__init__.py,sha256=-fAORvyFmG2iSvFOFDixmt9OTQRR58y89uhhu2EgbA8,46
 aiagents4pharma/talk2knowledgegraphs/configs/agents/t2kg_agent/default.yaml,sha256=ENCGROwYFpR6g4QD518h73sshdn3vPVpotBMk1QJcpU,4830
 aiagents4pharma/talk2knowledgegraphs/configs/app/__init__.py,sha256=fKfc3FR7g5KjY9b6jzrU6cwKTVVpkoVZQS3dvUowu34,69
@@ -91,26 +91,32 @@ aiagents4pharma/talk2knowledgegraphs/tests/test_tools_subgraph_extraction.py,sha
 aiagents4pharma/talk2knowledgegraphs/tests/test_tools_subgraph_summarization.py,sha256=oBqfspXXOxH04OQuPb8BCW0liIQTGKXtaPNSrPpQtFc,7597
 aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_embeddings.py,sha256=uYFoE_6zeU10_1mLLAHUr5c4S2XZMSc0Q_860o-KWEw,1517
 aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_huggingface.py,sha256=hzX84pheZdEsTtikF2KtBFiH44_xPjYXxLA6p4Ax1CY,1623
+aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_nim_molmim.py,sha256=LwtTZ-M7lHGxvRrGBXbyIT8AkA3T2OpeKqtNq3RK7Ik,2164
 aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_ollama.py,sha256=jn-TrPwF0aR9kVoerwkbMZa3U6Hc6HjV6Zoau4qSH4g,1834
 aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_sentencetransformer.py,sha256=Qxo6WeIDRy8aLh1tNKw0kSlzmUj3MtTak63oW2YwB24,1327
 aiagents4pharma/talk2knowledgegraphs/tests/test_utils_enrichments_enrichments.py,sha256=N6HRr4lWHXY7bTHe2uXJe4D_EG9WqZPibZne6qLl9_k,1447
 aiagents4pharma/talk2knowledgegraphs/tests/test_utils_enrichments_ollama.py,sha256=JhY7axvVULLywDJ2ctA-gob5YPeaJYWsaMNjHT6L9CU,3021
+aiagents4pharma/talk2knowledgegraphs/tests/test_utils_enrichments_pubchem.py,sha256=bk27KElJxOvKJ2RTz4ftleExQPMyWWS755KKmlImzbk,1241
 aiagents4pharma/talk2knowledgegraphs/tests/test_utils_kg_utils.py,sha256=pal76wi7WgQWUNk56BrzfFV8jKpbDaHHdbwtgx_gXLI,2410
+aiagents4pharma/talk2knowledgegraphs/tests/test_utils_pubchem_utils.py,sha256=C07YqUNYW7ofpKAzKh0lBovXKLvaiXFb3oJU6k1dvu4,411
 aiagents4pharma/talk2knowledgegraphs/tools/__init__.py,sha256=zpD4h7EYtyq0QNOqLd6bkxrPlPb2XN64ceI9ncgESrA,171
 aiagents4pharma/talk2knowledgegraphs/tools/graphrag_reasoning.py,sha256=OEuOFncDRdb7TQEGq4rkT5On-jI-R7Nt8K5EBzaND8w,5338
 aiagents4pharma/talk2knowledgegraphs/tools/load_arguments.py,sha256=zhmsRp-8vjB5rRekqTA07d3yb-42HWqng9dDMkvK6hM,623
 aiagents4pharma/talk2knowledgegraphs/tools/subgraph_extraction.py,sha256=te06QMFQfgJWrjaGrqpcOYeaV38jwm0KY_rXVSMHkeI,11468
 aiagents4pharma/talk2knowledgegraphs/tools/subgraph_summarization.py,sha256=mDSBOxopDfNhEJeU8fVI8b5lXTYrRzcc97aLbFgYSy4,4413
-aiagents4pharma/talk2knowledgegraphs/utils/__init__.py,sha256=Q9mzcSmkmhdnOn13fxGh1fNECYoUR5Y5CCuEJTIxwAI,167
+aiagents4pharma/talk2knowledgegraphs/utils/__init__.py,sha256=cZqb3LZLmBnmyAtWFv2Z-4uJvQmx0M4zKsfiWrlM3Pk,195
 aiagents4pharma/talk2knowledgegraphs/utils/kg_utils.py,sha256=6vQnPkeOWae_8jePjhma3sJuMTngy0I0tqzdFt6OqKg,2507
-aiagents4pharma/talk2knowledgegraphs/utils/embeddings/__init__.py,sha256=4TGK0XIVkkfGOyrSVwFQ-Lp-rzH9CCl-fWcqkFJKRLc,174
+aiagents4pharma/talk2knowledgegraphs/utils/pubchem_utils.py,sha256=IlrdGbRGD0IM7eMcpkOjuRjKNuH3lz_X8zN6RHwk61c,1340
+aiagents4pharma/talk2knowledgegraphs/utils/embeddings/__init__.py,sha256=POSDrSdFAWsBCueOPD-Fok-ARdTywJU1ivwpT9EU1Kw,199
 aiagents4pharma/talk2knowledgegraphs/utils/embeddings/embeddings.py,sha256=1nGznrAj-xT0xuSMBGz2dOujJ7M_IwSR84njxtxsy9A,2523
 aiagents4pharma/talk2knowledgegraphs/utils/embeddings/huggingface.py,sha256=2vi_elf6EgzfagFAO5QnL3a_aXZyN7B1EBziu44MTfM,3806
+aiagents4pharma/talk2knowledgegraphs/utils/embeddings/nim_molmim.py,sha256=XH6JNfmMS38UEU7UGJeeabHfRykharnQpQaqjO86OlQ,1537
 aiagents4pharma/talk2knowledgegraphs/utils/embeddings/ollama.py,sha256=8w0sjt3Ex5YJ_XvpKl9UbhdTiiaoMIarbPUxLBU-1Uw,2378
 aiagents4pharma/talk2knowledgegraphs/utils/embeddings/sentence_transformer.py,sha256=36iKlisOpMtGR5xfTAlSHXWvPqVC_Jbezod8kbBBMVg,2136
-aiagents4pharma/talk2knowledgegraphs/utils/enrichments/__init__.py,sha256=tW426knki2DBIHcWyF_K04iMMdbpIn_e_TpPmTgz2dI,113
+aiagents4pharma/talk2knowledgegraphs/utils/enrichments/__init__.py,sha256=JKGavA-umsGX3ng17_UYAvDBdbg-W-mPn8Q6JfP7J9U,143
 aiagents4pharma/talk2knowledgegraphs/utils/enrichments/enrichments.py,sha256=Bx8x6zzk5614ApWB90N_iv4_Y_Uq0-KwUeBwYSdQMU4,924
 aiagents4pharma/talk2knowledgegraphs/utils/enrichments/ollama.py,sha256=8eoxR-VHo0G7ReQIwje7xEhE-SJlHdef7_wJRpnvFIc,4116
+aiagents4pharma/talk2knowledgegraphs/utils/enrichments/pubchem_strings.py,sha256=qsVlDCGGDkUCv-R5_xFGhrtLS7P0CfagnM2qATwiOFM,1333
 aiagents4pharma/talk2knowledgegraphs/utils/extractions/__init__.py,sha256=7gwwtfzKhB8GuOBD47XRi0NprwEXkOzwNl5eeu-hDTI,86
 aiagents4pharma/talk2knowledgegraphs/utils/extractions/pcst.py,sha256=m5p0yoJb7I19ua5yeQfXPf7c4r6S1XPwttsrM7Qoy94,9336
 aiagents4pharma/talk2scholars/__init__.py,sha256=gphERyVKZHvOnMQsml7TIHlaIshHJ75R1J3FKExkfuY,120
@@ -152,8 +158,8 @@ aiagents4pharma/talk2scholars/tools/s2/query_results.py,sha256=EUfzRh5Qc_tMl5fDI
 aiagents4pharma/talk2scholars/tools/s2/retrieve_semantic_scholar_paper_id.py,sha256=Lg1L4HQCN2LaQEyWtLD73O67PMoXkPHi-Y8rCzHS0A4,2499
 aiagents4pharma/talk2scholars/tools/s2/search.py,sha256=mnBQWDuQ50UVw6B-bRuL8Ek1av-pEtdgzVMxpEA2BpI,4296
 aiagents4pharma/talk2scholars/tools/s2/single_paper_rec.py,sha256=xgnUj9W9JkeTvB2VJBJUAnia789GGNGqdqgJ_G16v2s,5120
-aiagents4pharma-1.21.0.dist-info/LICENSE,sha256=IcIbyB1Hyk5ZDah03VNQvJkbNk2hkBCDqQ8qtnCvB4Q,1077
-aiagents4pharma-1.21.0.dist-info/METADATA,sha256=YsjDHw3yfqfPClv0N3j35AObxHBhStDojFUUslyd_1Q,7757
-aiagents4pharma-1.21.0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-aiagents4pharma-1.21.0.dist-info/top_level.txt,sha256=-AH8rMmrSnJtq7HaAObS78UU-cTCwvX660dSxeM7a0A,16
-aiagents4pharma-1.21.0.dist-info/RECORD,,
+aiagents4pharma-1.22.0.dist-info/LICENSE,sha256=IcIbyB1Hyk5ZDah03VNQvJkbNk2hkBCDqQ8qtnCvB4Q,1077
+aiagents4pharma-1.22.0.dist-info/METADATA,sha256=qMCGdj4nouCcxKZOt2SwH53fDMIgPkoUTirEIvk2Mfs,7789
+aiagents4pharma-1.22.0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+aiagents4pharma-1.22.0.dist-info/top_level.txt,sha256=-AH8rMmrSnJtq7HaAObS78UU-cTCwvX660dSxeM7a0A,16
+aiagents4pharma-1.22.0.dist-info/RECORD,,