aiagents4pharma 1.11.0__py3-none-any.whl → 1.13.0__py3-none-any.whl

aiagents4pharma/configs/talk2biomodels/agents/t2b_agent/default.yaml

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

aiagents4pharma/talk2biomodels/agents/t2b_agent.py

@@ -17,6 +17,7 @@ from ..tools.simulate_model import SimulateModelTool
 from ..tools.custom_plotter import CustomPlotterTool
 from ..tools.ask_question import AskQuestionTool
 from ..tools.parameter_scan import ParameterScanTool
+from ..tools.steady_state import SteadyStateTool
 from ..states.state_talk2biomodels import Talk2Biomodels
 
 # Initialize logger
@@ -42,6 +43,7 @@ def get_app(uniq_id, llm_model='gpt-4o-mini'):
                     CustomPlotterTool(),
                     SearchModelsTool(),
                     GetModelInfoTool(),
+                    SteadyStateTool(),
                     ParameterScanTool()
                     ])
 

aiagents4pharma/talk2biomodels/states/state_talk2biomodels.py

@@ -22,3 +22,4 @@ class Talk2Biomodels(AgentState):
     sbml_file_path: Annotated[list, operator.add]
     dic_simulated_data: Annotated[list[dict], operator.add]
     dic_scanned_data: Annotated[list[dict], operator.add]
+    dic_steady_state_data: Annotated[list[dict], operator.add]

aiagents4pharma/talk2biomodels/tests/test_langgraph.py

@@ -98,7 +98,7 @@ def test_simulate_model_tool():
     # 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"
+    prompt = "Simulate model 64 and the uploaded model"
     # Invoke the agent
     app.invoke(
         {"messages": [HumanMessage(content=prompt)]},
@@ -145,10 +145,9 @@ def test_param_scan_tool():
     app.update_state(config, {"llm_model": "gpt-4o-mini"})
     prompt = """How will the value of Ab in model 537 change
             if the param kIL6Rbind is varied from 1 to 100 in steps of 10?
-            Set the initial `DoseQ2W` concentration to 300.
-            Reset the IL6{serum} concentration to 100 every 500 hours.
-            Assume that the model is simulated for 2016 hours with
-            an interval of 2016."""
+            Set the initial `DoseQ2W` concentration to 300. Also, reset
+            the IL6{serum} concentration to 100 every 500 hours and assume
+            that the model is simulated for 2016 hours with an interval of 2016."""
     # Invoke the agent
     app.invoke(
         {"messages": [HumanMessage(content=prompt)]},
@@ -181,11 +180,95 @@ def test_param_scan_tool():
     assert any((df["status"] == "success") &
                (df["name"] == "get_modelinfo"))
 
+def test_steady_state_tool():
+    '''
+    Test the steady_state tool.
+    '''
+    unique_id = 123
+    app = get_app(unique_id)
+    config = {"configurable": {"thread_id": unique_id}}
+    app.update_state(config, {"llm_model": "gpt-4o-mini"})
+    #########################################################
+    # In this case, we will test if the tool returns an error
+    # when the model does not achieve a steady state. The tool
+    # status should be "error".
+    prompt = """Run a steady state analysis of model 537."""
+    # Invoke the agent
+    app.invoke(
+        {"messages": [HumanMessage(content=prompt)]},
+        config=config
+    )
+    current_state = app.get_state(config)
+    reversed_messages = current_state.values["messages"][::-1]
+    tool_msg_status = None
+    for msg in reversed_messages:
+        # Assert that the status of the
+        # ToolMessage is "error"
+        if isinstance(msg, ToolMessage):
+            # print (msg)
+            tool_msg_status = msg.status
+            break
+    assert tool_msg_status == "error"
+    #########################################################
+    # In this case, we will test if the tool is indeed invoked
+    # successfully
+    prompt = """Run a steady state analysis of model 64.
+    Set the initial concentration of `Pyruvate` to 0.2. The
+    concentration of `NAD` resets to 100 every 2 time units."""
+    # Invoke the agent
+    app.invoke(
+        {"messages": [HumanMessage(content=prompt)]},
+        config=config
+    )
+    # Loop through the reversed messages until a
+    # ToolMessage is found.
+    current_state = app.get_state(config)
+    reversed_messages = current_state.values["messages"][::-1]
+    steady_state_invoked = False
+    for msg in reversed_messages:
+        # Assert that the message is a ToolMessage
+        # and its status is "error"
+        if isinstance(msg, ToolMessage):
+            print (msg)
+            if msg.name == "steady_state" and msg.status != "error":
+                steady_state_invoked = True
+                break
+    assert steady_state_invoked
+    #########################################################
+    # In this case, we will test if the `ask_question` tool is
+    # invoked upon asking a question about the already generated
+    # steady state results
+    prompt = """What is the Phosphoenolpyruvate concentration
+        at the steady state? Show onlyconcentration, rounded to
+        2 decimal places. For example, if the concentration is
+        0.123456, your response should be `0.12`. Do not return
+        any other information."""
+    # Invoke the agent
+    response = app.invoke(
+        {"messages": [HumanMessage(content=prompt)]},
+        config=config
+    )
+    assistant_msg = response["messages"][-1].content
+    current_state = app.get_state(config)
+    reversed_messages = current_state.values["messages"][::-1]
+    # Loop through the reversed messages until a
+    # ToolMessage is found.
+    ask_questool_invoked = False
+    for msg in reversed_messages:
+        # Assert that the message is a ToolMessage
+        # and its status is "error"
+        if isinstance(msg, ToolMessage):
+            if msg.name == "ask_question":
+                ask_questool_invoked = True
+                break
+    assert ask_questool_invoked
+    assert "0.06" in assistant_msg
+
 def test_integration():
     '''
     Test the integration of the tools.
     '''
-    unique_id = 123
+    unique_id = 1234567
     app = get_app(unique_id)
     config = {"configurable": {"thread_id": unique_id}}
     app.update_state(config, {"llm_model": "gpt-4o-mini"})
@@ -211,10 +294,8 @@ def test_integration():
     ##########################################
     # Update state
     app.update_state(config, {"llm_model": "gpt-4o-mini"})
-    prompt = "What is the concentration of CRP in serum at 1000 hours? "
-    # prompt += "Show only the concentration, rounded to 1 decimal place."
-    # prompt += "For example, if the concentration is 0.123456, "
-    # prompt += "your response should be `0.1`. Do not return any other information."
+    prompt = """What is the concentration of CRP
+            in serum after 1000 time points?"""
     # Test the tool get_modelinfo
     response = app.invoke(
                         {"messages": [HumanMessage(content=prompt)]},
@@ -271,8 +352,9 @@ def test_integration():
     # simulation results are available but
     # the species is not available
     ##########################################
-    prompt = "Plot the species `TP53`."
-
+    prompt = """Make a custom plot showing the
+        concentration of the species `TP53` over
+        time. Do not show any other species."""
     # Update state
     app.update_state(config, {"llm_model": "gpt-4o-mini"}
                     )

aiagents4pharma/talk2biomodels/tools/ask_question.py

@@ -5,7 +5,7 @@ Tool for asking a question about the simulation results.
 """
 
 import logging
-from typing import Type, Annotated
+from typing import Type, Annotated, Literal
 import pandas as pd
 from pydantic import BaseModel, Field
 from langchain_core.tools.base import BaseTool
@@ -22,9 +22,12 @@ 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.""")
+    question: str = Field(description="question about the simulation and steady state results")
+    experiment_name: str = Field(description="""Name assigned to the simulation
+                                            or steady state analysis when the tool 
+                                            simulate_model or steady_state is invoked.""")
+    question_context: Literal["simulation", "steady_state"] = Field(
+        description="Context of the question")
     state: Annotated[dict, InjectedState]
 
 # Note: It's important that every field has type hints.
@@ -32,41 +35,51 @@ class AskQuestionInput(BaseModel):
 # can lead to unexpected behavior.
 class AskQuestionTool(BaseTool):
     """
-    Tool for calculating the product of two numbers.
+    Tool for asking a question about the simulation or steady state results.
     """
     name: str = "ask_question"
-    description: str = "A tool to ask question about the simulation results."
+    description: str = """A tool to ask question about the
+                        simulation or steady state results."""
     args_schema: Type[BaseModel] = AskQuestionInput
     return_direct: bool = False
 
     def _run(self,
              question: str,
-             simulation_name: str,
+             experiment_name: str,
+             question_context: Literal["simulation", "steady_state"],
              state: Annotated[dict, InjectedState]) -> str:
         """
         Run the tool.
 
         Args:
-            question (str): The question to ask about the simulation results.
+            question (str): The question to ask about the simulation or steady state results.
             state (dict): The state of the graph.
-            simulation_name (str): The name assigned to the simulation.
+            experiment_name (str): The name assigned to the simulation or steady state analysis.
 
         Returns:
             str: The answer to the question.
         """
         logger.log(logging.INFO,
-                   "Calling ask_question tool %s, %s", question, simulation_name)
-        dic_simulated_data = {}
-        for data in state["dic_simulated_data"]:
+                   "Calling ask_question tool %s, %s, %s",
+                   question,
+                   question_context,
+                   experiment_name)
+        # print (f'Calling ask_question tool {question}, {question_context}, {experiment_name}')
+        if question_context == "steady_state":
+            dic_context = state["dic_steady_state_data"]
+        else:
+            dic_context = state["dic_simulated_data"]
+        dic_data = {}
+        for data in dic_context:
             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)
+                if key not in dic_data:
+                    dic_data[key] = []
+                dic_data[key] += [data[key]]
+        # print (dic_data)
+        df_data = pd.DataFrame.from_dict(dic_data)
         df = pd.DataFrame(
-                df_simulated_data[df_simulated_data['name'] == simulation_name]['data'].iloc[0]
-                )
+            df_data[df_data['name'] == experiment_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/steady_state.py

@@ -0,0 +1,208 @@
+#!/usr/bin/env python3
+
+"""
+Tool for parameter scan.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Type, Union, List, 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
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+@dataclass
+class TimeData:
+    """
+    Dataclass for storing the time data.
+    """
+    duration: Union[int, float] = 100
+    interval: Union[int, float] = 10
+
+@dataclass
+class SpeciesData:
+    """
+    Dataclass for storing the species data.
+    """
+    species_name: List[str] = Field(description="species name", default=[])
+    species_concentration: List[Union[int, float]] = Field(
+        description="initial species concentration",
+        default=[])
+
+@dataclass
+class TimeSpeciesNameConcentration:
+    """
+    Dataclass for storing the time, species name, and concentration data.
+    """
+    time: Union[int, float] = Field(description="time point where the event occurs")
+    species_name: str = Field(description="species name")
+    species_concentration: Union[int, float] = Field(
+        description="species concentration at the time point")
+
+@dataclass
+class ReocurringData:
+    """
+    Dataclass for species that reoccur. In other words, the concentration
+    of the species resets to a certain value after a certain time interval.
+    """
+    data: List[TimeSpeciesNameConcentration] = Field(
+        description="time, name, and concentration data of species that reoccur",
+        default=[])
+
+@dataclass
+class ArgumentData:
+    """
+    Dataclass for storing the argument data.
+    """
+    time_data: TimeData = Field(description="time data", default=None)
+    species_data: SpeciesData = Field(
+        description="species name and initial concentration data")
+    reocurring_data: ReocurringData = Field(
+        description="""Concentration and time data of species that reoccur
+            For example, a species whose concentration resets to a certain value
+            after a certain time interval""")
+    steadystate_name: str = Field(
+        description="""An AI assigned `_` separated name of
+        the steady state experiment based on human query""")
+
+def add_rec_events(model_object, reocurring_data):
+    """
+    Add reocurring events to the model.
+    """
+    for row in reocurring_data.data:
+        tp, sn, sc = row.time, row.species_name, row.species_concentration
+        basico.add_event(f'{sn}_{tp}',
+                            f'Time > {tp}',
+                            [[sn, str(sc)]],
+                            model=model_object.copasi_model)
+
+def run_steady_state(model_object, dic_species_data):
+    """
+    Run the steady state analysis.
+
+    Args:
+        model_object: The model object.
+        dic_species_data: 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_data)
+    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)
+    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
+                                   as well as the steady state data""",
+                                   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 for steady state analysis.
+    """
+    name: str = "steady_state"
+    description: str = """A tool to simulate a model and perform
+                        steady state analysisto answer questions
+                        about the steady state of species."""
+    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 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_data = {}
+        if arg_data:
+            # Prepare the dictionary of species data
+            if arg_data.species_data is not None:
+                dic_species_data = dict(zip(arg_data.species_data.species_name,
+                                            arg_data.species_data.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_data)
+        # Prepare the dictionary of scanned data
+        # that will be passed to the state of the graph
+        dic_steady_state_data = {
+            'name': arg_data.steadystate_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
+                                {arg_data.steadystate_name}
+                                are ready''',
+                        tool_call_id=tool_call_id
+                        )
+                    ],
+                }
+            )

aiagents4pharma/talk2knowledgegraphs/__init__.py

@@ -1,4 +1,5 @@
 '''
-This file is used to import the datasets, utils, and tools.
+This file is used to import the datasets and utils.
 '''
 from . import datasets
+from . import utils

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_enrichments_enrichments.py

@@ -0,0 +1,39 @@
+"""
+Test cases for utils/enrichments/enrichments.py
+"""
+
+from ..utils.enrichments.enrichments import Enrichments
+
+class TestEnrichments(Enrichments):
+    """Test implementation of the Enrichments interface for testing purposes."""
+
+    def enrich_documents(self, texts: list[str]) -> list[list[float]]:
+        return [
+            f"Additional text description of {text} as the input." for text in texts
+        ]
+
+    def enrich_documents_with_rag(self, texts, docs):
+        # Currently we don't have a RAG model to test this method.
+        # Thus, we will just call the enrich_documents method instead.
+        return self.enrich_documents(texts)
+
+def test_enrich_documents():
+    """Test enriching documents using the Enrichments interface."""
+    enrichments = TestEnrichments()
+    texts = ["text1", "text2"]
+    result = enrichments.enrich_documents(texts)
+    assert result == [
+        "Additional text description of text1 as the input.",
+        "Additional text description of text2 as the input.",
+    ]
+
+def test_enrich_documents_with_rag():
+    """Test enriching documents with RAG using the Enrichments interface."""
+    enrichments = TestEnrichments()
+    texts = ["text1", "text2"]
+    docs = ["doc1", "doc2"]
+    result = enrichments.enrich_documents_with_rag(texts, docs)
+    assert result == [
+        "Additional text description of text1 as the input.",
+        "Additional text description of text2 as the input.",
+    ]

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_enrichments_ollama.py

@@ -0,0 +1,117 @@
+"""
+Test cases for utils/enrichments/ollama.py
+"""
+
+import pytest
+import ollama
+from ..utils.enrichments.ollama import EnrichmentWithOllama
+
+@pytest.fixture(name="ollama_config")
+def fixture_ollama_config():
+    """Return a dictionary with Ollama configuration."""
+    return {
+        "model_name": "smollm2:360m",
+        "prompt_enrichment": """
+            Given the input as a list of strings, please return the list of addditional information of
+            each input terms using your prior knowledge.
+
+            Example:
+            Input: ['acetaminophen', 'aspirin']
+            Ouput: ['acetaminophen is a medication used to treat pain and fever', 
+            'aspirin is a medication used to treat pain, fever, and inflammation']
+            
+            Do not include any pretext as the output, only the list of strings enriched.
+
+            Input: {input}
+        """,
+        "temperature": 0.0,
+        "streaming": False,
+    }
+
+def test_no_model_ollama(ollama_config):
+    """Test the case when the Ollama model is not available."""
+    cfg = ollama_config
+    cfg_model = "smollm2:135m" # Choose a small model
+
+    # Delete the Ollama model
+    try:
+        ollama.delete(cfg_model)
+    except ollama.ResponseError:
+        pass
+
+    # Check if the model is available
+    with pytest.raises(
+        ValueError, match=f"Error: Pulled {cfg_model} model and restarted Ollama server."
+    ):
+        EnrichmentWithOllama(
+            model_name=cfg_model,
+            prompt_enrichment=cfg["prompt_enrichment"],
+            temperature=cfg["temperature"],
+            streaming=cfg["streaming"],
+        )
+    ollama.delete(cfg_model)
+
+def test_enrich_nodes_ollama(ollama_config):
+    """Test the Ollama textual enrichment class for node enrichment."""
+    # Prepare enrichment model
+    cfg = ollama_config
+    enr_model = EnrichmentWithOllama(
+        model_name=cfg["model_name"],
+        prompt_enrichment=cfg["prompt_enrichment"],
+        temperature=cfg["temperature"],
+        streaming=cfg["streaming"],
+    )
+
+    # Perform enrichment for nodes
+    nodes = ["Adalimumab", "Infliximab"]
+    enriched_nodes = enr_model.enrich_documents(nodes)
+    # Check the enriched nodes
+    assert len(enriched_nodes) == 2
+    assert all(
+        enriched_nodes[i] != nodes[i] for i in range(len(nodes))
+    )
+
+
+def test_enrich_relations_ollama(ollama_config):
+    """Test the Ollama textual enrichment class for relation enrichment."""
+    # Prepare enrichment model
+    cfg = ollama_config
+    enr_model = EnrichmentWithOllama(
+        model_name=cfg["model_name"],
+        prompt_enrichment=cfg["prompt_enrichment"],
+        temperature=cfg["temperature"],
+        streaming=cfg["streaming"],
+    )
+    # Perform enrichment for relations
+    relations = [
+        "IL23R-gene causation disease-inflammatory bowel diseases",
+        "NOD2-gene causation disease-inflammatory bowel diseases",
+    ]
+    enriched_relations = enr_model.enrich_documents(relations)
+    # Check the enriched relations
+    assert len(enriched_relations) == 2
+    assert all(
+        enriched_relations[i] != relations[i]
+        for i in range(len(relations))
+    )
+
+
+def test_enrich_ollama_rag(ollama_config):
+    """Test the Ollama textual enrichment class for enrichment with RAG (not implemented)."""
+    # Prepare enrichment model
+    cfg = ollama_config
+    enr_model = EnrichmentWithOllama(
+        model_name=cfg["model_name"],
+        prompt_enrichment=cfg["prompt_enrichment"],
+        temperature=cfg["temperature"],
+        streaming=cfg["streaming"],
+    )
+    # Perform enrichment for nodes
+    nodes = ["Adalimumab", "Infliximab"]
+    docs = [r"\path\to\doc1", r"\path\to\doc2"]
+    enriched_nodes = enr_model.enrich_documents_with_rag(nodes, docs)
+    # Check the enriched nodes
+    assert len(enriched_nodes) == 2
+    assert all(
+        enriched_nodes[i] != nodes[i] for i in range(len(nodes))
+    )

aiagents4pharma/talk2knowledgegraphs/utils/__init__.py

@@ -0,0 +1,5 @@
+'''
+This file is used to import utlities.
+'''
+from . import enrichments
+from . import embeddings

aiagents4pharma/talk2knowledgegraphs/utils/enrichments/__init__.py

@@ -0,0 +1,5 @@
+"""
+This package contains modules to use the enrichment model
+"""
+from . import enrichments
+from . import ollama

aiagents4pharma/talk2knowledgegraphs/utils/enrichments/enrichments.py

@@ -0,0 +1,36 @@
+"""
+Enrichments interface
+"""
+
+from abc import ABC, abstractmethod
+
+class Enrichments(ABC):
+    """Interface for enrichment models.
+
+    This is an interface meant for implementing text enrichment models.
+
+    Enrichment models are used to enrich node or relation features in a given knowledge graph.
+    """
+
+    @abstractmethod
+    def enrich_documents(self, texts: list[str]) -> list[list[str]]:
+        """Enrich documents.
+
+        Args:
+            texts: List of documents to enrich.
+
+        Returns:
+            List of enriched documents.
+        """
+
+    @abstractmethod
+    def enrich_documents_with_rag(self, texts: list[str], docs: list[str]) -> list[str]:
+        """Enrich documents with RAG.
+
+        Args:
+            texts: List of documents to enrich.
+            docs: List of reference documents to enrich the input texts.
+
+        Returns:
+            List of enriched documents with RAG.
+        """

aiagents4pharma/talk2knowledgegraphs/utils/enrichments/ollama.py

@@ -0,0 +1,123 @@
+#!/usr/bin/env python3
+
+"""
+Enrichment class using Ollama model based on LangChain Enrichment class.
+"""
+
+import time
+from typing import List
+import subprocess
+import ast
+import ollama
+from langchain_ollama import ChatOllama
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from .enrichments import Enrichments
+
+class EnrichmentWithOllama(Enrichments):
+    """
+    Enrichment class using Ollama model based on the Enrichment abstract class.
+    """
+    def __init__(
+        self,
+        model_name: str,
+        prompt_enrichment: str,
+        temperature: float,
+        streaming: bool,
+    ):
+        """
+        Initialize the EnrichmentWithOllama class.
+
+        Args:
+            model_name: The name of the Ollama model to be used.
+            prompt_enrichment: The prompt enrichment template.
+            temperature: The temperature for the Ollama model.
+            streaming: The streaming flag for the Ollama model.
+        """
+        # Setup the Ollama server
+        self.__setup(model_name)
+
+        # Set parameters
+        self.model_name = model_name
+        self.prompt_enrichment = prompt_enrichment
+        self.temperature = temperature
+        self.streaming = streaming
+
+        # Prepare prompt template
+        self.prompt_template = ChatPromptTemplate.from_messages(
+            [
+                ("system", self.prompt_enrichment),
+                ("human", "{input}"),
+            ]
+        )
+
+        # Prepare model
+        self.model = ChatOllama(
+            model=self.model_name,
+            temperature=self.temperature,
+            streaming=self.streaming,
+        )
+
+    def __setup(self, model_name: str) -> None:
+        """
+        Check if the Ollama model is available and run the Ollama server if needed.
+
+        Args:
+            model_name: The name of the Ollama model to be used.
+        """
+        try:
+            models_list = ollama.list()["models"]
+            if model_name not in [m['model'].replace(":latest", "") for m in models_list]:
+                ollama.pull(model_name)
+                time.sleep(30)
+                raise ValueError(f"Pulled {model_name} model")
+        except Exception as e:
+            with subprocess.Popen(
+                "ollama serve", shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+            ):
+                time.sleep(10)
+            raise ValueError(f"Error: {e} and restarted Ollama server.") from e
+
+    def enrich_documents(self, texts: List[str]) -> List[str]:
+        """
+        Enrich a list of input texts with additional textual features using OLLAMA model.
+        Important: Make sure the input is a list of texts based on the defined prompt template
+        with 'input' as the variable name.
+
+        Args:
+            texts: The list of texts to be enriched.
+
+        Returns:
+            The list of enriched texts.
+        """
+
+        # Perform enrichment
+        chain = self.prompt_template | self.model | StrOutputParser()
+
+        # Generate the enriched node
+        # Important: Make sure the input is a list of texts based on the defined prompt template
+        # with 'input' as the variable name
+        enriched_texts = chain.invoke({"input": "[" + ", ".join(texts) + "]"})
+
+        # Convert the enriched nodes to a list of dictionary
+        enriched_texts = ast.literal_eval(enriched_texts.replace("```", ""))
+
+        # Final check for the enriched texts
+        assert len(enriched_texts) == len(texts)
+
+        return enriched_texts
+
+    def enrich_documents_with_rag(self, texts, docs):
+        """
+        Enrich a list of input texts with additional textual features using OLLAMA model with RAG.
+        As of now, we don't have a RAG model to test this method yet.
+        Thus, we will just call the enrich_documents method instead.
+
+        Args:
+            texts: The list of texts to be enriched.
+            docs: The list of reference documents to enrich the input texts.
+        
+        Returns:
+            The list of enriched texts
+        """
+        return self.enrich_documents(texts)

{aiagents4pharma-1.11.0.dist-info → aiagents4pharma-1.13.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: aiagents4pharma
-Version: 1.11.0
+Version: 1.13.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
@@ -20,9 +20,11 @@ Requires-Dist: langchain-community==0.3.5
 Requires-Dist: langchain-core==0.3.31
 Requires-Dist: langchain-experimental==0.3.3
 Requires-Dist: langchain-openai==0.2.5
+Requires-Dist: langchain_ollama==0.2.2
 Requires-Dist: langgraph==0.2.66
 Requires-Dist: matplotlib==3.9.2
 Requires-Dist: openai==1.59.4
+Requires-Dist: ollama==0.4.6
 Requires-Dist: pandas==2.2.3
 Requires-Dist: plotly==5.24.1
 Requires-Dist: pydantic==2.9.2

{aiagents4pharma-1.11.0.dist-info → aiagents4pharma-1.13.0.dist-info}/RECORD

@@ -4,27 +4,28 @@ aiagents4pharma/configs/config.yaml,sha256=8y8uG6Dzx4-9jyb6hZ8r4lOJz5gA_sQhCiSCg
 aiagents4pharma/configs/talk2biomodels/__init__.py,sha256=5ah__-8XyRblwT0U1ByRigNjt_GyCheu7zce4aM-eZE,68
 aiagents4pharma/configs/talk2biomodels/agents/__init__.py,sha256=_ZoG8snICK2bidWtc2KOGs738LWg9_r66V9mOMnEb-E,71
 aiagents4pharma/configs/talk2biomodels/agents/t2b_agent/__init__.py,sha256=-fAORvyFmG2iSvFOFDixmt9OTQRR58y89uhhu2EgbA8,46
-aiagents4pharma/configs/talk2biomodels/agents/t2b_agent/default.yaml,sha256=yD7qZCneaM-JE5PdZjDmDoTRUdsFrzeCKZsBx1b-f20,293
+aiagents4pharma/configs/talk2biomodels/agents/t2b_agent/default.yaml,sha256=Oi89_BbxfQc6SGW1pC-hyZMqOIkiAMOlNwpCa4VCXk0,327
 aiagents4pharma/talk2biomodels/__init__.py,sha256=qUw3qXrENqSCLIKSLy_qtNPwPDTb1wdZ8fZispcHb3g,141
 aiagents4pharma/talk2biomodels/agents/__init__.py,sha256=sn5-fREjMdEvb-OUan3iOqrgYGjplNx3J8hYOaW0Po8,128
-aiagents4pharma/talk2biomodels/agents/t2b_agent.py,sha256=6Im4YFcdykN7wpEM8y9qi_x4lTg02WJpb0SEWh8TPLo,3188
+aiagents4pharma/talk2biomodels/agents/t2b_agent.py,sha256=8_2D3uknPjYbzKj7IDhC8xnz_HEvXck0b4RCJxoAxRs,3276
 aiagents4pharma/talk2biomodels/models/__init__.py,sha256=5fTHHm3PVloYPNKXbgNlcPgv3-u28ZquxGydFYDfhJA,122
 aiagents4pharma/talk2biomodels/models/basico_model.py,sha256=PH25FTOuUjsmw_UUxoRb-4kptOYpicEn4GqS0phS3nk,4807
 aiagents4pharma/talk2biomodels/models/sys_bio_model.py,sha256=JeoiGQAvQABHnG0wKR2XBmmxqQdtgO6kxaLDUTUmr1s,2001
 aiagents4pharma/talk2biomodels/states/__init__.py,sha256=YLg1-N0D9qyRRLRqwqfLCLAqZYDtMVZTfI8Y0b_4tbA,139
-aiagents4pharma/talk2biomodels/states/state_talk2biomodels.py,sha256=Dlsnh9dW1mCXTBXmlDAlOox7f4azFbLBG_2k3YPielM,824
+aiagents4pharma/talk2biomodels/states/state_talk2biomodels.py,sha256=XcrCsBbQgc4LCZ7ihpY5tcLQ82Mq17XhRXD1DW8Ep7U,887
 aiagents4pharma/talk2biomodels/tests/__init__.py,sha256=Jbw5tJxSrjGoaK5IX3pJWDCNzhrVQ10lkYq2oQ_KQD8,45
 aiagents4pharma/talk2biomodels/tests/test_basico_model.py,sha256=y82fpTJMPHwtXxlle1cGQ_2Bewwpxi0aJSVrVAYLhN0,2060
-aiagents4pharma/talk2biomodels/tests/test_langgraph.py,sha256=_71UZS1zucn6Nus4oCH7tINQVRvJEFnL0UIZ6-sUd3I,11967
+aiagents4pharma/talk2biomodels/tests/test_langgraph.py,sha256=QLAL4nmHrioTD-w-9OE0wQi5JdWJJ59PejNbDzCSvw4,15170
 aiagents4pharma/talk2biomodels/tests/test_sys_bio_model.py,sha256=HSmBBViMi0jYf4gWX21IbppAfDzG0nr_S3KtKS9fZVQ,2165
 aiagents4pharma/talk2biomodels/tools/__init__.py,sha256=SMTMlGHxTuJI7gjwGaTMv0XmilJ73-r5dp568hD3Fw0,266
-aiagents4pharma/talk2biomodels/tools/ask_question.py,sha256=uxCQ4ON8--D0ACPvT14t6x_aqm9LP6woBA4GM7bPXc4,3061
+aiagents4pharma/talk2biomodels/tools/ask_question.py,sha256=qpltsgyLFFwLYQeapQHASFRDCNiWsJkmTH_sUrfJ_Fg,3708
 aiagents4pharma/talk2biomodels/tools/custom_plotter.py,sha256=HWwKTX3o4dk0GcRVTO2hPrFSu98mtJ4TKC_hbHXOe1c,4018
 aiagents4pharma/talk2biomodels/tools/get_modelinfo.py,sha256=qA-4FOI-O728Nmn7s8JJ8HKwxvA9MZbst7NkPKTAMV4,5391
 aiagents4pharma/talk2biomodels/tools/load_biomodel.py,sha256=pyVzLQoMnuJYEwsjeOlqcUrbU1F1Z-pNlgkhFaoKpy0,689
 aiagents4pharma/talk2biomodels/tools/parameter_scan.py,sha256=aIyL_m46s3Q74ieJOZjZBM34VCjBKSMpEtckhdZofbE,12139
 aiagents4pharma/talk2biomodels/tools/search_models.py,sha256=Iq2ddofOOfZYtAurCISq3bAq5rbwB3l_rL1lgEFyFCI,2653
 aiagents4pharma/talk2biomodels/tools/simulate_model.py,sha256=sWmFVnVvJbdXXTqn_7gQl5UW0tv4FyU5yLXWLweLs_M,7059
+aiagents4pharma/talk2biomodels/tools/steady_state.py,sha256=itUXFTYO525D695LQlGb1ObwSuvHk0A5mXtgdCproqo,8102
 aiagents4pharma/talk2cells/__init__.py,sha256=zmOP5RAhabgKIQP-W4P4qKME2tG3fhAXM3MeO5_H8kE,120
 aiagents4pharma/talk2cells/agents/__init__.py,sha256=38nK2a_lEFRjO3qD6Fo9a3983ZCYat6hmJKWY61y2Mo,128
 aiagents4pharma/talk2cells/agents/scp_agent.py,sha256=gDMfhUNWHa_XWOqm1Ql6yLAdI_7bnIk5sRYn43H2sYk,3090
@@ -51,7 +52,7 @@ aiagents4pharma/talk2competitors/tools/s2/display_results.py,sha256=B8JJGohi1Eyx
 aiagents4pharma/talk2competitors/tools/s2/multi_paper_rec.py,sha256=FYLt47DAk6WOKfEk1Gj9zVvJGNyxA283PCp8IKW9U5M,4262
 aiagents4pharma/talk2competitors/tools/s2/search.py,sha256=pppjrQv5-8ep4fnqgTSBNgnbSnQsVIcNrRrH0p2TP1o,4025
 aiagents4pharma/talk2competitors/tools/s2/single_paper_rec.py,sha256=dAfUQxI7T5eu0eDxK8VAl7-JH0Wnw24CVkOQqwj-hXc,4810
-aiagents4pharma/talk2knowledgegraphs/__init__.py,sha256=SW7Ys2A4eXyFtizNPdSw91SHOPVUBGBsrCQ7TqwSUL0,91
+aiagents4pharma/talk2knowledgegraphs/__init__.py,sha256=4smVQoSMM6rflVnNkABqlDAAlSn4bYsq7rMVWjRGvis,103
 aiagents4pharma/talk2knowledgegraphs/datasets/__init__.py,sha256=L3gPuHskSegmtXskVrLIYr7FXe_ibKgJ2GGr1_Wok6k,173
 aiagents4pharma/talk2knowledgegraphs/datasets/biobridge_primekg.py,sha256=QlzDXmXREoa9MA6-GwzqRjdzndQeGBAF11Td6NFk_9Y,23426
 aiagents4pharma/talk2knowledgegraphs/datasets/dataset.py,sha256=-LaPLse8BkALqwFetNK7wch2dt9Dz6QKGKZKBKM6bIk,409
@@ -65,14 +66,19 @@ aiagents4pharma/talk2knowledgegraphs/tests/test_datasets_starkqa_primekg.py,sha2
 aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_embeddings.py,sha256=uYFoE_6zeU10_1mLLAHUr5c4S2XZMSc0Q_860o-KWEw,1517
 aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_huggingface.py,sha256=EINWyXg_3AMHF3WzFLhIUiFDuaEhTVHBvVAJr8VtMDg,1624
 aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_sentencetransformer.py,sha256=Qxo6WeIDRy8aLh1tNKw0kSlzmUj3MtTak63oW2YwB24,1327
-aiagents4pharma/talk2knowledgegraphs/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+aiagents4pharma/talk2knowledgegraphs/tests/test_utils_enrichments_enrichments.py,sha256=N6HRr4lWHXY7bTHe2uXJe4D_EG9WqZPibZne6qLl9_k,1447
+aiagents4pharma/talk2knowledgegraphs/tests/test_utils_enrichments_ollama.py,sha256=kMuB_vci6hKr2qJgXBmcje7yxeJ2nY2ImXw-NSJpts0,3912
+aiagents4pharma/talk2knowledgegraphs/utils/__init__.py,sha256=X8kSpmDTMkeE0fA5D0CWMsQ52YKoM5rRSXrjnali3IM,97
 aiagents4pharma/talk2knowledgegraphs/utils/kg_utils.py,sha256=6vQnPkeOWae_8jePjhma3sJuMTngy0I0tqzdFt6OqKg,2507
 aiagents4pharma/talk2knowledgegraphs/utils/embeddings/__init__.py,sha256=xRb0x7SoAb0nSVZYgjrqxWvENOMDuqIdL43NMjoOaCs,153
 aiagents4pharma/talk2knowledgegraphs/utils/embeddings/embeddings.py,sha256=1nGznrAj-xT0xuSMBGz2dOujJ7M_IwSR84njxtxsy9A,2523
 aiagents4pharma/talk2knowledgegraphs/utils/embeddings/huggingface.py,sha256=2vi_elf6EgzfagFAO5QnL3a_aXZyN7B1EBziu44MTfM,3806
 aiagents4pharma/talk2knowledgegraphs/utils/embeddings/sentence_transformer.py,sha256=36iKlisOpMtGR5xfTAlSHXWvPqVC_Jbezod8kbBBMVg,2136
-aiagents4pharma-1.11.0.dist-info/LICENSE,sha256=IcIbyB1Hyk5ZDah03VNQvJkbNk2hkBCDqQ8qtnCvB4Q,1077
-aiagents4pharma-1.11.0.dist-info/METADATA,sha256=GsvCuC24bJ_wwQP8aevOB4Eai8WcUIcr9kmN3EwYP_Y,8541
-aiagents4pharma-1.11.0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-aiagents4pharma-1.11.0.dist-info/top_level.txt,sha256=-AH8rMmrSnJtq7HaAObS78UU-cTCwvX660dSxeM7a0A,16
-aiagents4pharma-1.11.0.dist-info/RECORD,,
+aiagents4pharma/talk2knowledgegraphs/utils/enrichments/__init__.py,sha256=tW426knki2DBIHcWyF_K04iMMdbpIn_e_TpPmTgz2dI,113
+aiagents4pharma/talk2knowledgegraphs/utils/enrichments/enrichments.py,sha256=Bx8x6zzk5614ApWB90N_iv4_Y_Uq0-KwUeBwYSdQMU4,924
+aiagents4pharma/talk2knowledgegraphs/utils/enrichments/ollama.py,sha256=8eoxR-VHo0G7ReQIwje7xEhE-SJlHdef7_wJRpnvFIc,4116
+aiagents4pharma-1.13.0.dist-info/LICENSE,sha256=IcIbyB1Hyk5ZDah03VNQvJkbNk2hkBCDqQ8qtnCvB4Q,1077
+aiagents4pharma-1.13.0.dist-info/METADATA,sha256=gr1_xufeepNgQ_LbpMynm6gGpUyGNCcRd5RecsVAfiM,8609
+aiagents4pharma-1.13.0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+aiagents4pharma-1.13.0.dist-info/top_level.txt,sha256=-AH8rMmrSnJtq7HaAObS78UU-cTCwvX660dSxeM7a0A,16
+aiagents4pharma-1.13.0.dist-info/RECORD,,