aiagents4pharma 1.8.3__tar.gz → 1.10.0__tar.gz

{aiagents4pharma-1.8.3 → aiagents4pharma-1.10.0}/PKG-INFO

@@ -1,11 +1,11 @@
 Metadata-Version: 2.2
 Name: aiagents4pharma
-Version: 1.8.3
+Version: 1.10.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
 Classifier: Operating System :: OS Independent
-Requires-Python: >=3.10
+Requires-Python: >=3.12
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: copasi_basico==0.78
@@ -48,6 +48,9 @@ Requires-Dist: streamlit-feedback
 [![Talk2Cells](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2cells.yml/badge.svg)](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2cells.yml)
 [![Talk2KnowledgeGraphs](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2knowledgegraphs.yml/badge.svg)](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2knowledgegraphs.yml)
 [![Talk2Competitors](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2competitors.yml/badge.svg)](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2competitors.yml)
+![GitHub Release](https://img.shields.io/github/v/release/VirtualPatientEngine/AIAgents4Pharma)
+![Python Version from PEP 621 TOML](https://img.shields.io/python/required-version-toml?tomlFilePath=https%3A%2F%2Fraw.githubusercontent.com%2FVirtualPatientEngine%2FAIAgents4Pharma%2Frefs%2Fheads%2Fmain%2Fpyproject.toml)
+
 
 <h1 align="center" style="border-bottom: none;">🤖 AIAgents4Pharma</h1>
 
@@ -56,9 +59,9 @@ Welcome to **AIAgents4Pharma** – an open-source project by [Team VPE](https://
 Our toolkit currently consists of three intelligent agents, each designed to simplify and enhance access to specialized data in biology:
 
 - **Talk2BioModels**: Engage directly with mathematical models in systems biology.
-- **Talk2Cells** *(Work in progress)*: Query and analyze sequencing data with ease.
-- **Talk2KnowledgeGraphs** *(Work in progress)*: Access and explore complex biological knowledge graphs for insightful data connections.
-- **Talk2Competitors** *(Coming soon)*: Get recommendations for articles related to your choice. Download, query, and write/retrieve them to your reference manager (currently supporting Zotero).
+- **Talk2Cells** _(Work in progress)_: Query and analyze sequencing data with ease.
+- **Talk2KnowledgeGraphs** _(Work in progress)_: Access and explore complex biological knowledge graphs for insightful data connections.
+- **Talk2Competitors** _(Coming soon)_: Get recommendations for articles related to your choice. Download, query, and write/retrieve them to your reference manager (currently supporting Zotero).
 
 ---
 
@@ -72,15 +75,15 @@ Our toolkit currently consists of three intelligent agents, each designed to sim
 - Adjust parameters within the model to simulate different conditions.
 - Query simulation results.
 
-### 2. Talk2Cells *(Work in Progress)*
+### 2. Talk2Cells _(Work in Progress)_
 
 **Talk2Cells** is being developed to provide direct access to and analysis of sequencing data, such as RNA-Seq or DNA-Seq, using natural language.
 
-### 3. Talk2KnowledgeGraphs *(Work in Progress)*
+### 3. Talk2KnowledgeGraphs _(Work in Progress)_
 
 **Talk2KnowledgeGraphs** is an agent designed to enable interaction with biological knowledge graphs (KGs). KGs integrate vast amounts of structured biological data into a format that highlights relationships between entities, such as proteins, genes, and diseases.
 
-### 4. Talk2KnowledgeGraphs *(Coming soon)*
+### 4. Talk2Competitors _(Coming soon)_
 
 ## Getting Started
 
@@ -91,48 +94,60 @@ Our toolkit currently consists of three intelligent agents, each designed to sim
 - Required libraries specified in `requirements.txt`
 
 ### Installation
+
 #### Option 1: PyPI
-   ```bash
-   pip install aiagents4pharma
-   ```
+
+```bash
+pip install aiagents4pharma
+```
 
 Check out the tutorials on each agent for detailed instrcutions.
 
 #### Option 2: git
+
 1. **Clone the repository:**
+
    ```bash
    git clone https://github.com/VirtualPatientEngine/AIAgents4Pharma
    cd AIAgents4Pharma
    ```
 
 2. **Install dependencies:**
+
    ```bash
    pip install .
    ```
 
 3. **Initialize OPENAI_API_KEY**
+
    ```bash
-   export OPENAI_API_KEY = ....
+   export OPENAI_API_KEY=....
    ```
 
 4. **[Optional] Set up login credentials**
+
    ```bash
    vi .streamlit/secrets.toml
    ```
+
    and enter
+
    ```
    password='XXX'
    ```
-   Please note that the passowrd will be same for all the users.
+
+   Please note that the passoword will be same for all the users.
 
 5. **[Optional] Initialize LANGSMITH_API_KEY**
+
    ```bash
    export LANGCHAIN_TRACING_V2=true
    export LANGCHAIN_API_KEY=<your-api-key>
    ```
-   Please note that this will create a new tracing project in your Langsmith 
-   account with the name `<user_name>@<uuid>`, where `user_name` is the name 
-   you provided in the previous step. If you skip the previous step, it will 
+
+   Please note that this will create a new tracing project in your Langsmith
+   account with the name `<user_name>@<uuid>`, where `user_name` is the name
+   you provided in the previous step. If you skip the previous step, it will
    default to `default`. <uuid> will be the 128 bit unique ID created for the
    session.
 
@@ -164,6 +179,7 @@ We welcome contributions to AIAgents4Pharma! Here’s how you can help:
 5. **Open a pull request**
 
 ### Current Needs
+
 - **Beta testers** for Talk2BioModels.
 - **Developers** with experience in natural language processing, bioinformatics, or knowledge graphs for contributions to AIAgents4Pharma.
 
@@ -174,19 +190,22 @@ Check out our [CONTRIBUTING.md](CONTRIBUTING.md) for more information.
 ## Roadmap
 
 ### Completed
+
 - **Talk2BioModels**: Initial release with core capabilities for interacting with systems biology models.
 
 ### Planned
+
 - **User Interface**: Interactive web UI for all agents.
 - **Talk2Cells**: Integration of sequencing data analysis tools.
 - **Talk2KnowledgeGraphs**: Interface for biological knowledge graph interaction.
-- **Talk2Competitors**
+- **Talk2Competitors**: Interface for exploring articles
 
-We’re excited to bring AIAgents4Pharma to the bioinformatics and pharmaceutical research community. Together, let’s make data-driven biological research more accessible and insightful. 
+We’re excited to bring AIAgents4Pharma to the bioinformatics and pharmaceutical research community. Together, let’s make data-driven biological research more accessible and insightful.
 
 **Get Started** with AIAgents4Pharma today and transform the way you interact with biological data.
 
 ---
 
 ## Feedback
+
 Questions/Bug reports/Feature requests/Comments/Suggestions? We welcome all. Please use the `Isssues` tab 😀

{aiagents4pharma-1.8.3 → aiagents4pharma-1.10.0}/README.md

@@ -2,6 +2,9 @@
 [![Talk2Cells](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2cells.yml/badge.svg)](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2cells.yml)
 [![Talk2KnowledgeGraphs](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2knowledgegraphs.yml/badge.svg)](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2knowledgegraphs.yml)
 [![Talk2Competitors](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2competitors.yml/badge.svg)](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2competitors.yml)
+![GitHub Release](https://img.shields.io/github/v/release/VirtualPatientEngine/AIAgents4Pharma)
+![Python Version from PEP 621 TOML](https://img.shields.io/python/required-version-toml?tomlFilePath=https%3A%2F%2Fraw.githubusercontent.com%2FVirtualPatientEngine%2FAIAgents4Pharma%2Frefs%2Fheads%2Fmain%2Fpyproject.toml)
+
 
 <h1 align="center" style="border-bottom: none;">🤖 AIAgents4Pharma</h1>
 
@@ -10,9 +13,9 @@ Welcome to **AIAgents4Pharma** – an open-source project by [Team VPE](https://
 Our toolkit currently consists of three intelligent agents, each designed to simplify and enhance access to specialized data in biology:
 
 - **Talk2BioModels**: Engage directly with mathematical models in systems biology.
-- **Talk2Cells** *(Work in progress)*: Query and analyze sequencing data with ease.
-- **Talk2KnowledgeGraphs** *(Work in progress)*: Access and explore complex biological knowledge graphs for insightful data connections.
-- **Talk2Competitors** *(Coming soon)*: Get recommendations for articles related to your choice. Download, query, and write/retrieve them to your reference manager (currently supporting Zotero).
+- **Talk2Cells** _(Work in progress)_: Query and analyze sequencing data with ease.
+- **Talk2KnowledgeGraphs** _(Work in progress)_: Access and explore complex biological knowledge graphs for insightful data connections.
+- **Talk2Competitors** _(Coming soon)_: Get recommendations for articles related to your choice. Download, query, and write/retrieve them to your reference manager (currently supporting Zotero).
 
 ---
 
@@ -26,15 +29,15 @@ Our toolkit currently consists of three intelligent agents, each designed to sim
 - Adjust parameters within the model to simulate different conditions.
 - Query simulation results.
 
-### 2. Talk2Cells *(Work in Progress)*
+### 2. Talk2Cells _(Work in Progress)_
 
 **Talk2Cells** is being developed to provide direct access to and analysis of sequencing data, such as RNA-Seq or DNA-Seq, using natural language.
 
-### 3. Talk2KnowledgeGraphs *(Work in Progress)*
+### 3. Talk2KnowledgeGraphs _(Work in Progress)_
 
 **Talk2KnowledgeGraphs** is an agent designed to enable interaction with biological knowledge graphs (KGs). KGs integrate vast amounts of structured biological data into a format that highlights relationships between entities, such as proteins, genes, and diseases.
 
-### 4. Talk2KnowledgeGraphs *(Coming soon)*
+### 4. Talk2Competitors _(Coming soon)_
 
 ## Getting Started
 
@@ -45,48 +48,60 @@ Our toolkit currently consists of three intelligent agents, each designed to sim
 - Required libraries specified in `requirements.txt`
 
 ### Installation
+
 #### Option 1: PyPI
-   ```bash
-   pip install aiagents4pharma
-   ```
+
+```bash
+pip install aiagents4pharma
+```
 
 Check out the tutorials on each agent for detailed instrcutions.
 
 #### Option 2: git
+
 1. **Clone the repository:**
+
    ```bash
    git clone https://github.com/VirtualPatientEngine/AIAgents4Pharma
    cd AIAgents4Pharma
    ```
 
 2. **Install dependencies:**
+
    ```bash
    pip install .
    ```
 
 3. **Initialize OPENAI_API_KEY**
+
    ```bash
-   export OPENAI_API_KEY = ....
+   export OPENAI_API_KEY=....
    ```
 
 4. **[Optional] Set up login credentials**
+
    ```bash
    vi .streamlit/secrets.toml
    ```
+
    and enter
+
    ```
    password='XXX'
    ```
-   Please note that the passowrd will be same for all the users.
+
+   Please note that the passoword will be same for all the users.
 
 5. **[Optional] Initialize LANGSMITH_API_KEY**
+
    ```bash
    export LANGCHAIN_TRACING_V2=true
    export LANGCHAIN_API_KEY=<your-api-key>
    ```
-   Please note that this will create a new tracing project in your Langsmith 
-   account with the name `<user_name>@<uuid>`, where `user_name` is the name 
-   you provided in the previous step. If you skip the previous step, it will 
+
+   Please note that this will create a new tracing project in your Langsmith
+   account with the name `<user_name>@<uuid>`, where `user_name` is the name
+   you provided in the previous step. If you skip the previous step, it will
    default to `default`. <uuid> will be the 128 bit unique ID created for the
    session.
 
@@ -118,6 +133,7 @@ We welcome contributions to AIAgents4Pharma! Here’s how you can help:
 5. **Open a pull request**
 
 ### Current Needs
+
 - **Beta testers** for Talk2BioModels.
 - **Developers** with experience in natural language processing, bioinformatics, or knowledge graphs for contributions to AIAgents4Pharma.
 
@@ -128,19 +144,22 @@ Check out our [CONTRIBUTING.md](CONTRIBUTING.md) for more information.
 ## Roadmap
 
 ### Completed
+
 - **Talk2BioModels**: Initial release with core capabilities for interacting with systems biology models.
 
 ### Planned
+
 - **User Interface**: Interactive web UI for all agents.
 - **Talk2Cells**: Integration of sequencing data analysis tools.
 - **Talk2KnowledgeGraphs**: Interface for biological knowledge graph interaction.
-- **Talk2Competitors**
+- **Talk2Competitors**: Interface for exploring articles
 
-We’re excited to bring AIAgents4Pharma to the bioinformatics and pharmaceutical research community. Together, let’s make data-driven biological research more accessible and insightful. 
+We’re excited to bring AIAgents4Pharma to the bioinformatics and pharmaceutical research community. Together, let’s make data-driven biological research more accessible and insightful.
 
 **Get Started** with AIAgents4Pharma today and transform the way you interact with biological data.
 
 ---
 
 ## Feedback
+
 Questions/Bug reports/Feature requests/Comments/Suggestions? We welcome all. Please use the `Isssues` tab 😀

aiagents4pharma-1.10.0/aiagents4pharma/__init__.py

@@ -0,0 +1,11 @@
+"""
+This file is used to import aiagents4pharma modules.
+"""
+
+from . import (
+    configs,
+    talk2biomodels,
+    talk2cells,
+    talk2competitors,
+    talk2knowledgegraphs,
+)

{aiagents4pharma-1.8.3 → aiagents4pharma-1.10.0}/aiagents4pharma/configs/talk2biomodels/agents/t2b_agent/default.yaml

@@ -3,4 +3,6 @@ state_modifier: >
   You are Talk2BioModels agent.
   If the user asks for the uploaded model, 
   then pass the use_uploaded_model argument
-    as True.
+  as True. If the user asks for simulation, 
+  then suggest a value for the `simulation_name`
+  argument.

{aiagents4pharma-1.8.3 → aiagents4pharma-1.10.0}/aiagents4pharma/talk2biomodels/__init__.py

@@ -4,4 +4,4 @@ This file is used to import the models and tools.
 from . import models
 from . import tools
 from . import agents
-from . import states
+from . import states

{aiagents4pharma-1.8.3 → aiagents4pharma-1.10.0}/aiagents4pharma/talk2biomodels/states/state_talk2biomodels.py

@@ -20,5 +20,5 @@ class Talk2Biomodels(AgentState):
     # the operator for the sbml_file_path field.
     # https://langchain-ai.github.io/langgraph/troubleshooting/errors/INVALID_CONCURRENT_GRAPH_UPDATE/
     sbml_file_path: Annotated[list, operator.add]
-    dic_simulated_data: dict
+    dic_simulated_data: Annotated[list[dict], operator.add]
     llm_model: str

{aiagents4pharma-1.8.3 → aiagents4pharma-1.10.0}/aiagents4pharma/talk2biomodels/tests/test_langgraph.py

@@ -1,7 +1,8 @@
 '''
-Test cases
+Test cases for Talk2Biomodels.
 '''
 
+import pandas as pd
 from langchain_core.messages import HumanMessage, ToolMessage
 from ..agents.t2b_agent import get_app
 
@@ -13,7 +14,8 @@ def test_get_modelinfo_tool():
     app = get_app(unique_id)
     config = {"configurable": {"thread_id": unique_id}}
     # Update state
-    app.update_state(config,{"sbml_file_path": ["BIOMD0000000537.xml"]})
+    app.update_state(config,
+      {"sbml_file_path": ["aiagents4pharma/talk2biomodels/tests/BIOMD0000000449_url.xml"]})
     prompt = "Extract all relevant information from the uploaded model."
     # Test the tool get_modelinfo
     response = app.invoke(
@@ -56,26 +58,70 @@ def test_ask_question_tool():
 
     ##########################################
     # Test ask_question tool when simulation
-    # results are not available
+    # results are not available i.e. the
+    # simulation has not been run. In this
+    # case, the tool should return an error
     ##########################################
     # Update state
     app.update_state(config, {"llm_model": "gpt-4o-mini"})
+    # Define the prompt
     prompt = "Call the ask_question tool to answer the "
     prompt += "question: What is the concentration of CRP "
-    prompt += "in serum at 1000 hours?"
-
-    # Test the tool get_modelinfo
-    response = app.invoke(
-                        {"messages": [HumanMessage(content=prompt)]},
-                        config=config
-                    )
-    assistant_msg = response["messages"][-1].content
-    # Check if the assistant message is a string
-    assert isinstance(assistant_msg, str)
+    prompt += "in serum at 1000 hours? The simulation name "
+    prompt += "is `simulation_name`."
+    # Invoke the tool
+    app.invoke(
+            {"messages": [HumanMessage(content=prompt)]},
+            config=config
+        )
+    # Get the messages from the current state
+    # and reverse the order
+    current_state = app.get_state(config)
+    reversed_messages = current_state.values["messages"][::-1]
+    # Loop through the reversed messages until a
+    # ToolMessage is found.
+    for msg in reversed_messages:
+        # Assert that the message is a ToolMessage
+        # and its status is "error"
+        if isinstance(msg, ToolMessage):
+            assert msg.status == "error"
 
 def test_simulate_model_tool():
     '''
-    Test the simulate_model tool.
+    Test the simulate_model tool when simulating
+    multiple models.
+    '''
+    unique_id = 123
+    app = get_app(unique_id)
+    config = {"configurable": {"thread_id": unique_id}}
+    app.update_state(config, {"llm_model": "gpt-4o-mini"})
+    # Upload a model to the state
+    app.update_state(config,
+        {"sbml_file_path": ["aiagents4pharma/talk2biomodels/tests/BIOMD0000000449_url.xml"]})
+    prompt = "Simulate models 64 and the uploaded model"
+    # Invoke the agent
+    app.invoke(
+        {"messages": [HumanMessage(content=prompt)]},
+        config=config
+    )
+    current_state = app.get_state(config)
+    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 2
+    assert len(dic_simulated_data) == 2
+    # Check if the source of the first model is 64
+    assert dic_simulated_data[0]['source'] == 64
+    # Check if the source of the second model is upload
+    assert dic_simulated_data[1]['source'] == "upload"
+    # Check if the data of the first model contains
+    assert '1,3-bisphosphoglycerate' in dic_simulated_data[0]['data']
+    # Check if the data of the second model contains
+    assert 'mTORC2' in dic_simulated_data[1]['data']
+
+def test_integration():
+    '''
+    Test the integration of the tools.
     '''
     unique_id = 123
     app = get_app(unique_id)
@@ -138,9 +184,9 @@ def test_simulate_model_tool():
     reversed_messages = current_state.values["messages"][::-1]
     # Loop through the reversed messages
     # until a ToolMessage is found.
-    expected_artifact = ['CRP[serum]', 'CRPExtracellular']
-    expected_artifact += ['CRP Suppression (%)', 'CRP (% of baseline)']
-    expected_artifact += ['CRP[liver]']
+    expected_header = ['Time', 'CRP[serum]', 'CRPExtracellular']
+    expected_header += ['CRP Suppression (%)', 'CRP (% of baseline)']
+    expected_header += ['CRP[liver]']
     predicted_artifact = []
     for msg in reversed_messages:
         if isinstance(msg, ToolMessage):
@@ -150,9 +196,14 @@ def test_simulate_model_tool():
             if msg.name == "custom_plotter":
                 predicted_artifact = msg.artifact
                 break
-    # Check if the two artifacts are equal
-    # assert expected_artifact in predicted_artifact
-    assert set(expected_artifact).issubset(set(predicted_artifact))
+    # Convert the artifact into a pandas dataframe
+    # for easy comparison
+    df = pd.DataFrame(predicted_artifact)
+    # Extract the headers from the dataframe
+    predicted_header = df.columns.tolist()
+    # Check if the header is in the expected_header
+    # assert expected_header in predicted_artifact
+    assert set(expected_header).issubset(set(predicted_header))
     ##########################################
     # Test custom_plotter tool when the
     # simulation results are available but

{aiagents4pharma-1.8.3 → aiagents4pharma-1.10.0}/aiagents4pharma/talk2biomodels/tools/ask_question.py

@@ -23,6 +23,8 @@ class AskQuestionInput(BaseModel):
     Input schema for the AskQuestion tool.
     """
     question: str = Field(description="question about the simulation results")
+    simulation_name: str = Field(description="""Name assigned to the simulation
+                                 when the tool simulate_model was invoked.""")
     state: Annotated[dict, InjectedState]
 
 # Note: It's important that every field has type hints.
@@ -39,6 +41,7 @@ class AskQuestionTool(BaseTool):
 
     def _run(self,
              question: str,
+             simulation_name: str,
              state: Annotated[dict, InjectedState]) -> str:
         """
         Run the tool.
@@ -46,18 +49,24 @@ class AskQuestionTool(BaseTool):
         Args:
             question (str): The question to ask about the simulation results.
             state (dict): The state of the graph.
-            run_manager (Optional[CallbackManagerForToolRun]): The CallbackManagerForToolRun object.
+            simulation_name (str): The name assigned to the simulation.
 
         Returns:
             str: The answer to the question.
         """
         logger.log(logging.INFO,
-                   "Calling ask_question tool %s", question)
-        # Check if the simulation results are available
-        if 'dic_simulated_data' not in state:
-            return "Please run the simulation first before \
-                asking a question about the simulation results."
-        df = pd.DataFrame.from_dict(state['dic_simulated_data'])
+                   "Calling ask_question tool %s, %s", question, simulation_name)
+        dic_simulated_data = {}
+        for data in state["dic_simulated_data"]:
+            for key in data:
+                if key not in dic_simulated_data:
+                    dic_simulated_data[key] = []
+                dic_simulated_data[key] += [data[key]]
+        # print (dic_simulated_data)
+        df_simulated_data = pd.DataFrame.from_dict(dic_simulated_data)
+        df = pd.DataFrame(
+                df_simulated_data[df_simulated_data['name'] == simulation_name]['data'].iloc[0]
+                )
         prompt_content = None
         # if run_manager and 'prompt' in run_manager.metadata:
         #     prompt_content = run_manager.metadata['prompt']

{aiagents4pharma-1.8.3 → aiagents4pharma-1.10.0}/aiagents4pharma/talk2biomodels/tools/custom_plotter.py

@@ -6,14 +6,11 @@ Tool for plotting a custom figure.
 
 import logging
 from typing import Type, List, TypedDict, Annotated, Tuple, Union, Literal
-from typing import Type, List, TypedDict, Annotated, Tuple, Union, Literal
 from pydantic import BaseModel, Field
 import pandas as pd
-import pandas as pd
 from langchain_openai import ChatOpenAI
 from langchain_core.tools import BaseTool
 from langgraph.prebuilt import InjectedState
-from langgraph.prebuilt import InjectedState
 
 # Initialize logger
 logging.basicConfig(level=logging.INFO)
@@ -24,7 +21,7 @@ class CustomPlotterInput(BaseModel):
     Input schema for the PlotImage tool.
     """
     question: str = Field(description="Description of the plot")
-    state: Annotated[dict, InjectedState]
+    simulation_name: str = Field(description="Name assigned to the simulation")
     state: Annotated[dict, InjectedState]
 
 # Note: It's important that every field has type hints.
@@ -41,10 +38,10 @@ class CustomPlotterTool(BaseTool):
     description: str = "A tool to make custom plots of the simulation results"
     args_schema: Type[BaseModel] = CustomPlotterInput
     response_format: str = "content_and_artifact"
-    response_format: str = "content_and_artifact"
 
     def _run(self,
              question: str,
+             simulation_name: str,
              state: Annotated[dict, InjectedState]
              ) -> Tuple[str, Union[None, List[str]]]:
         """
@@ -53,17 +50,24 @@ class CustomPlotterTool(BaseTool):
         Args:
             question (str): The question about the custom plot.
             state (dict): The state of the graph.
-            question (str): The question about the custom plot.
-            state (dict): The state of the graph.
 
         Returns:
             str: The answer to the question
         """
         logger.log(logging.INFO, "Calling custom_plotter tool %s", question)
-        # Check if the simulation results are available
-        # if 'dic_simulated_data' not in state:
-        #     return "Please run the simulation first before plotting the figure.", None
-        df = pd.DataFrame.from_dict(state['dic_simulated_data'])
+        dic_simulated_data = {}
+        for data in state["dic_simulated_data"]:
+            for key in data:
+                if key not in dic_simulated_data:
+                    dic_simulated_data[key] = []
+                dic_simulated_data[key] += [data[key]]
+        # Create a pandas dataframe from the dictionary
+        df = pd.DataFrame.from_dict(dic_simulated_data)
+        # Get the simulated data for the current tool call
+        df = pd.DataFrame(
+                df[df['name'] == simulation_name]['data'].iloc[0]
+                )
+        # df = pd.DataFrame.from_dict(state['dic_simulated_data'])
         species_names = df.columns.tolist()
         # Exclude the time column
         species_names.remove('Time')
@@ -76,7 +80,8 @@ class CustomPlotterTool(BaseTool):
             A list of species based on user question.
             """
             relevant_species: Union[None, List[Literal[*species_names]]] = Field(
-                    description="List of species based on user question. If no relevant species are found, it will be None.")
+                    description="""List of species based on user question.
+                    If no relevant species are found, it will be None.""")
         # Create an instance of the LLM model
         llm = ChatOpenAI(model=state['llm_model'], temperature=0)
         llm_with_structured_output = llm.with_structured_output(CustomHeader)
@@ -90,5 +95,6 @@ class CustomPlotterTool(BaseTool):
         logger.info("Extracted species: %s", extracted_species)
         if len(extracted_species) == 0:
             return "No species found in the simulation results that matches the user prompt.", None
-        content = f"Plotted custom figure with species: {', '.join(extracted_species)}"
-        return content, extracted_species
+        # Include the time column
+        extracted_species.insert(0, 'Time')
+        return f"Custom plot {simulation_name}", df[extracted_species].to_dict(orient='records')

{aiagents4pharma-1.8.3 → aiagents4pharma-1.10.0}/aiagents4pharma/talk2biomodels/tools/get_modelinfo.py

@@ -25,12 +25,12 @@ class RequestedModelInfo:
     """
     Dataclass for storing the requested model information.
     """
-    species: bool = Field(description="Get species from the model.")
-    parameters: bool = Field(description="Get parameters from the model.")
-    compartments: bool = Field(description="Get compartments from the model.")
-    units: bool = Field(description="Get units from the model.")
-    description: bool = Field(description="Get description from the model.")
-    name: bool = Field(description="Get name from the model.")
+    species: bool = Field(description="Get species from the model.", default=False)
+    parameters: bool = Field(description="Get parameters from the model.", default=False)
+    compartments: bool = Field(description="Get compartments from the model.", default=False)
+    units: bool = Field(description="Get units from the model.", default=False)
+    description: bool = Field(description="Get description from the model.", default=False)
+    name: bool = Field(description="Get name from the model.", default=False)
 
 class GetModelInfoInput(BaseModel):
     """