aiagents4pharma 1.17.1__py3-none-any.whl → 1.19.0__py3-none-any.whl

aiagents4pharma/talk2biomodels/agents/t2b_agent.py

@@ -8,6 +8,7 @@ import logging
 from typing import Annotated
 import hydra
 from langchain_openai import ChatOpenAI
+from langchain_core.language_models.chat_models import BaseChatModel
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.graph import START, StateGraph
 from langgraph.prebuilt import create_react_agent, ToolNode, InjectedState
@@ -26,7 +27,8 @@ from ..states.state_talk2biomodels import Talk2Biomodels
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-def get_app(uniq_id, llm_model='gpt-4o-mini'):
+def get_app(uniq_id,
+            llm_model: BaseChatModel = ChatOpenAI(model='gpt-4o-mini', temperature=0)):
     '''
     This function returns the langraph app.
     '''
@@ -51,8 +53,6 @@ def get_app(uniq_id, llm_model='gpt-4o-mini'):
                     QueryArticle()
                 ])
 
-    # Define the model
-    llm = ChatOpenAI(model=llm_model, temperature=0)
     # Load hydra configuration
     logger.log(logging.INFO, "Load Hydra configuration for Talk2BioModels agent.")
     with hydra.initialize(version_base=None, config_path="../configs"):
@@ -62,7 +62,7 @@ def get_app(uniq_id, llm_model='gpt-4o-mini'):
     logger.log(logging.INFO, "state_modifier: %s", cfg.state_modifier)
     # Create the agent
     model = create_react_agent(
-                llm,
+                llm_model,
                 tools=tools,
                 state_schema=Talk2Biomodels,
                 state_modifier=cfg.state_modifier,

aiagents4pharma/talk2biomodels/configs/tools/ask_question/default.yaml

@@ -10,22 +10,14 @@ steady_state_prompt: >
 
   Here are some instructions to help you answer questions:
 
-  1. Before you answer any question, follow the plan and solve
-  technique. Start by understanding the question, then plan your
-  approach to solve the question, and finally solve the question
-  by following the plan. Always give a brief explanation of your
-  answer to the user.
+  1. If the user wants to know the time taken by the model to reach
+  steady state, you should look at the `steady_state_transition_time`
+  column of the data for the model species.
+  
+  2. The highest value in the column `steady_state_transition_time`
+  is the time taken by the model to reach steady state.
 
-  2. If the user wants to know the time taken by the model to reach
-  steady state, you should look at the steady_state_transition_time
-  column of the data for the model species. The highest value in
-  this column is the time taken by the model to reach steady state.
-
-  3. To get accurate results, trim the data to the relevant columns
-  before performing any calculations. This will help you avoid
-  errors in your calculations, and ignore irrelevant data.
-
-  4. Please use the units provided below to answer the questions.
+  3. Please use the units provided below to answer the questions.
 simulation_prompt: >
   Following is the information about the data frame:
   1. First column is the time column, and the rest of the columns

aiagents4pharma/talk2biomodels/states/state_talk2biomodels.py

@@ -7,6 +7,8 @@ This is the state file for the Talk2BioModels agent.
 from typing import Annotated
 import operator
 from langgraph.prebuilt.chat_agent_executor import AgentState
+from langchain_core.language_models import BaseChatModel
+from langchain_core.embeddings import Embeddings
 
 def add_data(data1: dict, data2: dict) -> dict:
     """
@@ -26,7 +28,8 @@ class Talk2Biomodels(AgentState):
     """
     The state for the Talk2BioModels agent.
     """
-    llm_model: str
+    llm_model: BaseChatModel
+    text_embedding_model: Embeddings
     pdf_file_name: str
     # A StateGraph may receive a concurrent updates
     # which is not supported by the StateGraph. Hence,

aiagents4pharma/talk2biomodels/tests/test_ask_question.py

@@ -3,6 +3,7 @@ Test cases for Talk2Biomodels.
 '''
 
 from langchain_core.messages import HumanMessage, ToolMessage
+from langchain_openai import ChatOpenAI
 from ..agents.t2b_agent import get_app
 
 def test_ask_question_tool():
@@ -10,7 +11,7 @@ def test_ask_question_tool():
     Test the ask_question tool without the simulation results.
     '''
     unique_id = 12345
-    app = get_app(unique_id, llm_model='gpt-4o-mini')
+    app = get_app(unique_id)
     config = {"configurable": {"thread_id": unique_id}}
 
     ##########################################
@@ -20,7 +21,8 @@ def test_ask_question_tool():
     # case, the tool should return an error
     ##########################################
     # Update state
-    app.update_state(config, {"llm_model": "gpt-4o-mini"})
+    app.update_state(config,
+            {"llm_model": ChatOpenAI(model='gpt-4o-mini', temperature=0)})
     # Define the prompt
     prompt = "Call the ask_question tool to answer the "
     prompt += "question: What is the concentration of CRP "

aiagents4pharma/talk2biomodels/tests/test_get_annotation.py

@@ -5,6 +5,7 @@ Test cases for Talk2Biomodels get_annotation tool.
 import random
 import pytest
 from langchain_core.messages import HumanMessage, ToolMessage
+from langchain_openai import ChatOpenAI
 from ..agents.t2b_agent import get_app
 from ..tools.get_annotation import prepare_content_msg
 
@@ -16,7 +17,9 @@ def make_graph_fixture():
     unique_id = random.randint(1000, 9999)
     graph = get_app(unique_id)
     config = {"configurable": {"thread_id": unique_id}}
-    graph.update_state(config, {"llm_model": "gpt-4o-mini"})
+    graph.update_state(config, {"llm_model": ChatOpenAI(model='gpt-4o-mini',
+                                    temperature=0)
+                                })
     return graph, config
 
 def test_no_model_provided(make_graph):
@@ -85,7 +88,6 @@ def test_invalid_species_provided(make_graph):
                 # (likely due to an invalid species).
                 test_condition = True
                 break
-    # assert test_condition
     assert test_condition
 
 def test_invalid_and_valid_species_provided(make_graph):

aiagents4pharma/talk2biomodels/tests/test_integration.py

@@ -4,8 +4,11 @@ Test cases for Talk2Biomodels.
 
 import pandas as pd
 from langchain_core.messages import HumanMessage, ToolMessage
+from langchain_openai import ChatOpenAI
 from ..agents.t2b_agent import get_app
 
+LLM_MODEL = ChatOpenAI(model='gpt-4o-mini', temperature=0)
+
 def test_integration():
     '''
     Test the integration of the tools.
@@ -13,7 +16,7 @@ def test_integration():
     unique_id = 1234567
     app = get_app(unique_id)
     config = {"configurable": {"thread_id": unique_id}}
-    app.update_state(config, {"llm_model": "gpt-4o-mini"})
+    app.update_state(config, {"llm_model": LLM_MODEL})
     # ##########################################
     # ## Test simulate_model tool
     # ##########################################
@@ -34,7 +37,7 @@ def test_integration():
     # results are available
     ##########################################
     # Update state
-    app.update_state(config, {"llm_model": "gpt-4o-mini"})
+    app.update_state(config, {"llm_model": LLM_MODEL})
     prompt = """What is the concentration of CRP in serum after 100 hours?
     Round off the value to 2 decimal places."""
     # Test the tool get_modelinfo
@@ -49,12 +52,15 @@ def test_integration():
 
     ##########################################
     # Test custom_plotter tool when the
-    # simulation results are available
+    # simulation results are available but
+    # the species is not available
     ##########################################
-    prompt = "Plot only CRP related species."
-
+    prompt = """Call the custom_plotter tool to make a plot
+        showing only species `TP53` and `Pyruvate`. Let me
+        know if these species were not found. Do not
+        invoke any other tool."""
     # Update state
-    app.update_state(config, {"llm_model": "gpt-4o-mini"}
+    app.update_state(config, {"llm_model": LLM_MODEL}
                     )
     # Test the tool get_modelinfo
     response = app.invoke(
@@ -66,11 +72,8 @@ def test_integration():
     # Get the messages from the current state
     # and reverse the order
     reversed_messages = current_state.values["messages"][::-1]
-    # Loop through the reversed messages
-    # until a ToolMessage is found.
-    expected_header = ['Time', 'CRP{serum}', 'CRPExtracellular']
-    expected_header += ['CRP Suppression (%)', 'CRP (% of baseline)']
-    expected_header += ['CRP{liver}']
+    # Loop through the reversed messages until a
+    # ToolMessage is found.
     predicted_artifact = []
     for msg in reversed_messages:
         if isinstance(msg, ToolMessage):
@@ -80,24 +83,17 @@ def test_integration():
             if msg.name == "custom_plotter":
                 predicted_artifact = msg.artifact
                 break
-    # 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))
+    # Check if the the predicted artifact is `None`
+    assert predicted_artifact is None
+
     ##########################################
     # Test custom_plotter tool when the
-    # simulation results are available but
-    # the species is not available
+    # simulation results are available
     ##########################################
-    prompt = """Make a custom plot showing the
-        concentration of the species `TP53` over
-        time. Do not show any other species."""
+    prompt = "Plot only CRP related species."
+
     # Update state
-    app.update_state(config, {"llm_model": "gpt-4o-mini"}
+    app.update_state(config, {"llm_model": LLM_MODEL}
                     )
     # Test the tool get_modelinfo
     response = app.invoke(
@@ -105,13 +101,15 @@ def test_integration():
                         config=config
                     )
     assistant_msg = response["messages"][-1].content
-    # print (response["messages"])
     current_state = app.get_state(config)
     # Get the messages from the current state
     # and reverse the order
     reversed_messages = current_state.values["messages"][::-1]
-    # Loop through the reversed messages until a
-    # ToolMessage is found.
+    # Loop through the reversed messages
+    # until a ToolMessage is found.
+    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):
@@ -121,5 +119,11 @@ def test_integration():
             if msg.name == "custom_plotter":
                 predicted_artifact = msg.artifact
                 break
-    # Check if the the predicted artifact is `None`
-    assert predicted_artifact is None
+    # 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))

aiagents4pharma/talk2biomodels/tests/test_query_article.py

@@ -5,6 +5,7 @@ Test cases for Talk2Biomodels query_article tool.
 from pydantic import BaseModel, Field
 from langchain_core.messages import HumanMessage, ToolMessage
 from langchain_openai import ChatOpenAI
+from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings
 from ..agents.t2b_agent import get_app
 
 class Article(BaseModel):
@@ -21,8 +22,10 @@ def test_query_article_with_an_article():
     app = get_app(unique_id)
     config = {"configurable": {"thread_id": unique_id}}
     # Update state by providing the pdf file name
+    # and the text embedding model
     app.update_state(config,
-      {"pdf_file_name": "aiagents4pharma/talk2biomodels/tests/article_on_model_537.pdf"})
+      {"pdf_file_name": "aiagents4pharma/talk2biomodels/tests/article_on_model_537.pdf",
+       "text_embedding_model": NVIDIAEmbeddings(model='nvidia/llama-3.2-nv-embedqa-1b-v2')})
     prompt = "What is the title of the article?"
     # Test the tool query_article
     response = app.invoke(
@@ -55,6 +58,9 @@ def test_query_article_without_an_article():
     app = get_app(unique_id)
     config = {"configurable": {"thread_id": unique_id}}
     prompt = "What is the title of the uploaded article?"
+    # Update state by providing the text embedding model
+    app.update_state(config,
+      {"text_embedding_model": NVIDIAEmbeddings(model='nvidia/llama-3.2-nv-embedqa-1b-v2')})
     # Test the tool query_article
     app.invoke(
             {"messages": [HumanMessage(content=prompt)]},

aiagents4pharma/talk2biomodels/tests/test_search_models.py

@@ -3,6 +3,7 @@ Test cases for Talk2Biomodels search models tool.
 '''
 
 from langchain_core.messages import HumanMessage
+from langchain_nvidia_ai_endpoints import ChatNVIDIA
 from ..agents.t2b_agent import get_app
 
 def test_search_models_tool():
@@ -13,7 +14,8 @@ def test_search_models_tool():
     app = get_app(unique_id)
     config = {"configurable": {"thread_id": unique_id}}
     # Update state
-    app.update_state(config, {"llm_model": "gpt-4o-mini"})
+    app.update_state(config,
+            {"llm_model": ChatNVIDIA(model="meta/llama-3.3-70b-instruct")})
     prompt = "Search for models on Crohn's disease."
     # Test the tool get_modelinfo
     response = app.invoke(

aiagents4pharma/talk2biomodels/tests/test_steady_state.py

@@ -3,8 +3,11 @@ Test cases for Talk2Biomodels steady state tool.
 '''
 
 from langchain_core.messages import HumanMessage, ToolMessage
+from langchain_openai import ChatOpenAI
 from ..agents.t2b_agent import get_app
 
+LLM_MODEL = ChatOpenAI(model='gpt-4o-mini', temperature=0)
+
 def test_steady_state_tool():
     '''
     Test the steady_state tool.
@@ -12,7 +15,7 @@ def test_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"})
+    app.update_state(config, {"llm_model": LLM_MODEL})
     #########################################################
     # In this case, we will test if the tool returns an error
     # when the model does not achieve a steady state. The tool
@@ -37,8 +40,8 @@ def test_steady_state_tool():
     #########################################################
     # 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
+    prompt = """Bring model 64 to a steady state. 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(

aiagents4pharma/talk2biomodels/tools/ask_question.py

@@ -12,7 +12,6 @@ import pandas as pd
 from pydantic import BaseModel, Field
 from langchain_core.tools.base import BaseTool
 from langchain_experimental.agents import create_pandas_dataframe_agent
-from langchain_openai import ChatOpenAI
 from langgraph.prebuilt import InjectedState
 
 # Initialize logger
@@ -101,7 +100,7 @@ class AskQuestionTool(BaseTool):
         prompt_content += f"{basico.model_info.get_model_units()}\n\n"
         # Create a pandas dataframe agent
         df_agent = create_pandas_dataframe_agent(
-                        ChatOpenAI(model=state['llm_model']),
+                        state['llm_model'],
                         allow_dangerous_code=True,
                         agent_type='tool-calling',
                         df=df,

aiagents4pharma/talk2biomodels/tools/custom_plotter.py

@@ -5,10 +5,9 @@ Tool for plotting a custom figure.
 """
 
 import logging
-from typing import Type, List, TypedDict, Annotated, Tuple, Union, Literal
+from typing import Type, Annotated, List, Tuple, Union, Literal
 from pydantic import BaseModel, Field
 import pandas as pd
-from langchain_openai import ChatOpenAI
 from langchain_core.tools import BaseTool
 from langgraph.prebuilt import InjectedState
 
@@ -71,30 +70,44 @@ class CustomPlotterTool(BaseTool):
         species_names = df.columns.tolist()
         # Exclude the time column
         species_names.remove('Time')
+        logging.log(logging.INFO, "Species names: %s", species_names)
         # In the following code, we extract the species
         # from the user question. We use Literal to restrict
         # the species names to the ones available in the
         # simulation results.
-        class CustomHeader(TypedDict):
+        class CustomHeader(BaseModel):
             """
             A list of species based on user question.
+
+            This is a Pydantic model that restricts the species
+            names to the ones available in the simulation results.
+            
+            If no species is relevant, set the attribute
+            `relevant_species` to None.
             """
             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="This is a list of species based on the user question."
+                    "It is restricted to the species available in the simulation results."
+                    "If no species is relevant, set this attribute to None."
+                    "If the user asks for very specific species (for example, using the"
+                    "keyword `only` in the question), set this attribute to correspond "
+                    "to the species available in the simulation results, otherwise set it to None."
+                    )
         # Create an instance of the LLM model
-        llm = ChatOpenAI(model=state['llm_model'], temperature=0)
+        logging.log(logging.INFO, "LLM model: %s", state['llm_model'])
+        llm = state['llm_model']
         llm_with_structured_output = llm.with_structured_output(CustomHeader)
         results = llm_with_structured_output.invoke(question)
+        if results.relevant_species is None:
+            raise ValueError("No species found in the simulation results \
+                             that matches the user prompt.")
         extracted_species = []
         # Extract the species from the results
         # that are available in the simulation results
-        for species in results['relevant_species']:
+        for species in results.relevant_species:
             if species in species_names:
                 extracted_species.append(species)
-        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
+        logging.info("Extracted species: %s", 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_annotation.py

@@ -5,7 +5,7 @@ This module contains the `GetAnnotationTool` for fetching species annotations
 based on the provided model and species names.
 """
 import math
-from typing import List, Annotated, Type, TypedDict, Union, Literal
+from typing import List, Annotated, Type, Union, Literal
 import logging
 from dataclasses import dataclass
 import hydra
@@ -17,7 +17,7 @@ from langgraph.prebuilt import InjectedState
 from langchain_core.tools.base import BaseTool
 from langchain_core.tools.base import InjectedToolCallId
 from langchain_core.messages import ToolMessage
-from langchain_openai import ChatOpenAI
+# from langchain_openai import ChatOpenAI
 from .load_biomodel import ModelData, load_biomodel
 from ..api.uniprot import search_uniprot_labels
 from ..api.ols import search_ols_labels
@@ -49,7 +49,7 @@ def extract_relevant_species_names(model_object, arg_data, state):
     all_species_names = df_species.index.tolist()
 
     # Define a structured output for the LLM model
-    class CustomHeader(TypedDict):
+    class CustomHeader(BaseModel):
         """
         A list of species based on user question.
         """
@@ -58,17 +58,21 @@ def extract_relevant_species_names(model_object, arg_data, state):
                 If no relevant species are found, it must be None.""")
 
     # Create an instance of the LLM model
-    llm = ChatOpenAI(model=state['llm_model'], temperature=0)
+    llm = state['llm_model']
     # Get the structured output from the LLM model
     llm_with_structured_output = llm.with_structured_output(CustomHeader)
     # Define the question for the LLM model using the prompt
     question = cfg.prompt
     question += f'Here is the user question: {arg_data.user_question}'
     # Invoke the LLM model with the user question
-    dic = llm_with_structured_output.invoke(question)
+    results = llm_with_structured_output.invoke(question)
+    logging.info("Results from the LLM model: %s", results)
+    # Check if the returned species names are empty
+    if not results.relevant_species:
+        raise ValueError("Model does not contain the requested species.")
     extracted_species = []
     # Extract all the species names from the model
-    for species in dic['relevant_species']:
+    for species in results.relevant_species:
         if species in all_species_names:
             extracted_species.append(species)
     logger.info("Extracted species: %s", extracted_species)
@@ -136,10 +140,7 @@ class GetAnnotationTool(BaseTool):
 
         # Extract relevant species names based on the user question
         list_species_names = extract_relevant_species_names(model_object, arg_data, state)
-
-        # Check if the returned species names are empty
-        if not list_species_names:
-            raise ValueError("Model does not contain the requested species.")
+        print (list_species_names)
 
         (annotations_df,
          species_without_description) = self._fetch_annotations(list_species_names)

aiagents4pharma/talk2biomodels/tools/query_article.py

@@ -9,7 +9,6 @@ from typing import Type, Annotated
 from pydantic import BaseModel, Field
 from langchain_core.tools import BaseTool
 from langchain_core.vectorstores import InMemoryVectorStore
-from langchain_openai.embeddings import OpenAIEmbeddings
 from langchain_community.document_loaders import PyPDFLoader
 from langgraph.prebuilt import InjectedState
 
@@ -51,8 +50,13 @@ class QueryArticle(BaseTool):
         pages = []
         for page in loader.lazy_load():
             pages.append(page)
+        # Set up text embedding model
+        text_embedding_model = state['text_embedding_model']
+        logging.info("Loaded text embedding model %s", text_embedding_model)
         # Create a vector store from the pages
-        vector_store = InMemoryVectorStore.from_documents(pages, OpenAIEmbeddings())
+        vector_store = InMemoryVectorStore.from_documents(
+                                            pages,
+                                            text_embedding_model)
         # Search the article with the question
         docs = vector_store.similarity_search(question)
         # Return the content of the pages

aiagents4pharma/talk2biomodels/tools/search_models.py

@@ -5,14 +5,18 @@ Tool for searching models based on search query.
 """
 
 from typing import Type, Annotated
+import logging
 from pydantic import BaseModel, Field
 from basico import biomodels
 from langchain_core.tools import BaseTool
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
-from langchain_openai import ChatOpenAI
 from langgraph.prebuilt import InjectedState
 
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 class SearchModelsInput(BaseModel):
     """
     Input schema for the search models tool.
@@ -43,8 +47,10 @@ class SearchModelsTool(BaseTool):
         Returns:
             dict: The answer to the question in the form of a dictionary.
         """
+        logger.log(logging.INFO, "Searching models with the query and model: %s, %s",
+                   query, state['llm_model'])
         search_results = biomodels.search_for_model(query)
-        llm = ChatOpenAI(model=state['llm_model'])
+        llm = state['llm_model']
         # Check if run_manager's metadata has the key 'prompt_content'
         prompt_content = f'''
                         Convert the input into a table.

aiagents4pharma/talk2knowledgegraphs/__init__.py

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

aiagents4pharma/talk2knowledgegraphs/agents/__init__.py

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

aiagents4pharma/talk2knowledgegraphs/agents/t2kg_agent.py

@@ -0,0 +1,85 @@
+'''
+This is the agent file for the Talk2KnowledgeGraphs agent.
+'''
+
+import logging
+from typing import Annotated
+import hydra
+from langchain_ollama import ChatOllama
+from langchain_core.language_models.chat_models import BaseChatModel
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph import START, StateGraph
+from langgraph.prebuilt import create_react_agent, ToolNode, InjectedState
+from ..tools.subgraph_extraction import SubgraphExtractionTool
+from ..tools.subgraph_summarization import SubgraphSummarizationTool
+from ..tools.graphrag_reasoning import GraphRAGReasoningTool
+from ..states.state_talk2knowledgegraphs import Talk2KnowledgeGraphs
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def get_app(uniq_id, llm_model: BaseChatModel=ChatOllama(model='llama3.2:1b', temperature=0.0)):
+    '''
+    This function returns the langraph app.
+    '''
+    def agent_t2kg_node(state: Annotated[dict, InjectedState]):
+        '''
+        This function calls the model.
+        '''
+        logger.log(logging.INFO, "Calling t2kg_agent node with thread_id %s", uniq_id)
+        response = model.invoke(state, {"configurable": {"thread_id": uniq_id}})
+
+        return response
+
+    # Load hydra configuration
+    logger.log(logging.INFO, "Load Hydra configuration for Talk2KnowledgeGraphs agent.")
+    with hydra.initialize(version_base=None, config_path="../configs"):
+        cfg = hydra.compose(config_name='config',
+                            overrides=['agents/t2kg_agent=default'])
+        cfg = cfg.agents.t2kg_agent
+
+    # Define the tools
+    subgraph_extraction = SubgraphExtractionTool()
+    subgraph_summarization = SubgraphSummarizationTool()
+    graphrag_reasoning = GraphRAGReasoningTool()
+    tools = ToolNode([
+                    subgraph_extraction,
+                    subgraph_summarization,
+                    graphrag_reasoning,
+                    ])
+
+    # Create the agent
+    model = create_react_agent(
+                llm_model,
+                tools=tools,
+                state_schema=Talk2KnowledgeGraphs,
+                state_modifier=cfg.state_modifier,
+                checkpointer=MemorySaver()
+            )
+
+    # Define a new graph
+    workflow = StateGraph(Talk2KnowledgeGraphs)
+
+    # Define the two nodes we will cycle between
+    workflow.add_node("agent_t2kg", agent_t2kg_node)
+
+    # Set the entrypoint as the first node
+    # This means that this node is the first one called
+    workflow.add_edge(START, "agent_t2kg")
+
+    # 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 with thread_id %s and llm_model %s",
+               uniq_id,
+               llm_model)
+
+    return app

aiagents4pharma/talk2knowledgegraphs/configs/__init__.py

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