Skip to content

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
16
17
18
19
20
21
class SearchModelsInput(BaseModel):
    """
    Input schema for the search models tool.
    """
    query: str = Field(description="Search models query", default=None)
    state: Annotated[dict, InjectedState]

SearchModelsTool

Bases: BaseTool

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

Source code in aiagents4pharma/talk2biomodels/tools/search_models.py
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
class SearchModelsTool(BaseTool):
    """
    Tool for returning the search results based on the search query.
    """
    name: str = "search_models"
    description: str = "Search models in the BioMmodels database based on keywords."
    args_schema: Type[BaseModel] = SearchModelsInput
    return_direct: bool = True

    def _run(self,
             query: str,
             state: Annotated[dict, InjectedState]) -> dict:
        """
        Run the tool.

        Args:
            query (str): The search query.

        Returns:
            dict: The answer to the question in the form of a dictionary.
        """
        search_results = biomodels.search_for_model(query)
        llm = ChatOpenAI(model=state['llm_model'])
        # Check if run_manager's metadata has the key 'prompt_content'
        prompt_content = f'''
                        Convert the input into a table.

                        The table must contain the following columns:
                        - #
                        - BioModel ID
                        - BioModel Name
                        - Format
                        - Submission Date

                        Additional Guidelines:
                        - The column # must contain the row number starting from 1.
                        - Embed the url for each BioModel ID in the table 
                        in the first column in the markdown format.
                        - The Submission Date must be in the format YYYY-MM-DD.

                        Input:
                        {input}.
                        '''
        prompt_template = ChatPromptTemplate.from_messages(
            [("system", prompt_content),
             ("user", "{input}")]
        )
        parser = StrOutputParser()
        chain = prompt_template | llm | parser
        return chain.invoke({"input": search_results})

_run(query, state)

Run the tool.

Parameters:

Name Type Description Default
query str

The search query.

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
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
def _run(self,
         query: str,
         state: Annotated[dict, InjectedState]) -> dict:
    """
    Run the tool.

    Args:
        query (str): The search query.

    Returns:
        dict: The answer to the question in the form of a dictionary.
    """
    search_results = biomodels.search_for_model(query)
    llm = ChatOpenAI(model=state['llm_model'])
    # Check if run_manager's metadata has the key 'prompt_content'
    prompt_content = f'''
                    Convert the input into a table.

                    The table must contain the following columns:
                    - #
                    - BioModel ID
                    - BioModel Name
                    - Format
                    - Submission Date

                    Additional Guidelines:
                    - The column # must contain the row number starting from 1.
                    - Embed the url for each BioModel ID in the table 
                    in the first column in the markdown format.
                    - The Submission Date must be in the format YYYY-MM-DD.

                    Input:
                    {input}.
                    '''
    prompt_template = ChatPromptTemplate.from_messages(
        [("system", prompt_content),
         ("user", "{input}")]
    )
    parser = StrOutputParser()
    chain = prompt_template | llm | parser
    return chain.invoke({"input": search_results})