aiagents4pharma 1.24.0__py3-none-any.whl → 1.25.0__py3-none-any.whl

aiagents4pharma/__init__.py

@@ -7,4 +7,5 @@ from . import (
     talk2cells,
     talk2scholars,
     talk2knowledgegraphs,
+    talk2aiagents4pharma
 )

aiagents4pharma/talk2aiagents4pharma/__init__.py

@@ -0,0 +1,6 @@
+'''
+This file is used to import the models and tools.
+'''
+from . import agents
+from . import states
+from . import configs

aiagents4pharma/talk2aiagents4pharma/agents/__init__.py

@@ -0,0 +1,5 @@
+'''
+This file is used to import all the modules in the package.
+'''
+# import everything from the module
+from . import main_agent

aiagents4pharma/talk2aiagents4pharma/agents/main_agent.py

@@ -0,0 +1,52 @@
+#/usr/bin/env python3
+
+'''
+This is the main agent file for the AIAgents4Pharma.
+'''
+
+import logging
+import hydra
+from langgraph_supervisor import create_supervisor
+from langchain_openai import ChatOpenAI
+from langchain_core.language_models.chat_models import BaseChatModel
+from langgraph.checkpoint.memory import MemorySaver
+from ...talk2biomodels.agents.t2b_agent import get_app as get_app_t2b
+from ...talk2knowledgegraphs.agents.t2kg_agent import get_app as get_app_t2kg
+from ..states.state_talk2aiagents4pharma import Talk2AIAgents4Pharma
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def get_app(uniq_id,
+            llm_model: BaseChatModel = ChatOpenAI(model='gpt-4o-mini', temperature=0)):
+    '''
+    This function returns the langraph app.
+    '''
+    # Load hydra configuration
+    logger.log(logging.INFO, "Launching AIAgents4Pharma_Agent with thread_id %s", uniq_id)
+    with hydra.initialize(version_base=None, config_path="../configs"):
+        cfg = hydra.compose(config_name='config',
+                            overrides=['agents/main_agent=default'])
+        cfg = cfg.agents.main_agent
+    logger.log(logging.INFO, "System_prompt of T2AA4P: %s", cfg.system_prompt)
+    # Create supervisor workflow
+    workflow = create_supervisor(
+        [
+            get_app_t2b(uniq_id, llm_model),    # Talk2BioModels
+            get_app_t2kg(uniq_id, llm_model)    # Talk2KnowledgeGraphs
+        ],
+        model=llm_model,
+        state_schema=Talk2AIAgents4Pharma,
+        # Full history is needed to extract
+        # the tool artifacts
+        output_mode="full_history",
+        add_handoff_back_messages=False,
+        prompt=cfg.system_prompt
+    )
+
+    # Compile and run
+    app = workflow.compile(checkpointer=MemorySaver(),
+                           name="AIAgents4Pharma_Agent")
+
+    return app

aiagents4pharma/talk2aiagents4pharma/configs/__init__.py

@@ -0,0 +1,5 @@
+'''
+Import all the modules in the package
+'''
+
+from . import agents

aiagents4pharma/talk2aiagents4pharma/configs/agents/__init__.py

@@ -0,0 +1,5 @@
+'''
+Import all the modules in the package
+'''
+
+from . import main_agent

aiagents4pharma/talk2aiagents4pharma/configs/agents/main_agent/default.yaml

@@ -0,0 +1,12 @@
+_target_: agents.main_agent.get_app
+system_prompt: >
+  You are Talk2AIAgents4Pharma agent.
+  You are managing a team of the following 2 agents:
+   
+  1. Talk2Biomodels (T2B) agent: This agent can operate
+  on mathematical models of biological systems. This
+  agent can also query an uploaded document/pdf/article. 
+   
+  2. Talk2KnowledgeGraphs (T2KG) agent: This agent can
+  reason over a knowledge graph of biological entities
+  and their relationships.

aiagents4pharma/talk2aiagents4pharma/configs/config.yaml

@@ -0,0 +1,3 @@
+defaults:
+  - _self_
+  - agents/main_agent: default

aiagents4pharma/talk2aiagents4pharma/states/__init__.py

@@ -0,0 +1,4 @@
+'''
+This file is used to import all the models in the package.
+'''
+from . import state_talk2aiagents4pharma

aiagents4pharma/talk2aiagents4pharma/states/state_talk2aiagents4pharma.py

@@ -0,0 +1,16 @@
+"""
+This is the state file for the Talk2AIAgents4Pharma agent.
+"""
+
+from ...talk2biomodels.states.state_talk2biomodels import Talk2Biomodels
+from ...talk2knowledgegraphs.states.state_talk2knowledgegraphs import Talk2KnowledgeGraphs
+
+class Talk2AIAgents4Pharma(Talk2Biomodels,
+                           Talk2KnowledgeGraphs):
+    """
+    The state for the Talk2AIAgents4Pharma agent.
+
+    This class inherits from the classes:
+    1. Talk2Biomodels
+    2. Talk2KnowledgeGraphs
+    """

aiagents4pharma/talk2aiagents4pharma/tests/__init__.py

@@ -0,0 +1,3 @@
+'''
+This module contains the test cases.
+'''

aiagents4pharma/talk2aiagents4pharma/tests/test_main_agent.py

@@ -0,0 +1,111 @@
+'''
+Test Talk2AIAgents4Pharma supervisor agent.
+'''
+
+import pytest
+from langchain_core.messages import HumanMessage
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+from ..agents.main_agent import get_app
+
+# Define the data path for the test files of Talk2KnowledgeGraphs agent
+DATA_PATH = "aiagents4pharma/talk2knowledgegraphs/tests/files"
+
+@pytest.fixture(name="input_dict")
+def input_dict_fixture():
+    """
+    Input dictionary fixture for Talk2AIAgents4Pharma agent,
+    which is partly inherited from the Talk2KnowledgeGraphs agent.
+    """
+    input_dict = {
+        "topk_nodes": 3,
+        "topk_edges": 3,
+        "uploaded_files": [],
+        "dic_source_graph": [
+            {
+                "name": "PrimeKG",
+                "kg_pyg_path": f"{DATA_PATH}/primekg_ibd_pyg_graph.pkl",
+                "kg_text_path": f"{DATA_PATH}/primekg_ibd_text_graph.pkl",
+            }
+        ],
+        "dic_extracted_graph": []
+    }
+
+    return input_dict
+
+def test_main_agent_invokes_t2kg(input_dict):
+    """
+    In the following test, we will ask the main agent (supervisor)
+    to list drugs that target the gene Interleukin-6. We will check
+    if the Talk2KnowledgeGraphs agent is invoked. We will do so by
+    checking the state of the Talk2AIAgents4Pharma agent, which is
+    partly inherited from the Talk2KnowledgeGraphs agent
+
+    Args:
+        input_dict: Input dictionary
+    """
+    # Prepare LLM and embedding model
+    input_dict["llm_model"] = ChatOpenAI(model="gpt-4o-mini", temperature=0.0)
+    input_dict["embedding_model"] = OpenAIEmbeddings(model="text-embedding-3-small")
+
+    # Setup the app
+    unique_id = 12345
+    app = get_app(unique_id, llm_model=input_dict["llm_model"])
+    config = {"configurable": {"thread_id": unique_id}}
+    # Update state
+    app.update_state(
+        config,
+        input_dict,
+    )
+    prompt = "List drugs that target the gene Interleukin-6"
+
+    # Invoke the agent
+    response = app.invoke({"messages": [HumanMessage(content=prompt)]}, config=config)
+
+    # Check assistant message
+    assistant_msg = response["messages"][-1].content
+    assert isinstance(assistant_msg, str)
+
+    # Check extracted subgraph dictionary
+    current_state = app.get_state(config)
+    dic_extracted_graph = current_state.values["dic_extracted_graph"][0]
+    assert isinstance(dic_extracted_graph, dict)
+    assert dic_extracted_graph["graph_source"] == "PrimeKG"
+    assert dic_extracted_graph["topk_nodes"] == 3
+    assert dic_extracted_graph["topk_edges"] == 3
+    assert isinstance(dic_extracted_graph["graph_dict"], dict)
+    assert len(dic_extracted_graph["graph_dict"]["nodes"]) > 0
+    assert len(dic_extracted_graph["graph_dict"]["edges"]) > 0
+    assert isinstance(dic_extracted_graph["graph_text"], str)
+    # Check summarized subgraph
+    assert isinstance(dic_extracted_graph["graph_summary"], str)
+
+def test_main_agent_invokes_t2b():
+    '''
+    In the following test, we will ask the main agent (supervisor)
+    to simulate a model. And we will check if the Talk2BioModels
+    agent is invoked. We will do so by checking the state of the
+    Talk2AIAgents4Pharma agent, which is partly inherited from the
+    Talk2BioModels agent.
+    '''
+    unique_id = 123
+    app = get_app(unique_id)
+    config = {"configurable": {"thread_id": unique_id}}
+    prompt = "Simulate model 64"
+    # Invoke the agent
+    app.invoke(
+        {"messages": [HumanMessage(content=prompt)]},
+        config=config
+    )
+    # Get the state of the Talk2AIAgents4Pharma agent
+    current_state = app.get_state(config)
+    # Check if the dic_simulated_data is in the state
+    dic_simulated_data = current_state.values["dic_simulated_data"]
+    # Check if the dic_simulated_data is a list
+    assert isinstance(dic_simulated_data, list)
+    # Check if the length of the dic_simulated_data is 1
+    assert len(dic_simulated_data) == 1
+    # Check if the source of the model is 64
+    assert dic_simulated_data[0]['source'] == 64
+    # Check if the data of the model contains
+    # '1,3-bisphosphoglycerate'
+    assert '1,3-bisphosphoglycerate' in dic_simulated_data[0]['data']

aiagents4pharma/talk2biomodels/agents/t2b_agent.py

@@ -87,7 +87,8 @@ def get_app(uniq_id,
     # meaning you can use it as you would any other runnable.
     # Note that we're (optionally) passing the memory
     # when compiling the graph
-    app = workflow.compile(checkpointer=checkpointer)
+    app = workflow.compile(checkpointer=checkpointer,
+                           name="T2B_Agent")
     logger.log(logging.INFO,
                "Compiled the graph with thread_id %s and llm_model %s",
                uniq_id,

aiagents4pharma/talk2knowledgegraphs/agents/t2kg_agent.py

@@ -76,7 +76,8 @@ def get_app(uniq_id, llm_model: BaseChatModel=ChatOllama(model='llama3.2:1b', te
     # meaning you can use it as you would any other runnable.
     # Note that we're (optionally) passing the memory
     # when compiling the graph
-    app = workflow.compile(checkpointer=checkpointer)
+    app = workflow.compile(checkpointer=checkpointer,
+                           name="T2KG_Agent")
     logger.log(logging.INFO,
                "Compiled the graph with thread_id %s and llm_model %s",
                uniq_id,

{aiagents4pharma-1.24.0.dist-info → aiagents4pharma-1.25.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: aiagents4pharma
-Version: 1.24.0
+Version: 1.25.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
@@ -23,7 +23,7 @@ Requires-Dist: langchain-experimental==0.3.3
 Requires-Dist: langchain-nvidia-ai-endpoints==0.3.9
 Requires-Dist: langchain-openai==0.2.5
 Requires-Dist: langchain_ollama==0.2.3
-Requires-Dist: langgraph==0.2.66
+Requires-Dist: langgraph_supervisor==0.0.4
 Requires-Dist: matplotlib==3.9.2
 Requires-Dist: openai==1.59.4
 Requires-Dist: ollama==0.4.7
@@ -83,6 +83,7 @@ Our toolkit currently consists of the following agents:
 - **Talk2KnowledgeGraphs** _(v1 in progress)_: Access and explore complex biological knowledge graphs for insightful data connections.
 - **Talk2Scholars** _(v1 in progress)_: Get recommendations for articles related to your choice. Download, query, and write/retrieve them to your reference manager (currently supporting Zotero).
 - **Talk2Cells** _(v1 in progress)_: Query and analyze sequencing data with ease.
+- **Talk2AIAgents4Pharma** _(v1 in progress)_: Converse with all the agents above (currently supports T2B and T2KG)
 
 ![AIAgents4Pharma](docs/assets/AIAgents4Pharma.png)
 
@@ -104,49 +105,47 @@ Check out the tutorials on each agent for detailed instrcutions.
 
 _Both `Talk2Biomodels` and `Talk2Scholars` are now available on Docker Hub._
 
-#### **Running Talk2Biomodels**
+1. **Pull the Docker images**
 
-1. **Pull the Docker image**
    ```bash
    docker pull virtualpatientengine/talk2biomodels
    ```
-2. **Run the container**
+
+   ```bash
+   docker pull virtualpatientengine/talk2scholars
+   ```
+
+2. **Run the containers**
+
    ```bash
    docker run -d \
+     --name talk2biomodels \
      -e OPENAI_API_KEY=<your_openai_api_key> \
      -e NVIDIA_API_KEY=<your_nvidia_api_key> \
      -p 8501:8501 \
      virtualpatientengine/talk2biomodels
    ```
-3. **Access the Web App**  
-    Open your browser and go to:
-   ```
-   http://localhost:8501
-   ```
-   _You can create a free account at NVIDIA and apply for their
-   free credits [here](https://build.nvidia.com/explore/discover)._
-
-#### **Running Talk2Scholars**
 
-1. **Pull the Docker image**
-   ```bash
-   docker pull virtualpatientengine/talk2scholars
-   ```
-2. **Run the container**
    ```bash
    docker run -d \
+     --name talk2scholars \
      -e OPENAI_API_KEY=<your_openai_api_key> \
      -e ZOTERO_API_KEY=<your_zotero_api_key> \
      -e ZOTERO_USER_ID=<your_zotero_user_id> \
      -p 8501:8501 \
      virtualpatientengine/talk2scholars
    ```
+
 3. **Access the Web App**  
-   Open your browser and go to:
+    Open your browser and go to:
+
    ```
    http://localhost:8501
    ```
 
+   _You can create a free account at NVIDIA and apply for their
+   free credits [here](https://build.nvidia.com/explore/discover)._
+
 #### **Notes**
 
 - Ensure you **replace `<your_openai_api_key>`, `<your_nvidia_api_key>`, `<your_zotero_api_key>`, and `<your_zotero_user_id>`** with your actual credentials.
@@ -192,10 +191,10 @@ _Both `Talk2Biomodels` and `Talk2Scholars` are now available on Docker Hub._
    ```
 
    _Please note that this will create a new tracing project in your Langsmith
-   account with the name `T2X-xxxx`, where `X` can be `B` (Biomodels), `S` (Scholars),
-   `KG` (KnowledgeGraphs), or `C` (Cells). If you skip the previous step, it will
-   default to the name `default`. `xxxx` will be the 4-digit ID created for the
-   session._
+   account with the name `T2X-xxxx`, where `X` can be `AA4P` (Main Agent),
+   `B` (Biomodels), `S` (Scholars), `KG` (KnowledgeGraphs), or `C` (Cells).
+   If you skip the previous step, it will default to the name `default`.
+   `xxxx` will be the 4-digit ID created for the session._
 
 6. **Launch the app:**
    ```bash

{aiagents4pharma-1.24.0.dist-info → aiagents4pharma-1.25.0.dist-info}/RECORD

@@ -1,7 +1,18 @@
-aiagents4pharma/__init__.py,sha256=Ua9fqYW5gV1SZ0nOyOMd4T3wTlBui1-mrlJzFUQLFgY,161
+aiagents4pharma/__init__.py,sha256=7xnvthMuxYurECWvyb0_yaPJ18SsqvmkKPTCxgbQlNQ,186
+aiagents4pharma/talk2aiagents4pharma/__init__.py,sha256=KOPe8cZ-ROQ6EGn8FeejRFUPokayKRMTgiXFyOpZGoA,122
+aiagents4pharma/talk2aiagents4pharma/agents/__init__.py,sha256=Cc1RitlLGzJ5d_dCSUdguZH417nlJux1qVCVS2M9t30,129
+aiagents4pharma/talk2aiagents4pharma/agents/main_agent.py,sha256=oVHYFXHmwq8y6e41JvmRYDXvCjItqY7setpcC2U2y6Q,1870
+aiagents4pharma/talk2aiagents4pharma/configs/__init__.py,sha256=5ah__-8XyRblwT0U1ByRigNjt_GyCheu7zce4aM-eZE,68
+aiagents4pharma/talk2aiagents4pharma/configs/config.yaml,sha256=VnbMbVSYfCh68cHZ0JLu00UjOUmapejN3EsN3lnBXtU,51
+aiagents4pharma/talk2aiagents4pharma/configs/agents/__init__.py,sha256=zrJcq-4m0YUKfSlRGC8KzBmEooaASKuL_Y75yDp-ZoA,72
+aiagents4pharma/talk2aiagents4pharma/configs/agents/main_agent/default.yaml,sha256=CRi4t2VGYfDouDYb6id5_qLm-6hLkp5WHCZJwsJlYM0,455
+aiagents4pharma/talk2aiagents4pharma/states/__init__.py,sha256=3wSvCpM29oqvVjhbhabm7FNm9Zt0rHO5tEn63YW6doc,108
+aiagents4pharma/talk2aiagents4pharma/states/state_talk2aiagents4pharma.py,sha256=NxujEBDKubvpV9UG7ERTDRB6psr0XnObCNHyztLAhgo,485
+aiagents4pharma/talk2aiagents4pharma/tests/__init__.py,sha256=Jbw5tJxSrjGoaK5IX3pJWDCNzhrVQ10lkYq2oQ_KQD8,45
+aiagents4pharma/talk2aiagents4pharma/tests/test_main_agent.py,sha256=H8jcSZlzyt0wkijF_hgX8p63moXejbdP_W2xrWGTu3g,4122
 aiagents4pharma/talk2biomodels/__init__.py,sha256=1cq1HX2xoi_a0nDPuXYoSTrnL26OHQBW3zXNwwwjFO0,181
 aiagents4pharma/talk2biomodels/agents/__init__.py,sha256=sn5-fREjMdEvb-OUan3iOqrgYGjplNx3J8hYOaW0Po8,128
-aiagents4pharma/talk2biomodels/agents/t2b_agent.py,sha256=bbktkTinp5VHHmAZEs1yMPN3Bu7rnyxRq2Cb7xVe8Gw,3477
+aiagents4pharma/talk2biomodels/agents/t2b_agent.py,sha256=1NxSQbh5wbDpnap4mHwzSIjjHP1xakE2d1Pg5bHtDrE,3522
 aiagents4pharma/talk2biomodels/api/__init__.py,sha256=_GmDQqDLYpsUPUeE1nBNlT5AI9oTXIcqgOfNfvmonqA,123
 aiagents4pharma/talk2biomodels/api/kegg.py,sha256=QzYDAfJ16E7tbHGxP8ZNWRizMkMRS_HJuucueXEC1Gg,2943
 aiagents4pharma/talk2biomodels/api/ols.py,sha256=qq0Qy-gJDxanQW-HfCChDsTQsY1M41ua8hMlTnfuzrA,2202
@@ -62,7 +73,7 @@ aiagents4pharma/talk2cells/tools/scp_agent/display_studies.py,sha256=6q59gh_NQai
 aiagents4pharma/talk2cells/tools/scp_agent/search_studies.py,sha256=MLe-twtFnOu-P8P9diYq7jvHBHbWFRRCZLcfpUzqPMg,2806
 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/agents/t2kg_agent.py,sha256=GdogSyTJa0LlSWPqDWdO41EXcS-PThatblPWpuHih-M,3125
 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
@@ -175,8 +186,8 @@ aiagents4pharma/talk2scholars/tools/s2/search.py,sha256=i5KMFJWK31CjYtVT1McJpLzg
 aiagents4pharma/talk2scholars/tools/s2/single_paper_rec.py,sha256=7PoZfcstxDThWX6NYOgxN_9M_nwgMPAALch8OmjraVY,5568
 aiagents4pharma/talk2scholars/tools/zotero/__init__.py,sha256=1UW4r5ECvAwYpo1Fjf7lQPO--M8I85baYCHocFOAq4M,53
 aiagents4pharma/talk2scholars/tools/zotero/zotero_read.py,sha256=NJ65fAJ4u2Zq15uvEajVOhI4QnNvyqA6FHPaEDqvMw0,4321
-aiagents4pharma-1.24.0.dist-info/LICENSE,sha256=IcIbyB1Hyk5ZDah03VNQvJkbNk2hkBCDqQ8qtnCvB4Q,1077
-aiagents4pharma-1.24.0.dist-info/METADATA,sha256=N3OSAZ9173G-Lh37d02tKfhqeCoQ8Q6yj8MtbNet0IM,9474
-aiagents4pharma-1.24.0.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
-aiagents4pharma-1.24.0.dist-info/top_level.txt,sha256=-AH8rMmrSnJtq7HaAObS78UU-cTCwvX660dSxeM7a0A,16
-aiagents4pharma-1.24.0.dist-info/RECORD,,
+aiagents4pharma-1.25.0.dist-info/LICENSE,sha256=IcIbyB1Hyk5ZDah03VNQvJkbNk2hkBCDqQ8qtnCvB4Q,1077
+aiagents4pharma-1.25.0.dist-info/METADATA,sha256=GcI8CTr7TUVIlC35PVwhhBi-5LEn6bTd2r3FVhtHCC4,9469
+aiagents4pharma-1.25.0.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
+aiagents4pharma-1.25.0.dist-info/top_level.txt,sha256=-AH8rMmrSnJtq7HaAObS78UU-cTCwvX660dSxeM7a0A,16
+aiagents4pharma-1.25.0.dist-info/RECORD,,