From 55f3ceb0edda32cdd0b2ece9fa4e724925024ae3 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Sat, 22 Jun 2024 18:02:22 -0300 Subject: [PATCH] refactor: Update PromptComponent post_code_processing method Refactor the post_code_processing method in the PromptComponent class to improve code readability and maintainability. The changes include: - Import the process_prompt_template function from langflow.base.prompts.api_utils. - Import the update_template_values function from langflow.template.utils. - Call the process_prompt_template function to update the template value. - Update the frontend_node with any values set in the current_build_config. - Remove unnecessary comments and whitespace. These changes enhance the overall quality of the code and make it easier to understand and maintain. --- .../langflow/components/prompts/Prompt.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/src/backend/base/langflow/components/prompts/Prompt.py b/src/backend/base/langflow/components/prompts/Prompt.py index 6761cf1fe..8f8d3d6a2 100644 --- a/src/backend/base/langflow/components/prompts/Prompt.py +++ b/src/backend/base/langflow/components/prompts/Prompt.py @@ -1,6 +1,8 @@ +from langflow.base.prompts.api_utils import process_prompt_template from langflow.custom import Component from langflow.io import Output, PromptInput from langflow.schema.message import Message +from langflow.template.utils import update_template_values class PromptComponent(Component): @@ -23,3 +25,20 @@ class PromptComponent(Component): prompt = await Message.from_template_and_variables(**self._attributes) self.status = prompt.text return prompt + + def post_code_processing(self, new_build_config: dict, current_build_config: dict): + """ + This function is called after the code validation is done. + """ + frontend_node = super().post_code_processing(new_build_config, current_build_config) + template = frontend_node["template"]["template"]["value"] + _ = process_prompt_template( + template=template, + name="template", + custom_fields=frontend_node["custom_fields"], + frontend_node_template=frontend_node["template"], + ) + # Now that template is updated, we need to grab any values that were set in the current_build_config + # and update the frontend_node with those values + update_template_values(frontend_template=frontend_node, raw_template=current_build_config["template"]) + return frontend_node