aiagents4pharma 1.8.3__py3-none-any.whl → 1.10.0__py3-none-any.whl

aiagents4pharma/__init__.py

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

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/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/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/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/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/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/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):
     """

aiagents4pharma/talk2biomodels/tools/simulate_model.py

@@ -52,10 +52,10 @@ class TimeSpeciesNameConcentration:
 class RecurringData:
     """
     Dataclass for storing the species and time data 
-    on recurring basis.
+    on reocurring basis.
     """
     data: List[TimeSpeciesNameConcentration] = Field(
-        description="species and time data on recurring basis",
+        description="species and time data on reocurring basis",
         default=None)
 
 @dataclass
@@ -68,12 +68,15 @@ class ArgumentData:
         description="species name and initial concentration data",
         default=None)
     recurring_data: RecurringData = Field(
-        description="species and time data on recurring basis",
+        description="species and time data on reocurring basis",
         default=None)
+    simulation_name: str = Field(
+        description="""An AI assigned `_` separated name of
+        the simulation based on human query""")
 
 def add_rec_events(model_object, recurring_data):
     """
-    Add recurring events to the model.
+    Add reocurring events to the model.
     """
     for row in recurring_data.data:
         tp, sn, sc = row.time, row.species_name, row.species_concentration
@@ -86,9 +89,12 @@ class SimulateModelInput(BaseModel):
     """
     Input schema for the SimulateModel tool.
     """
-    sys_bio_model: ModelData = Field(description="model data", default=None)
-    arg_data: ArgumentData = Field(description="time, species, and recurring data",
-                                        default=None)
+    sys_bio_model: ModelData = Field(description="model data",
+                                     default=None)
+    arg_data: ArgumentData = Field(description=
+                                   """time, species, and reocurring data
+                                   as well as the simulation name""",
+                                   default=None)
     tool_call_id: Annotated[str, InjectedToolCallId]
     state: Annotated[dict, InjectedState]
 
@@ -153,12 +159,20 @@ class SimulateModelTool(BaseTool):
             interval=interval
             )
 
+        dic_simulated_data = {
+            'name': arg_data.simulation_name,
+            'source': sys_bio_model.biomodel_id if sys_bio_model.biomodel_id else 'upload',
+            'tool_call_id': tool_call_id,
+            'data': df.to_dict()
+        }
+
         # Prepare the dictionary of updated state for the model
         dic_updated_state_for_model = {}
         for key, value in {
-                        "model_id": [sys_bio_model.biomodel_id],
-                        "sbml_file_path": [sbml_file_path],
-                        }.items():
+            "model_id": [sys_bio_model.biomodel_id],
+            "sbml_file_path": [sbml_file_path],
+            "dic_simulated_data": [dic_simulated_data],
+            }.items():
             if value:
                 dic_updated_state_for_model[key] = value
 
@@ -166,11 +180,11 @@ class SimulateModelTool(BaseTool):
         return Command(
                 update=dic_updated_state_for_model|{
                 # update the state keys
-                "dic_simulated_data": df.to_dict(),
+                # "dic_simulated_data": df.to_dict(),
                 # update the message history
                 "messages": [
                     ToolMessage(
-                        content="Simulation results are ready.",
+                        content=f"Simulation results of {arg_data.simulation_name}",
                         tool_call_id=tool_call_id
                         )
                     ],

aiagents4pharma/talk2competitors/__init__.py

@@ -0,0 +1,5 @@
+"""
+This file is used to import all the modules in the package.
+"""
+
+from . import agents, config, state, tests, tools

aiagents4pharma/talk2competitors/agents/__init__.py

@@ -0,0 +1,6 @@
+'''
+This file is used to import all the modules in the package.
+'''
+
+from . import main_agent
+from . import s2_agent

aiagents4pharma/talk2competitors/agents/main_agent.py

@@ -0,0 +1,130 @@
+#!/usr/bin/env python3
+
+"""
+Main agent for the talk2competitors app.
+"""
+
+import logging
+from typing import Literal
+from dotenv import load_dotenv
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import AIMessage
+from langchain_openai import ChatOpenAI
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph import END, START, StateGraph
+from langgraph.types import Command
+from ..agents import s2_agent
+from ..config.config import config
+from ..state.state_talk2competitors import Talk2Competitors
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+load_dotenv()
+
+def make_supervisor_node(llm: BaseChatModel) -> str:
+    """
+    Creates a supervisor node following LangGraph patterns.
+
+    Args:
+        llm (BaseChatModel): The language model to use for generating responses.
+
+    Returns:
+        str: The supervisor node function.
+    """
+    # options = ["FINISH", "s2_agent"]
+
+    def supervisor_node(state: Talk2Competitors) -> Command[Literal["s2_agent", "__end__"]]:
+        """
+        Supervisor node that routes to appropriate sub-agents.
+
+        Args:
+            state (Talk2Competitors): The current state of the conversation.
+
+        Returns:
+            Command[Literal["s2_agent", "__end__"]]: The command to execute next.
+        """
+        logger.info("Supervisor node called")
+
+        messages = [{"role": "system", "content": config.MAIN_AGENT_PROMPT}] + state[
+            "messages"
+        ]
+        response = llm.invoke(messages)
+        goto = (
+            "FINISH"
+            if not any(
+                kw in state["messages"][-1].content.lower()
+                for kw in ["search", "paper", "find"]
+            )
+            else "s2_agent"
+        )
+
+        if goto == "FINISH":
+            return Command(
+                goto=END,
+                update={
+                    "messages": state["messages"]
+                    + [AIMessage(content=response.content)],
+                    "is_last_step": True,
+                    "current_agent": None,
+                },
+            )
+
+        return Command(
+            goto="s2_agent",
+            update={
+                "messages": state["messages"],
+                "is_last_step": False,
+                "current_agent": "s2_agent",
+            },
+        )
+
+    return supervisor_node
+
+def get_app(thread_id: str, llm_model ='gpt-4o-mini') -> StateGraph:
+    """
+    Returns the langraph app with hierarchical structure.
+
+    Args:
+        thread_id (str): The thread ID for the conversation.
+
+    Returns:
+        The compiled langraph app.
+    """
+    def call_s2_agent(state: Talk2Competitors) -> Command[Literal["__end__"]]:
+        """
+        Node for calling the S2 agent.
+
+        Args:
+            state (Talk2Competitors): The current state of the conversation.
+
+        Returns:
+            Command[Literal["__end__"]]: The command to execute next.
+        """
+        logger.info("Calling S2 agent")
+        app = s2_agent.get_app(thread_id, llm_model)
+        response = app.invoke(state)
+        logger.info("S2 agent completed")
+        return Command(
+            goto=END,
+            update={
+                "messages": response["messages"],
+                "papers": response.get("papers", []),
+                "is_last_step": True,
+                "current_agent": "s2_agent",
+            },
+        )
+    llm = ChatOpenAI(model=llm_model, temperature=0)
+    workflow = StateGraph(Talk2Competitors)
+
+    supervisor = make_supervisor_node(llm)
+    workflow.add_node("supervisor", supervisor)
+    workflow.add_node("s2_agent", call_s2_agent)
+
+    # Define edges
+    workflow.add_edge(START, "supervisor")
+    workflow.add_edge("s2_agent", END)
+
+    app = workflow.compile(checkpointer=MemorySaver())
+    logger.info("Main agent workflow compiled")
+    return app

aiagents4pharma/talk2competitors/agents/s2_agent.py

@@ -0,0 +1,75 @@
+#/usr/bin/env python3
+
+'''
+Agent for interacting with Semantic Scholar
+'''
+
+import logging
+from dotenv import load_dotenv
+from langchain_openai import ChatOpenAI
+from langgraph.graph import START, StateGraph
+from langgraph.prebuilt import create_react_agent
+from langgraph.checkpoint.memory import MemorySaver
+from ..config.config import config
+from ..state.state_talk2competitors import Talk2Competitors
+# from ..tools.s2 import s2_tools
+from ..tools.s2.search import search_tool
+from ..tools.s2.display_results import display_results
+from ..tools.s2.single_paper_rec import get_single_paper_recommendations
+from ..tools.s2.multi_paper_rec import get_multi_paper_recommendations
+
+load_dotenv()
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def get_app(uniq_id, llm_model='gpt-4o-mini'):
+    '''
+    This function returns the langraph app.
+    '''
+    def agent_s2_node(state: Talk2Competitors):
+        '''
+        This function calls the model.
+        '''
+        logger.log(logging.INFO, "Creating Agent_S2 node with thread_id %s", uniq_id)
+        response = model.invoke(state, {"configurable": {"thread_id": uniq_id}})
+        return response
+
+    # Define the tools
+    tools = [search_tool,
+            display_results,
+            get_single_paper_recommendations,
+            get_multi_paper_recommendations]
+
+    # Create the LLM
+    llm = ChatOpenAI(model=llm_model, temperature=0)
+    model = create_react_agent(
+                            llm,
+                            tools=tools,
+                            state_schema=Talk2Competitors,
+                            state_modifier=config.S2_AGENT_PROMPT,
+                            checkpointer=MemorySaver()
+                        )
+
+    # Define a new graph
+    workflow = StateGraph(Talk2Competitors)
+
+    # Define the two nodes we will cycle between
+    workflow.add_node("agent_s2", agent_s2_node)
+
+    # Set the entrypoint as `agent`
+    # This means that this node is the first one called
+    workflow.add_edge(START, "agent_s2")
+
+    # Initialize memory to persist state between graph runs
+    checkpointer = MemorySaver()
+
+    # Finally, we compile it!
+    # This compiles it into a LangChain Runnable,
+    # 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)
+    logger.log(logging.INFO, "Compiled the graph")
+
+    return app

aiagents4pharma/talk2competitors/config/__init__.py

@@ -0,0 +1,5 @@
+"""
+This package contains configuration settings and prompts used by various AI agents
+"""
+
+from . import config