Search models

Tool for searching models based on search query.

SearchModelsInput

Bases: BaseModel

Input schema for the search models tool.

Source code in aiagents4pharma/talk2biomodels/tools/search_models.py

class SearchModelsInput(BaseModel):
    """
    Input schema for the search models tool.
    """

    query: str = Field(description="Search models query", default=None)
    num_query: int = Field(description="Top number of models to search", default=10, le=100)
    tool_call_id: Annotated[str, InjectedToolCallId]

SearchModelsTool

Bases: BaseTool

Tool for returning the search results based on the search query.

Source code in aiagents4pharma/talk2biomodels/tools/search_models.py

class SearchModelsTool(BaseTool):
    """
    Tool for returning the search results based on the search query.
    """

    name: str = "search_models"
    description: str = "Search for only manually curated models in "
    "the BioMmodels database based on keywords."
    args_schema: type[BaseModel] = SearchModelsInput
    return_direct: bool = False

    def _run(
        self,
        tool_call_id: Annotated[str, InjectedToolCallId],
        query: str = None,
        num_query: int = 10,
    ) -> dict:
        """
        Run the tool.

        Args:
            query (str): The search query.
            num_query (int): The number of models to search.
            tool_call_id (str): The tool call ID.

        Returns:
            dict: The answer to the question in the form of a dictionary.
        """
        logger.log(
            logging.INFO,
            "Searching models with the query and number %s, %s",
            query,
            num_query,
        )
        # Search for models based on the query
        search_results = biomodels.search_for_model(query, num_results=num_query)
        # Convert the search results to a pandas DataFrame
        df = pd.DataFrame(search_results)
        # Prepare a message to return
        first_n = min(3, len(search_results))
        content = f"Found {len(search_results)} manually curated models"
        content += f" for the query: {query}."
        # Pass the first 3 models to the LLM
        # to avoid hallucinations
        content += f" Here is the summary of the first {first_n} models:"
        for i in range(first_n):
            content += (
                f"\nModel {i + 1}: {search_results[i]['name']} (ID: {search_results[i]['id']})"
            )
        # Return the updated state of the tool
        return Command(
            update={
                # update the message history
                "messages": [
                    ToolMessage(
                        content=content,
                        tool_call_id=tool_call_id,
                        artifact={"dic_data": df.to_dict(orient="records")},
                    )
                ],
            }
        )

description = 'Search for only manually curated models in ' class-attribute instance-attribute

the BioMmodels database based on keywords.

_run(tool_call_id, query=None, num_query=10)

Run the tool.

Parameters:

Name	Type	Description	Default
query	str	The search query.	None
num_query	int	The number of models to search.	10
tool_call_id	str	The tool call ID.	required

Returns:

Name	Type	Description
dict	dict	The answer to the question in the form of a dictionary.

Source code in aiagents4pharma/talk2biomodels/tools/search_models.py

def _run(
    self,
    tool_call_id: Annotated[str, InjectedToolCallId],
    query: str = None,
    num_query: int = 10,
) -> dict:
    """
    Run the tool.

    Args:
        query (str): The search query.
        num_query (int): The number of models to search.
        tool_call_id (str): The tool call ID.

    Returns:
        dict: The answer to the question in the form of a dictionary.
    """
    logger.log(
        logging.INFO,
        "Searching models with the query and number %s, %s",
        query,
        num_query,
    )
    # Search for models based on the query
    search_results = biomodels.search_for_model(query, num_results=num_query)
    # Convert the search results to a pandas DataFrame
    df = pd.DataFrame(search_results)
    # Prepare a message to return
    first_n = min(3, len(search_results))
    content = f"Found {len(search_results)} manually curated models"
    content += f" for the query: {query}."
    # Pass the first 3 models to the LLM
    # to avoid hallucinations
    content += f" Here is the summary of the first {first_n} models:"
    for i in range(first_n):
        content += (
            f"\nModel {i + 1}: {search_results[i]['name']} (ID: {search_results[i]['id']})"
        )
    # Return the updated state of the tool
    return Command(
        update={
            # update the message history
            "messages": [
                ToolMessage(
                    content=content,
                    tool_call_id=tool_call_id,
                    artifact={"dic_data": df.to_dict(orient="records")},
                )
            ],
        }
    )