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

aiagents4pharma/talk2biomodels/__init__.py

@@ -5,3 +5,4 @@ from . import models
 from . import tools
 from . import agents
 from . import states
+from . import api

aiagents4pharma/talk2biomodels/agents/t2b_agent.py

@@ -15,6 +15,7 @@ from ..tools.search_models import SearchModelsTool
 from ..tools.get_modelinfo import GetModelInfoTool
 from ..tools.simulate_model import SimulateModelTool
 from ..tools.custom_plotter import CustomPlotterTool
+from ..tools.get_annotation import GetAnnotationTool
 from ..tools.ask_question import AskQuestionTool
 from ..tools.parameter_scan import ParameterScanTool
 from ..tools.steady_state import SteadyStateTool
@@ -44,8 +45,9 @@ def get_app(uniq_id, llm_model='gpt-4o-mini'):
                     SearchModelsTool(),
                     GetModelInfoTool(),
                     SteadyStateTool(),
-                    ParameterScanTool()
-                    ])
+                    ParameterScanTool(),
+                    GetAnnotationTool()
+                ])
 
     # Define the model
     llm = ChatOpenAI(model=llm_model, temperature=0)

aiagents4pharma/talk2biomodels/api/__init__.py

@@ -0,0 +1,6 @@
+'''
+This file is used to import the modules in the package.
+'''
+from . import uniprot
+from . import ols
+from . import kegg

aiagents4pharma/talk2biomodels/api/kegg.py

@@ -0,0 +1,83 @@
+"""
+This module contains the API for fetching Kegg database
+"""
+import re
+from typing import List, Dict
+import requests
+
+def fetch_from_api(base_url: str, query: str) -> str:
+    """Fetch data from the given API endpoint."""
+    try:
+        response = requests.get(base_url + query, timeout=10)
+        response.raise_for_status()
+        return response.text
+    except requests.exceptions.RequestException as e:
+        print(f"Error fetching data for query {query}: {e}")
+        return ""
+
+def fetch_kegg_names(ids: List[str], batch_size: int = 10) -> Dict[str, str]:
+    """
+    Fetch the names of multiple KEGG entries using the KEGG REST API in batches.
+
+    Args:
+        ids (List[str]): List of KEGG IDs.
+        batch_size (int): Maximum number of IDs to include in a single request.
+
+    Returns:
+        Dict[str, str]: A mapping of KEGG IDs to their names.
+    """
+    if not ids:
+        return {}
+
+    base_url = "https://rest.kegg.jp/get/"
+    entry_name_map = {}
+
+    # Process IDs in batches
+    for i in range(0, len(ids), batch_size):
+        batch = ids[i:i + batch_size]
+        query = "+".join(batch)
+        entry_data = fetch_from_api(base_url, query)
+
+        # if not entry_data:
+        #     continue
+        entries = entry_data.split("///")
+        for entry in entries:
+            if not entry.strip():
+                continue
+            lines = entry.strip().split("\n")
+            entry_line = next((line for line in lines
+                                if line.startswith("ENTRY")), None)
+            name_line = next((line for line in lines
+                                if line.startswith("NAME")), None)
+
+            # if not entry_line and not name_line:
+            #     continue
+            entry_id = entry_line.split()[1]
+            # Split multiple names in the NAME field and clean them
+            names = [
+                re.sub(r'[^a-zA-Z0-9\s]', '', name).strip()
+                for name in name_line.replace("NAME", "").strip().split(";")
+            ]
+            # Join cleaned names into a single string
+            entry_name_map[entry_id] = " ".join(names).strip()
+
+    return entry_name_map
+
+def fetch_kegg_annotations(data: List[Dict[str, str]],
+                           batch_size: int = 10) -> Dict[str, Dict[str, str]]:
+    """Fetch KEGG entry descriptions grouped by database type."""
+    grouped_data = {}
+    for entry in data:
+        db_type = entry["Database"].lower()
+        grouped_data.setdefault(db_type, []).append(entry["Id"])
+
+    results = {}
+    for db_type, ids in grouped_data.items():
+        results[db_type] = fetch_kegg_names(ids, batch_size=batch_size)
+
+    return results
+
+# def get_protein_name_or_label(data: List[Dict[str, str]],
+#                               batch_size: int = 10) -> Dict[str, Dict[str, str]]:
+#     """Fetch descriptions for KEGG-related identifiers."""
+#     return fetch_kegg_annotations(data, batch_size=batch_size)

aiagents4pharma/talk2biomodels/api/ols.py

@@ -0,0 +1,72 @@
+"""
+This module contains the API for fetching ols database
+"""
+from typing import List, Dict
+import requests
+
+def fetch_from_ols(term: str) -> str:
+    """
+    Fetch the label for a single term from OLS.
+
+    Args:
+        term (str): The term in the format "ONTOLOGY:TERM_ID".
+
+    Returns:
+        str: The label for the term or an error message.
+    """
+    try:
+        ontology, _ = term.split(":")
+        base_url = f"https://www.ebi.ac.uk/ols4/api/ontologies/{ontology.lower()}/terms"
+        params = {"obo_id": term}
+        response = requests.get(
+            base_url,
+            params=params,
+            headers={"Accept": "application/json"},
+            timeout=10
+        )
+        response.raise_for_status()
+        data = response.json()
+        label = '-'
+        # Extract and return the label
+        if "_embedded" in data and "terms" in data["_embedded"] \
+             and len(data["_embedded"]["terms"]) > 0:
+            label = data["_embedded"]["terms"][0].get("label", "Label not found")
+        return label
+    except (requests.exceptions.RequestException, KeyError, IndexError) as e:
+        return f"Error: {str(e)}"
+
+def fetch_ols_labels(terms: List[str]) -> Dict[str, str]:
+    """
+    Fetch labels for multiple terms from OLS.
+
+    Args:
+        terms (List[str]): A list of terms in the format "ONTOLOGY:TERM_ID".
+
+    Returns:
+        Dict[str, str]: A mapping of term IDs to their labels or error messages.
+    """
+    results = {}
+    for term in terms:
+        results[term] = fetch_from_ols(term)
+    return results
+
+def search_ols_labels(data: List[Dict[str, str]]) -> Dict[str, Dict[str, str]]:
+    """
+    Fetch OLS annotations grouped by ontology type.
+
+    Args:
+        data (List[Dict[str, str]]): A list of dictionaries containing 'Id' and 'Database'.
+
+    Returns:
+        Dict[str, Dict[str, str]]: A mapping of ontology type to term labels.
+    """
+    grouped_data = {}
+    for entry in data:
+        ontology = entry["Database"].lower()
+        grouped_data.setdefault(ontology, []).append(entry["Id"])
+
+    results = {}
+    for ontology, terms in grouped_data.items():
+        results[ontology] = fetch_ols_labels(terms)
+
+    return results

aiagents4pharma/talk2biomodels/api/uniprot.py

@@ -0,0 +1,35 @@
+"""
+This module contains the API for fetching uniprot database
+"""
+from typing import List, Dict
+import requests
+
+def search_uniprot_labels(identifiers: List[str]) -> Dict[str, str]:
+    """
+    Fetch protein names or labels for a list of UniProt identifiers by making sequential requests.
+
+    Args:
+        identifiers (List[str]): A list of UniProt identifiers.
+
+    Returns:
+        Dict[str, str]: A mapping of UniProt identifiers to their protein names or error messages.
+    """
+    results = {}
+    base_url = "https://www.uniprot.org/uniprot/"
+
+    for identifier in identifiers:
+        url = f"{base_url}{identifier}.json"
+        try:
+            response = requests.get(url, timeout=10)
+            response.raise_for_status()
+            data = response.json()
+            protein_name = (
+                data.get('proteinDescription', {})
+                .get('recommendedName', {})
+                .get('fullName', {})
+                .get('value', 'Name not found')
+            )
+            results[identifier] = protein_name
+        except requests.exceptions.RequestException as e:
+            results[identifier] = f"Error: {str(e)}"
+    return results

aiagents4pharma/talk2biomodels/states/state_talk2biomodels.py

@@ -23,3 +23,4 @@ class Talk2Biomodels(AgentState):
     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]
+    dic_annotations_data : Annotated[list[dict], operator.add]

aiagents4pharma/talk2biomodels/tests/test_api.py

@@ -0,0 +1,57 @@
+'''
+Test cases for Talk2Biomodels.
+'''
+
+from ..api.uniprot import search_uniprot_labels
+from ..api.ols import fetch_from_ols
+from ..api.kegg import fetch_kegg_names, fetch_from_api
+
+def test_search_uniprot_labels():
+    '''
+    Test the search_uniprot_labels function.
+    '''
+    # "P61764" = Positive result, "P0000Q" = negative result
+    identifiers = ["P61764", "P0000Q"]
+    results = search_uniprot_labels(identifiers)
+    assert results["P61764"] == "Syntaxin-binding protein 1"
+    assert results["P0000Q"].startswith("Error: 400")
+
+def test_fetch_from_ols():
+    '''
+    Test the fetch_from_ols function.
+    '''
+    term_1 = "GO:0005886" #Positive result
+    term_2 = "GO:ABC123" #Negative result
+    label_1 = fetch_from_ols(term_1)
+    label_2 = fetch_from_ols(term_2)
+    assert isinstance(label_1, str), f"Expected string, got {type(label_1)}"
+    assert isinstance(label_2, str), f"Expected string, got {type(label_2)}"
+    assert label_1 == "plasma membrane"
+    assert label_2.startswith("Error: 404")
+
+def test_fetch_kegg_names():
+    '''
+    Test the fetch_kegg_names function.
+    '''
+    ids = ["C00001", "C00002"]
+    results = fetch_kegg_names(ids)
+    assert results["C00001"] == "H2O"
+    assert results["C00002"] == "ATP"
+
+    # Try with an empty list
+    results = fetch_kegg_names([])
+    assert not results
+
+def test_fetch_from_api():
+    '''
+    Test the fetch_from_api function.
+    '''
+    base_url = "https://rest.kegg.jp/get/"
+    query = "C00001"
+    entry_data = fetch_from_api(base_url, query)
+    assert entry_data.startswith("ENTRY       C00001")
+
+    # Try with an invalid query
+    query = "C0000Q"
+    entry_data = fetch_from_api(base_url, query)
+    assert not entry_data

aiagents4pharma/talk2biomodels/tests/test_get_annotation.py

@@ -0,0 +1,173 @@
+'''
+Test cases for Talk2Biomodels.
+'''
+import random
+import pytest
+from langchain_core.messages import HumanMessage, ToolMessage
+from ..agents.t2b_agent import get_app
+from ..tools.get_annotation import prepare_content_msg
+
+@pytest.fixture(name="make_graph")
+def make_graph_fixture():
+    '''
+    Create an instance of the talk2biomodels agent.
+    '''
+    unique_id = random.randint(1000, 9999)
+    graph = get_app(unique_id)
+    config = {"configurable": {"thread_id": unique_id}}
+    return graph, config
+
+def test_species_list(make_graph):
+    '''
+    Test the tool by passing species names.
+    '''
+    # Test with a valid species name
+    app, config = make_graph
+    prompt = "Extract annotations of species IL6 in model 537."
+    app.invoke(
+                {"messages": [HumanMessage(content=prompt)]},
+                config=config
+            )
+    current_state = app.get_state(config)
+    # print (current_state.values["dic_annotations_data"])
+    dic_annotations_data = current_state.values["dic_annotations_data"]
+
+    # The assert statement checks if IL6 is present in the returned annotations.
+    assert dic_annotations_data[0]['data']["Species Name"][0] == "IL6"
+
+    # Test with an invalid species name
+    app, config = make_graph
+    prompt = "Extract annotations of species NADH in model 537."
+    app.invoke(
+        {"messages": [HumanMessage(content=prompt)]},
+        config=config
+    )
+    current_state = app.get_state(config)
+    reversed_messages = current_state.values["messages"][::-1]
+    # Loop through the reversed messages until a
+    # ToolMessage is found.
+
+    test_condition = False
+    for msg in reversed_messages:
+        # Assert that the one of the messages is a ToolMessage
+        # and its artifact is None.
+        if isinstance(msg, ToolMessage) and msg.name == "get_annotation":
+            #If a ToolMessage exists and artifact is None (meaning no valid annotation was found)
+            #and the rejected species (NADH) is mentioned, the test passes.
+            if msg.artifact is None and 'NADH' in msg.content:
+                #If artifact is None, it means no annotation was found
+                # (likely due to an invalid species).
+                #If artifact contains data, the tool successfully retrieved annotations.
+                test_condition = True
+                break
+    # assert test_condition
+    assert test_condition, "Expected rejection message for NADH but did not find it."
+
+    # Test with an invalid species name and a valid species name
+    app, config = make_graph
+    prompt = "Extract annotations of species NADH, NAD, and IL7 in model 64."
+    app.invoke(
+        {"messages": [HumanMessage(content=prompt)]},
+        config=config
+    )
+    current_state = app.get_state(config)
+    # dic_annotations_data = current_state.values["dic_annotations_data"]
+    reversed_messages = current_state.values["messages"][::-1]
+    # Loop through the reversed messages until a
+    # ToolMessage is found.
+    artifact_was_none = False
+    for msg in reversed_messages:
+        # Assert that the one of the messages is a ToolMessage
+        # and its artifact is None.
+        if isinstance(msg, ToolMessage) and msg.name == "get_annotation":
+            # print (msg.artifact, msg.content)
+
+            if msg.artifact is True and 'IL7' in msg.content:
+                artifact_was_none = True
+                break
+    assert artifact_was_none
+
+def test_all_species(make_graph):
+    '''
+    Test the tool by asking for annotations of all species is specific models.
+
+    model 12 contains species with no URL.
+    model 20 contains species without description.
+    model 56 contains species with database outside of KEGG, UniProt, and OLS.
+    '''
+    # Test valid models
+    for model_id in [12, 20, 56]:
+        app, config = make_graph
+        prompt = f"Extract annotations of all species model {model_id}."
+        # Test the tool get_modelinfo
+        app.invoke({"messages": [HumanMessage(content=prompt)]},
+                            config=config
+                        )
+        #print(response["messages"])
+        # assistant_msg = response["messages"][-1].content
+
+        current_state = app.get_state(config)
+
+        reversed_messages = current_state.values["messages"][::-1]
+        # Coveres all of the use cases for the expecetd sting on all the species
+        test_condition = False
+        for msg in reversed_messages:
+            if isinstance(msg, ToolMessage) and msg.name == "get_annotation":
+                if model_id == 12:
+                # For model 12:
+                # Expect a successful extraction (artifact is True) and that the content
+                # matches what is returned by prepare_content_msg for species ['lac'].
+                    if (msg.artifact is True and msg.content == prepare_content_msg(['lac'],[])
+                        and msg.status=="success"):
+                        test_condition = True
+                        break
+
+                if model_id == 20:
+                # For model 20:
+                # Expect an error message containing a note that species extraction failed.
+                    if ("Unable to extract species from the model"
+                        in msg.content and msg.status == "error"):
+                        test_condition = True
+                        break
+
+                if model_id == 56:
+                    # For model 56:
+                    # Expect a successful extraction (artifact is True) and that the content
+                    # matches for for missing description ['ORI'].
+                    if (msg.artifact is True and
+                    msg.content == prepare_content_msg([],['ORI'])
+                    and msg.status == "success"):
+                        test_condition = True
+                        break
+
+    # Retrieve the dictionary that holds all the annotation data from the app's state
+    dic_annotations_data = current_state.values["dic_annotations_data"]
+
+    assert isinstance(dic_annotations_data, list),\
+        f"Expected a list for model {model_id}, got {type(dic_annotations_data)}"
+    assert len(dic_annotations_data) > 0,\
+        f"Expected species data for model {model_id}, but got empty list"
+    assert test_condition # Expected output is validated
+
+    # Test case where no model is specified
+    app, config = make_graph
+    prompt = "Extract annotations of all species."
+    app.invoke({"messages": [HumanMessage(content=prompt)]},
+                        config=config
+                    )
+    current_state = app.get_state(config)
+    # dic_annotations_data = current_state.values["dic_annotations_data"]
+    reversed_messages = current_state.values["messages"][::-1]
+    print(reversed_messages)
+
+    test_condition = False
+    for msg in reversed_messages:
+        # Assert that the one of the messages is a ToolMessage
+        if isinstance(msg, ToolMessage) and msg.name == "get_annotation":
+            if "Error:" in msg.content and msg.status == "error":
+                test_condition = True
+                break
+    # Loop through the reversed messages until a
+    # ToolMessage is found.
+    # Ensure the system correctly informs the user to specify a model
+    assert test_condition, "Expected error message when no model is specified was not found."

aiagents4pharma/talk2biomodels/tools/__init__.py

@@ -7,4 +7,6 @@ from . import ask_question
 from . import custom_plotter
 from . import get_modelinfo
 from . import parameter_scan
+from . import steady_state
 from . import load_biomodel
+from . import get_annotation

aiagents4pharma/talk2biomodels/tools/get_annotation.py

@@ -0,0 +1,289 @@
+#!/usr/bin/env python3
+
+"""
+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
+import logging
+from pydantic import BaseModel, Field
+import basico
+import pandas as pd
+from langgraph.types import Command
+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 .load_biomodel import ModelData, load_biomodel
+from ..api.uniprot import search_uniprot_labels
+from ..api.ols import search_ols_labels
+from ..api.kegg import fetch_kegg_annotations
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+ols_ontology_abbreviations = {'pato', 'chebi', 'sbo', 'fma', 'pr','go'}
+
+def prepare_content_msg(species_not_found: List[str],
+                        species_without_description: List[str]):
+    """
+    Prepare the content message.
+    """
+    content = 'Successfully extracted annotations for the species.'
+    if species_not_found:
+        content += f'''The following species do not exist, and
+                        hence their annotations were not extracted:
+                        {', '.join(species_not_found)}.'''
+    if species_without_description:
+        content += f'''The descriptions for the following species
+                        were not found:
+                        {", ".join(species_without_description)}.'''
+    return content
+
+class GetAnnotationInput(BaseModel):
+    """
+    Input schema for annotation tool.
+    """
+    sys_bio_model: ModelData = Field(description="model data")
+    tool_call_id: Annotated[str, InjectedToolCallId]
+    list_species_names: List[str] = Field(
+        default=[],
+        description='''List of species names to fetch annotations for.
+                      If not provided, annotations for all
+                      species in the model will be fetched.'''
+    )
+    state: Annotated[dict, InjectedState]
+
+class GetAnnotationTool(BaseTool):
+    """
+    Tool for fetching species annotations based on the provided model and species names.
+    """
+    name: str = "get_annotation"
+    description: str = '''A tool to extract annotations for a list of species names
+                        based on the provided model. Annotations include
+                        the species name, description, database, ID, link,
+                        and qualifier. The tool can handle multiple species
+                        in a single invoke.'''
+    args_schema: Type[BaseModel] = GetAnnotationInput
+    return_direct: bool = False
+
+    def _run(self,
+             tool_call_id: Annotated[str, InjectedToolCallId],
+             state: Annotated[dict, InjectedState],
+             list_species_names: List[str] = None,
+             sys_bio_model: ModelData = None) -> str:
+        """
+        Run the tool.
+        """
+        logger.info("Running the GetAnnotationTool tool for species %s", list_species_names)
+
+        # Prepare the model object
+        sbml_file_path = state['sbml_file_path'][-1] if state['sbml_file_path'] else None
+        model_object = load_biomodel(sys_bio_model, sbml_file_path=sbml_file_path)
+
+        # Extract all the species names from the model
+        df_species = basico.model_info.get_species(model=model_object.copasi_model)
+
+        if df_species is None:
+            # for example this may happen with model 20
+            raise ValueError("Unable to extract species from the model.")
+        # Fetch annotations for the species names
+        list_species_names = list_species_names or df_species.index.tolist()
+
+        (annotations_df,
+         species_not_found,
+         species_without_description) = self._fetch_annotations(list_species_names)
+
+        # Check if annotations are empty
+        # If empty, return a message
+        if annotations_df.empty:
+            logger.warning("The annotations dataframe is empty.")
+            return prepare_content_msg(species_not_found, species_without_description)
+
+        # Process annotations
+        annotations_df = self._process_annotations(annotations_df)
+
+        # Prepare the simulated data
+        dic_annotations_data = {
+            'source': sys_bio_model.biomodel_id if sys_bio_model.biomodel_id else 'upload',
+            'tool_call_id': tool_call_id,
+            'data': annotations_df.to_dict()
+        }
+
+        # Update the state with the annotations data
+        dic_updated_state_for_model = {}
+        for key, value in {
+            "model_id": [sys_bio_model.biomodel_id],
+            "sbml_file_path": [sbml_file_path],
+            "dic_annotations_data": [dic_annotations_data]
+        }.items():
+            if value:
+                dic_updated_state_for_model[key] = value
+
+        return Command(
+            update=dic_updated_state_for_model | {
+                "messages": [
+                    ToolMessage(
+                        content=prepare_content_msg(species_not_found,
+                                                    species_without_description),
+                        artifact=True,
+                        tool_call_id=tool_call_id
+                    )
+                ]
+            }
+        )
+
+    def _fetch_annotations(self, list_species_names: List[str]) -> tuple:
+        """
+        Fetch annotations for the given species names from the model.
+        In this method, we fetch the MIRIAM annotations for the species names.
+        If the annotation is not found, we add the species to the list of
+        species not found. If the annotation is found, we extract the descriptions
+        from the annotation and add them to the data list.
+
+        Args:
+            list_species_names (List[str]): List of species names to fetch annotations for.
+
+        Returns:
+            tuple: A tuple containing the annotations dataframe, species not found list,
+                   and description not found list.
+        """
+        species_not_found = []
+        description_not_found = []
+        data = []
+
+        # Loop through the species names
+        for species in list_species_names:
+            # Get the MIRIAM annotation for the species
+            annotation = basico.get_miriam_annotation(name=species)
+            # If the annotation is not found, add the species to the list
+            if annotation is None:
+                species_not_found.append(species)
+                continue
+
+            # Extract the descriptions from the annotation
+            descriptions = annotation.get("descriptions", [])
+
+            if descriptions == []:
+                description_not_found.append(species)
+                continue
+
+            # Loop through the descriptions and add them to the data list
+            for desc in descriptions:
+                data.append({
+                    "Species Name": species,
+                    "Link": desc["id"],
+                    "Qualifier": desc["qualifier"]
+                })
+
+        # Create a dataframe from the data list
+        annotations_df = pd.DataFrame(data)
+
+        # Return the annotations dataframe and the species not found list
+        return annotations_df, species_not_found, description_not_found
+
+    def _process_annotations(self, annotations_df: pd.DataFrame) -> pd.DataFrame:
+        """
+        Process annotations dataframe to add additional information.
+        In this method, we add a new column for the ID, a new column for the database,
+        and a new column for the description. We then reorder the columns and process
+        the link to format it correctly.
+
+        Args:
+            annotations_df (pd.DataFrame): Annotations dataframe to process.
+
+        Returns:
+            pd.DataFrame: Processed annotations dataframe
+        """
+        logger.info("Processing annotations.")
+        # Add a new column for the ID
+        # Get the ID from the link key
+        annotations_df['Id'] = annotations_df['Link'].str.split('/').str[-1]
+
+        # Add a new column for the database
+        # Get the database from the link key
+        annotations_df['Database'] = annotations_df['Link'].str.split('/').str[-2]
+
+        # Fetch descriptions for the IDs based on the database type
+        # by qyerying the respective APIs
+        identifiers = annotations_df[['Id', 'Database']].to_dict(orient='records')
+        descriptions = self._fetch_descriptions(identifiers)
+
+        # Add a new column for the description
+        # Get the description from the descriptions dictionary
+        # based on the ID. If the description is not found, use '-'
+        annotations_df['Description'] = annotations_df['Id'].apply(lambda x:
+                                                                   descriptions.get(x, '-'))
+        # annotations_df.index = annotations_df.index + 1
+
+        # Reorder the columns
+        annotations_df = annotations_df[
+            ["Species Name", "Description", "Database", "Id", "Link", "Qualifier"]
+        ]
+
+        # Process the link to format it correctly
+        annotations_df["Link"] = annotations_df["Link"].apply(self._process_link)
+
+        # Return the processed annotations dataframe
+        return annotations_df
+
+    def _process_link(self, link: str) -> str:
+        """
+        Process link to format it correctly.
+        """
+        for ols_ontology_abbreviation in ols_ontology_abbreviations:
+            if ols_ontology_abbreviation +'/' in link:
+                link = link.replace(f"{ols_ontology_abbreviation}/", "")
+        if "kegg.compound" in link:
+            link = link.replace("kegg.compound/", "kegg.compound:")
+        return link
+
+    def _fetch_descriptions(self, data: List[dict[str, str]]) -> dict[str, str]:
+        """
+        Fetch protein names or labels based on the database type.
+        """
+        logger.info("Fetching descriptions for the IDs.")
+        results = {}
+        grouped_data = {}
+
+        # In the following loop, we create a dictionary with database as the key
+        # and a list of identifiers as the value. If either the database or the
+        # identifier is NaN, we set it to None.
+        for entry in data:
+            identifier = entry.get('Id')
+            database = entry.get('Database')
+            # Check if database is NaN
+            if isinstance(database, float):
+                if math.isnan(database):
+                    database = None
+                    results[identifier or "unknown"] = "-"
+            else:
+                database = database.lower()
+                grouped_data.setdefault(database, []).append(identifier)
+
+        # In the following loop, we fetch the descriptions for the identifiers
+        # based on the database type.
+        # Constants
+
+        for database, identifiers in grouped_data.items():
+            if database == 'uniprot':
+                results.update(search_uniprot_labels(identifiers))
+            elif database in ols_ontology_abbreviations:
+                annotations = search_ols_labels([
+                        {"Id": id_, "Database": database}
+                        for id_ in identifiers
+                    ])
+                for identifier in identifiers:
+                    results[identifier] = annotations.get(database, {}).get(identifier, "-")
+            elif database == 'kegg.compound':
+                data = [{"Id": identifier, "Database": "kegg.compound"}
+                        for identifier in identifiers]
+                annotations = fetch_kegg_annotations(data)
+                for identifier in identifiers:
+                    results[identifier] = annotations.get(database, {}).get(identifier, "-")
+            else:
+                for identifier in identifiers:
+                    results[identifier] = "-"
+        return results

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: aiagents4pharma
-Version: 1.13.0
+Version: 1.14.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

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

@@ -5,21 +5,28 @@ aiagents4pharma/configs/talk2biomodels/__init__.py,sha256=5ah__-8XyRblwT0U1ByRig
 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=Oi89_BbxfQc6SGW1pC-hyZMqOIkiAMOlNwpCa4VCXk0,327
-aiagents4pharma/talk2biomodels/__init__.py,sha256=qUw3qXrENqSCLIKSLy_qtNPwPDTb1wdZ8fZispcHb3g,141
+aiagents4pharma/talk2biomodels/__init__.py,sha256=2ICwVh1u07SZv31Jd2DKHobauOxWNWY29_Gqq3kOnNQ,159
 aiagents4pharma/talk2biomodels/agents/__init__.py,sha256=sn5-fREjMdEvb-OUan3iOqrgYGjplNx3J8hYOaW0Po8,128
-aiagents4pharma/talk2biomodels/agents/t2b_agent.py,sha256=8_2D3uknPjYbzKj7IDhC8xnz_HEvXck0b4RCJxoAxRs,3276
+aiagents4pharma/talk2biomodels/agents/t2b_agent.py,sha256=13aSlBZBWtjXOLq7c99u33c923fi2Ab0VW--eX5gF-o,3366
+aiagents4pharma/talk2biomodels/api/__init__.py,sha256=_GmDQqDLYpsUPUeE1nBNlT5AI9oTXIcqgOfNfvmonqA,123
+aiagents4pharma/talk2biomodels/api/kegg.py,sha256=QzYDAfJ16E7tbHGxP8ZNWRizMkMRS_HJuucueXEC1Gg,2943
+aiagents4pharma/talk2biomodels/api/ols.py,sha256=qq0Qy-gJDxanQW-HfCChDsTQsY1M41ua8hMlTnfuzrA,2202
+aiagents4pharma/talk2biomodels/api/uniprot.py,sha256=aPUAVBR7UYXDuuhDpKezAK2aTMzo-NxFYFq6C0W5u6U,1175
 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=XcrCsBbQgc4LCZ7ihpY5tcLQ82Mq17XhRXD1DW8Ep7U,887
+aiagents4pharma/talk2biomodels/states/state_talk2biomodels.py,sha256=44vRJwc2NlVLQQr1Smipr02YzXHOyqUVSNlg7rFjli0,950
 aiagents4pharma/talk2biomodels/tests/__init__.py,sha256=Jbw5tJxSrjGoaK5IX3pJWDCNzhrVQ10lkYq2oQ_KQD8,45
+aiagents4pharma/talk2biomodels/tests/test_api.py,sha256=7Kz2r5F5tjmn3F0LoM33oP-21W633936YHiyf5toGg0,1716
 aiagents4pharma/talk2biomodels/tests/test_basico_model.py,sha256=y82fpTJMPHwtXxlle1cGQ_2Bewwpxi0aJSVrVAYLhN0,2060
+aiagents4pharma/talk2biomodels/tests/test_get_annotation.py,sha256=lgdWNl6g1hiTcUbcmgn2bUk5_-8EUpSNa0MMpIMGeDA,7301
 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/__init__.py,sha256=ZiOdSFaeHW6y3hdtBfsKf0vSb3MuCLuy9MDyjARggb4,322
 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_annotation.py,sha256=Ifbbz08YFI1ifAy3t0tYkr45-k7inO8lZePvCSe5ZfA,11835
 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
@@ -77,8 +84,8 @@ aiagents4pharma/talk2knowledgegraphs/utils/embeddings/sentence_transformer.py,sh
 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,,
+aiagents4pharma-1.14.0.dist-info/LICENSE,sha256=IcIbyB1Hyk5ZDah03VNQvJkbNk2hkBCDqQ8qtnCvB4Q,1077
+aiagents4pharma-1.14.0.dist-info/METADATA,sha256=yGCN0sRc2dixfRh-UAPk2uLZLBjb5WUaSd0VwROU7qY,8609
+aiagents4pharma-1.14.0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+aiagents4pharma-1.14.0.dist-info/top_level.txt,sha256=-AH8rMmrSnJtq7HaAObS78UU-cTCwvX660dSxeM7a0A,16
+aiagents4pharma-1.14.0.dist-info/RECORD,,