aiagents4pharma 1.8.0__py3-none-any.whl → 1.15.0__py3-none-any.whl

aiagents4pharma/talk2biomodels/tools/steady_state.py

@@ -0,0 +1,167 @@
+#!/usr/bin/env python3
+
+"""
+Tool for parameter scan.
+"""
+
+import logging
+from typing import Type, Annotated
+import basico
+from pydantic import BaseModel, Field
+from langgraph.types import Command
+from langgraph.prebuilt import InjectedState
+from langchain_core.tools import BaseTool
+from langchain_core.messages import ToolMessage
+from langchain_core.tools.base import InjectedToolCallId
+from .load_biomodel import ModelData, load_biomodel
+from .load_arguments import ArgumentData, add_rec_events
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def run_steady_state(model_object,
+                     dic_species_to_be_analyzed_before_experiment):
+    """
+    Run the steady state analysis.
+
+    Args:
+        model_object: The model object.
+        dic_species_to_be_analyzed_before_experiment: Dictionary of species data.
+
+    Returns:
+        DataFrame: The results of the steady state analysis.
+    """
+    # Update the fixed model species and parameters
+    # These are the initial conditions of the model
+    # set by the user
+    model_object.update_parameters(dic_species_to_be_analyzed_before_experiment)
+    logger.log(logging.INFO, "Running steady state analysis")
+    # Run the steady state analysis
+    output = basico.task_steadystate.run_steadystate(model=model_object.copasi_model)
+    if output == 0:
+        logger.error("Steady state analysis failed")
+        raise ValueError("A steady state was not found")
+    logger.log(logging.INFO, "Steady state analysis successful")
+    # Store the steady state results in a DataFrame
+    df_steady_state = basico.model_info.get_species(model=model_object.copasi_model).reset_index()
+    # print (df_steady_state)
+    # Rename the column name to species_name
+    df_steady_state.rename(columns={'name': 'species_name'},
+                           inplace=True)
+    # Rename the column concentration to steady_state_concentration
+    df_steady_state.rename(columns={'concentration': 'steady_state_concentration'},
+                           inplace=True)
+    # Rename the column transition_time to steady_state_transition_time
+    df_steady_state.rename(columns={'transition_time': 'steady_state_transition_time'},
+                           inplace=True)
+    # Drop some columns
+    df_steady_state.drop(columns=
+                         [
+                            'initial_particle_number',
+                            'initial_expression',
+                            'expression',
+                            'particle_number',
+                            'type',
+                            'particle_number_rate',
+                            'key',
+                            'sbml_id',
+                            'display_name'],
+                            inplace=True)
+    logger.log(logging.INFO, "Steady state results with shape %s", df_steady_state.shape)
+    return df_steady_state
+
+class SteadyStateInput(BaseModel):
+    """
+    Input schema for the steady state tool.
+    """
+    sys_bio_model: ModelData = Field(description="model data",
+                                     default=None)
+    arg_data: ArgumentData = Field(
+        description="time, species, and reocurring data"
+                " that must be set before the steady state analysis"
+                " as well as the experiment name", default=None)
+    tool_call_id: Annotated[str, InjectedToolCallId]
+    state: Annotated[dict, InjectedState]
+
+# Note: It's important that every field has type hints. BaseTool is a
+# Pydantic class and not having type hints can lead to unexpected behavior.
+class SteadyStateTool(BaseTool):
+    """
+    Tool to bring a model to steady state.
+    """
+    name: str = "steady_state"
+    description: str = "A tool to bring a model to steady state."
+    args_schema: Type[BaseModel] = SteadyStateInput
+
+    def _run(self,
+        tool_call_id: Annotated[str, InjectedToolCallId],
+        state: Annotated[dict, InjectedState],
+        sys_bio_model: ModelData = None,
+        arg_data: ArgumentData = None
+    ) -> Command:
+        """
+        Run the tool.
+
+        Args:
+            tool_call_id (str): The tool call ID. This is injected by the system.
+            state (dict): The state of the tool.
+            sys_bio_model (ModelData): The model data.
+            arg_data (ArgumentData): The argument data.
+
+        Returns:
+            Command: The updated state of the tool.
+        """
+        logger.log(logging.INFO, "Calling the steady_state tool %s, %s",
+                   sys_bio_model, arg_data)
+        # print (f'Calling steady_state tool {sys_bio_model}, {arg_data}, {tool_call_id}')
+        sbml_file_path = state['sbml_file_path'][-1] if len(state['sbml_file_path']) > 0 else None
+        model_object = load_biomodel(sys_bio_model,
+                                  sbml_file_path=sbml_file_path)
+        # Prepare the dictionary of species data
+        # that will be passed to the simulate method
+        # of the BasicoModel class
+        dic_species_to_be_analyzed_before_experiment = {}
+        if arg_data:
+            # Prepare the dictionary of species data
+            if arg_data.species_to_be_analyzed_before_experiment is not None:
+                dic_species_to_be_analyzed_before_experiment = dict(
+                    zip(arg_data.species_to_be_analyzed_before_experiment.species_name,
+                        arg_data.species_to_be_analyzed_before_experiment.species_concentration))
+            # Add reocurring events (if any) to the model
+            if arg_data.reocurring_data is not None:
+                add_rec_events(model_object, arg_data.reocurring_data)
+        # Run the parameter scan
+        df_steady_state = run_steady_state(model_object,
+                                           dic_species_to_be_analyzed_before_experiment)
+        # Prepare the dictionary of scanned data
+        # that will be passed to the state of the graph
+        dic_steady_state_data = {
+            'name': arg_data.experiment_name,
+            'source': sys_bio_model.biomodel_id if sys_bio_model.biomodel_id else 'upload',
+            'tool_call_id': tool_call_id,
+            'data': df_steady_state.to_dict(orient='records')
+        }
+        # 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],
+            "dic_steady_state_data": [dic_steady_state_data]
+            }.items():
+            if value:
+                dic_updated_state_for_model[key] = value
+        # Return the updated state
+        return Command(
+                update=dic_updated_state_for_model|{
+                # Update the message history
+                "messages": [
+                ToolMessage(
+                        content=f"Steady state analysis of"
+                                f" {arg_data.experiment_name}"
+                                " was successful.",
+                        tool_call_id=tool_call_id
+                        )
+                    ],
+                }
+            )

aiagents4pharma/talk2cells/tests/scp_agent/test_scp_agent.py

@@ -0,0 +1,23 @@
+'''
+Test cases for the search_studies
+'''
+
+# from ..tools.search_studies import search_studies
+from aiagents4pharma.talk2cells.agents.scp_agent import get_app
+from langchain_core.messages import HumanMessage
+
+def test_agent_scp():
+    '''
+    Test the agent_scp.
+    '''
+    unique_id = 12345
+    app = get_app(unique_id)
+    config = {"configurable": {"thread_id": unique_id}}
+    prompt = "Search for studies on Crohns Disease."
+    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)

aiagents4pharma/talk2cells/tools/scp_agent/__init__.py

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

aiagents4pharma/talk2cells/tools/scp_agent/display_studies.py

@@ -0,0 +1,25 @@
+#!/usr/bin/env python3
+
+'''
+This tool is used to display the table of studies.
+'''
+
+import logging
+from typing import Annotated
+from langchain_core.tools import tool
+from langgraph.prebuilt import InjectedState
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+@tool('display_studies')
+def display_studies(state: Annotated[dict, InjectedState]):
+    """
+    Display the table of studies.
+
+    Args:
+        state (dict): The state of the agent.
+    """
+    logger.log(logging.INFO, "Calling the tool display_studies")
+    return state["search_table"]

aiagents4pharma/talk2cells/tools/scp_agent/search_studies.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+
+'''
+A tool to fetch studies from the Single Cell Portal.
+'''
+
+import logging
+from typing import Annotated
+import requests
+from langchain_core.tools import tool
+from langchain_core.tools.base import InjectedToolCallId
+from langchain_core.messages import ToolMessage
+from langgraph.types import Command
+import pandas as pd
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+@tool('search_studies')
+def search_studies(search_term: str,
+                   tool_call_id: Annotated[str, InjectedToolCallId],
+                   limit: int = 5):
+    """
+    Fetch studies from single cell portal
+
+    Args:
+        search_term (str): The search term to use. Example: "COVID-19", "cancer", etc.
+        limit (int): The number of papers to return. Default is 5.
+
+    """
+    logger.log(logging.INFO, "Calling the tool search_studies")
+    scp_endpoint = 'https://singlecell.broadinstitute.org/single_cell/api/v1/search?type=study'
+    # params = {'terms': search_term, 'facets': 'MONDO_0005011'}
+    params = {'terms': search_term}
+    status_code = 0
+    while status_code != 200:
+        # Make a GET request to the single cell portal
+        search_response = requests.get(scp_endpoint,
+                                       params=params,
+                                       timeout=10,
+                                       verify=False)
+        status_code = search_response.status_code
+        logger.log(logging.INFO, "Status code %s received from SCP")
+
+    # Select the columns to display in the table
+    selected_columns = ["study_source", "name", "study_url", "gene_count", "cell_count"]
+
+    # Extract the data from the response
+    # with the selected columns
+    df = pd.DataFrame(search_response.json()['studies'])[selected_columns]
+
+    # Convert column 'Study Name' into clickable
+    # hyperlinks from the column 'Study URL'
+    scp_api_url = 'https://singlecell.broadinstitute.org'
+    df['name'] = df.apply(
+            lambda x: f"<a href=\"{scp_api_url}/{x['study_url']}\">{x['name']}</a>",
+            axis=1)
+
+    # Excldue the column 'Study URL' from the dataframe
+    df = df.drop(columns=['study_url'])
+
+    # Add a new column a the beginning of the dataframe with row numbers
+    df.insert(0, 'S/N', range(1, 1 + len(df)))
+
+    # Update the state key 'search_table' with the dataframe in markdown format
+    return Command(
+        update={
+            # update the state keys
+            "search_table": df.to_markdown(tablefmt="grid"),
+            # update the message history
+            "messages": [
+                ToolMessage(
+                    f"Successfully fetched {limit} studies on {search_term}.",
+                    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

aiagents4pharma/talk2competitors/config/config.py

@@ -0,0 +1,110 @@
+"""Configuration module for AI agents handling paper searches and recommendations."""
+
+
+# pylint: disable=R0903
+class Config:
+    """Configuration class containing prompts for AI agents.
+
+    This class stores prompt templates used by various AI agents in the system,
+    particularly for academic paper searches and recommendations.
+    """
+
+    MAIN_AGENT_PROMPT = (
+        "You are a supervisory AI agent that routes user queries to specialized tools.\n"
+        "Your task is to select the most appropriate tool based on the user's request.\n\n"
+        "Available tools and their capabilities:\n\n"
+        "1. semantic_scholar_agent:\n"
+        "   - Search for academic papers and research\n"
+        "   - Get paper recommendations\n"
+        "   - Find similar papers\n"
+        "   USE FOR: Any queries about finding papers, academic research, "
+        "or getting paper recommendations\n\n"
+        "ROUTING GUIDELINES:\n\n"
+        "ALWAYS route to semantic_scholar_agent for:\n"
+        "- Finding academic papers\n"
+        "- Searching research topics\n"
+        "- Getting paper recommendations\n"
+        "- Finding similar papers\n"
+        "- Any query about academic literature\n\n"
+        "Approach:\n"
+        "1. Identify the core need in the user's query\n"
+        "2. Select the most appropriate tool based on the guidelines above\n"
+        "3. If unclear, ask for clarification\n"
+        "4. For multi-step tasks, focus on the immediate next step\n\n"
+        "Remember:\n"
+        "- Be decisive in your tool selection\n"
+        "- Focus on the immediate task\n"
+        "- Default to semantic_scholar_agent for any paper-finding tasks\n"
+        "- Ask for clarification if the request is ambiguous\n\n"
+        "When presenting paper search results, always use this exact format:\n\n"
+        "Remember to:\n"
+        "- Always remember to add the url\n"
+        "- Put URLs on the title line itself as markdown\n"
+        "- Maintain consistent spacing and formatting"
+    )
+
+    S2_AGENT_PROMPT = (
+        "You are a specialized academic research assistant with access to the following tools:\n\n"
+        "1. search_papers:\n"
+        "   USE FOR: General paper searches\n"
+        "   - Enhances search terms automatically\n"
+        "   - Adds relevant academic keywords\n"
+        "   - Focuses on recent research when appropriate\n\n"
+        "2. get_single_paper_recommendations:\n"
+        "   USE FOR: Finding papers similar to a specific paper\n"
+        "   - Takes a single paper ID\n"
+        "   - Returns related papers\n\n"
+        "3. get_multi_paper_recommendations:\n"
+        "   USE FOR: Finding papers similar to multiple papers\n"
+        "   - Takes multiple paper IDs\n"
+        "   - Finds papers related to all inputs\n\n"
+        "GUIDELINES:\n\n"
+        "For paper searches:\n"
+        "- Enhance search terms with academic language\n"
+        "- Include field-specific terminology\n"
+        '- Add "recent" or "latest" when appropriate\n'
+        "- Keep queries focused and relevant\n\n"
+        "For paper recommendations:\n"
+        "- Identify paper IDs (40-character hexadecimal strings)\n"
+        "- Use single_paper_recommendations for one ID\n"
+        "- Use multi_paper_recommendations for multiple IDs\n\n"
+        "Best practices:\n"
+        "1. Start with a broad search if no paper IDs are provided\n"
+        "2. Look for paper IDs in user input\n"
+        "3. Enhance search terms for better results\n"
+        "4. Consider the academic context\n"
+        "5. Be prepared to refine searches based on feedback\n\n"
+        "Remember:\n"
+        "- Always select the most appropriate tool\n"
+        "- Enhance search queries naturally\n"
+        "- Consider academic context\n"
+        "- Focus on delivering relevant results\n\n"
+        "IMPORTANT GUIDELINES FOR PAPER RECOMMENDATIONS:\n\n"
+        "For Multiple Papers:\n"
+        "- When getting recommendations for multiple papers, always use "
+        "get_multi_paper_recommendations tool\n"
+        "- DO NOT call get_single_paper_recommendations multiple times\n"
+        "- Always pass all paper IDs in a single call to get_multi_paper_recommendations\n"
+        '- Use for queries like "find papers related to both/all papers" or '
+        '"find similar papers to these papers"\n\n'
+        "For Single Paper:\n"
+        "- Use get_single_paper_recommendations when focusing on one specific paper\n"
+        "- Pass only one paper ID at a time\n"
+        '- Use for queries like "find papers similar to this paper" or '
+        '"get recommendations for paper X"\n'
+        "- Do not use for multiple papers\n\n"
+        "Examples:\n"
+        '- For "find related papers for both papers":\n'
+        "  ✓ Use get_multi_paper_recommendations with both paper IDs\n"
+        "  × Don't make multiple calls to get_single_paper_recommendations\n\n"
+        '- For "find papers related to the first paper":\n'
+        "  ✓ Use get_single_paper_recommendations with just that paper's ID\n"
+        "  × Don't use get_multi_paper_recommendations\n\n"
+        "Remember:\n"
+        "- Be precise in identifying which paper ID to use for single recommendations\n"
+        "- Don't reuse previous paper IDs unless specifically requested\n"
+        "- For fresh paper recommendations, always use the original paper ID"
+    )
+
+
+config = Config()

aiagents4pharma/talk2competitors/state/__init__.py

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

aiagents4pharma/talk2competitors/state/state_talk2competitors.py

@@ -0,0 +1,32 @@
+"""
+This is the state file for the talk2comp agent.
+"""
+
+import logging
+from typing import Annotated, Any, Dict, Optional
+
+from langgraph.prebuilt.chat_agent_executor import AgentState
+from typing_extensions import NotRequired, Required
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def replace_dict(existing: Dict[str, Any], new: Dict[str, Any]) -> Dict[str, Any]:
+    """Replace the existing dict with the new one."""
+    logger.info("Updating existing state %s with the state dict: %s", existing, new)
+    return new
+
+
+class Talk2Competitors(AgentState):
+    """
+    The state for the talk2comp agent, inheriting from AgentState.
+    """
+
+    papers: Annotated[Dict[str, Any], replace_dict]  # Changed from List to Dict
+    search_table: NotRequired[str]
+    next: str  # Required for routing in LangGraph
+    current_agent: NotRequired[Optional[str]]
+    is_last_step: Required[bool]  # Required field for LangGraph
+    llm_model: str