Skip to content

Save model

Tool for saving models.

SaveModelInput

Bases: BaseModel

Input schema for the save model tool.

Source code in aiagents4pharma/talk2biomodels/tools/save_model.py
23
24
25
26
27
28
29
30
31
32
33
34
35
36
class SaveModelInput(BaseModel):
    """
    Input schema for the save model tool.
    """

    path_to_folder: str = Field(
        description="Path to folder to save the model. Keep it to . if not provided.", default="."
    )
    output_filename: str = Field(
        description="Filename to save the model as. Default is 'saved_model.xml'.",
        default="saved_model.xml",
    )
    tool_call_id: Annotated[str, InjectedToolCallId]
    state: Annotated[dict, InjectedState]

SaveModelTool

Bases: BaseTool

Tool for saving a model.

Source code in aiagents4pharma/talk2biomodels/tools/save_model.py
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
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
class SaveModelTool(BaseTool):
    """
    Tool for saving a model.
    """

    name: str = "save_model"
    description: str = "A tool to save the current biomodel to a \
                        user specified path with the default filename\
                         'saved_model.xml'"
    args_schema: type[BaseModel] = SaveModelInput
    return_direct: bool = False

    def _run(
        self,
        tool_call_id: Annotated[str, InjectedToolCallId],
        state: Annotated[dict, InjectedState],
        path_to_folder: str = ".",
        output_filename: str = "saved_model.xml",
    ) -> Command:
        """
        Run the tool.

        Args:
            path (str): The path to save the model.
            tool_call_id (str): The tool call ID.

        Returns:

        """
        logger.log(
            logging.INFO,
            "Saving model to path: %s with filename: %s",
            path_to_folder,
            output_filename,
        )
        # Check if path does not exist
        if not os.path.exists(path_to_folder):
            content = f"Error: Path {path_to_folder} does not exist."
            logger.error(content)
        else:
            logger.info("Saving now")
            # Save the model to the specified path
            with open(os.path.join(path_to_folder, output_filename), "w", encoding="utf-8") as f:
                f.write(state["model_as_string"][-1])
            content = f"Model saved successfully to {path_to_folder}/{output_filename}."
            logger.info(content)
        # Return the updated state of the tool
        return Command(
            update={
                # update the message history
                "messages": [
                    ToolMessage(
                        content=content,
                        tool_call_id=tool_call_id,
                    )
                ],
            }
        )

_run(tool_call_id, state, path_to_folder='.', output_filename='saved_model.xml')

Run the tool.

Parameters:

Name Type Description Default
path str

The path to save the model.

required
tool_call_id str

The tool call ID.

required

Returns:

Source code in aiagents4pharma/talk2biomodels/tools/save_model.py
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
def _run(
    self,
    tool_call_id: Annotated[str, InjectedToolCallId],
    state: Annotated[dict, InjectedState],
    path_to_folder: str = ".",
    output_filename: str = "saved_model.xml",
) -> Command:
    """
    Run the tool.

    Args:
        path (str): The path to save the model.
        tool_call_id (str): The tool call ID.

    Returns:

    """
    logger.log(
        logging.INFO,
        "Saving model to path: %s with filename: %s",
        path_to_folder,
        output_filename,
    )
    # Check if path does not exist
    if not os.path.exists(path_to_folder):
        content = f"Error: Path {path_to_folder} does not exist."
        logger.error(content)
    else:
        logger.info("Saving now")
        # Save the model to the specified path
        with open(os.path.join(path_to_folder, output_filename), "w", encoding="utf-8") as f:
            f.write(state["model_as_string"][-1])
        content = f"Model saved successfully to {path_to_folder}/{output_filename}."
        logger.info(content)
    # Return the updated state of the tool
    return Command(
        update={
            # update the message history
            "messages": [
                ToolMessage(
                    content=content,
                    tool_call_id=tool_call_id,
                )
            ],
        }
    )