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')}
                        )
                    ],
                }
        )