From 72a16ef60732acefc2ebb13a2b5ab5cd611f714f Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Sat, 9 Dec 2023 23:23:06 -0300 Subject: [PATCH] Refactor code by removing unused to_dict() methods --- src/backend/langflow/template/field/base.py | 42 ++++--------- .../langflow/template/frontend_node/agents.py | 61 ++++++------------- .../langflow/template/frontend_node/base.py | 8 +-- .../langflow/template/frontend_node/chains.py | 3 - .../frontend_node/custom_components.py | 3 - .../formatter/field_formatters.py | 4 +- .../template/frontend_node/prompts.py | 16 +---- .../langflow/template/frontend_node/tools.py | 13 +--- .../langflow/template/template/base.py | 5 +- 9 files changed, 41 insertions(+), 114 deletions(-) diff --git a/src/backend/langflow/template/field/base.py b/src/backend/langflow/template/field/base.py index 2c24e75bc..ef88ff0f5 100644 --- a/src/backend/langflow/template/field/base.py +++ b/src/backend/langflow/template/field/base.py @@ -1,11 +1,11 @@ -from abc import ABC -from typing import Any, Optional, Union +from typing import Any, Optional -from pydantic import BaseModel, Field, field_serializer, model_serializer +from pydantic import BaseModel, ConfigDict, Field, field_serializer -class TemplateFieldCreator(BaseModel, ABC): - field_type: str = Field(default="str", alias="type") +class TemplateField(BaseModel): + model_config = ConfigDict() + field_type: str = Field(default="str", serialization_alias="type") """The type of field this is. Default is a string.""" required: bool = False @@ -14,7 +14,7 @@ class TemplateFieldCreator(BaseModel, ABC): placeholder: str = "" """A placeholder string for the field. Default is an empty string.""" - is_list: bool = Field(default=False, alias="list") + is_list: bool = Field(default=False, serialization_alias="list") """Defines if the field is a list. Default is False.""" show: bool = True @@ -23,19 +23,19 @@ class TemplateFieldCreator(BaseModel, ABC): multiline: bool = False """Defines if the field will allow the user to open a text editor. Default is False.""" - value: Any = None + value: Any = "" """The value of the field. Default is None.""" - file_types: list[str] = Field(default=[], alias="fileTypes") + file_types: list[str] = Field(default=[], serialization_alias="fileTypes") """List of file types associated with the field. Default is an empty list. (duplicate)""" - file_path: Union[str, None] = None + file_path: Optional[str] = "" """The file path of the field if it is a file. Defaults to None.""" password: bool = False """Specifies if the field is a password. Defaults to False.""" - options: list[str] = None + options: Optional[list[str]] = None """List of options for the field. Only used when is_list=True. Default is an empty list.""" name: str = "" @@ -59,31 +59,11 @@ class TemplateFieldCreator(BaseModel, ABC): refresh: Optional[bool] = None """Specifies if the field should be refreshed. Defaults to False.""" - @model_serializer(mode="wrap") - def serialize_model(self, handler): - # This will be the result of model_dump or dict() - # so we need to build a dict to return - result = handler(self) - result["value"] = self.value - return result - - - - # for key in list(result.keys()): - # if result[key] is None or result[key] == [] and key != "value": - # del result[key] - # return result - def to_dict(self): return self.model_dump(by_alias=True, exclude_none=True) - @field_serializer("file_path") def serialize_file_path(self, value): if self.field_type == "file": return value - return None - - -class TemplateField(TemplateFieldCreator): - pass + return "" diff --git a/src/backend/langflow/template/frontend_node/agents.py b/src/backend/langflow/template/frontend_node/agents.py index 852f8d09f..b2a08d944 100644 --- a/src/backend/langflow/template/frontend_node/agents.py +++ b/src/backend/langflow/template/frontend_node/agents.py @@ -28,17 +28,17 @@ class SQLAgentNode(FrontendNode): type_name="sql_agent", fields=[ TemplateField( - field_type="str", # pyright: ignore + field_type="str", # pyright: ignore required=True, placeholder="", - is_list=False, # pyright: ignore + is_list=False, # pyright: ignore show=True, multiline=False, value="", name="database_uri", ), TemplateField( - field_type="BaseLanguageModel", # pyright: ignore + field_type="BaseLanguageModel", # pyright: ignore required=True, show=True, name="llm", @@ -49,9 +49,6 @@ class SQLAgentNode(FrontendNode): description: str = """Construct an SQL agent from an LLM and tools.""" base_classes: list[str] = ["AgentExecutor"] - def to_dict(self): - return super().to_dict() - class VectorStoreRouterAgentNode(FrontendNode): name: str = "VectorStoreRouterAgent" @@ -59,14 +56,14 @@ class VectorStoreRouterAgentNode(FrontendNode): type_name="vectorstorerouter_agent", fields=[ TemplateField( - field_type="VectorStoreRouterToolkit", # pyright: ignore + field_type="VectorStoreRouterToolkit", # pyright: ignore required=True, show=True, name="vectorstoreroutertoolkit", display_name="Vector Store Router Toolkit", ), TemplateField( - field_type="BaseLanguageModel", # pyright: ignore + field_type="BaseLanguageModel", # pyright: ignore required=True, show=True, name="llm", @@ -77,9 +74,6 @@ class VectorStoreRouterAgentNode(FrontendNode): description: str = """Construct an agent from a Vector Store Router.""" base_classes: list[str] = ["AgentExecutor"] - def to_dict(self): - return super().to_dict() - class VectorStoreAgentNode(FrontendNode): name: str = "VectorStoreAgent" @@ -87,14 +81,14 @@ class VectorStoreAgentNode(FrontendNode): type_name="vectorstore_agent", fields=[ TemplateField( - field_type="VectorStoreInfo", # pyright: ignore + field_type="VectorStoreInfo", # pyright: ignore required=True, show=True, name="vectorstoreinfo", display_name="Vector Store Info", ), TemplateField( - field_type="BaseLanguageModel", # pyright: ignore + field_type="BaseLanguageModel", # pyright: ignore required=True, show=True, name="llm", @@ -105,9 +99,6 @@ class VectorStoreAgentNode(FrontendNode): description: str = """Construct an agent from a Vector Store.""" base_classes: list[str] = ["AgentExecutor"] - def to_dict(self): - return super().to_dict() - class SQLDatabaseNode(FrontendNode): name: str = "SQLDatabase" @@ -115,9 +106,9 @@ class SQLDatabaseNode(FrontendNode): type_name="sql_database", fields=[ TemplateField( - field_type="str", # pyright: ignore + field_type="str", # pyright: ignore required=True, - is_list=False, # pyright: ignore + is_list=False, # pyright: ignore show=True, multiline=False, value="", @@ -128,9 +119,6 @@ class SQLDatabaseNode(FrontendNode): description: str = """SQLAlchemy wrapper around a database.""" base_classes: list[str] = ["SQLDatabase"] - def to_dict(self): - return super().to_dict() - class CSVAgentNode(FrontendNode): name: str = "CSVAgent" @@ -138,15 +126,15 @@ class CSVAgentNode(FrontendNode): type_name="csv_agent", fields=[ TemplateField( - field_type="file", # pyright: ignore + field_type="file", # pyright: ignore required=True, show=True, name="path", value="", - file_types=[".csv"], # pyright: ignore + file_types=[".csv"], # pyright: ignore ), TemplateField( - field_type="BaseLanguageModel", # pyright: ignore + field_type="BaseLanguageModel", # pyright: ignore required=True, show=True, name="llm", @@ -157,9 +145,6 @@ class CSVAgentNode(FrontendNode): description: str = """Construct a CSV agent from a CSV and tools.""" base_classes: list[str] = ["AgentExecutor"] - def to_dict(self): - return super().to_dict() - class InitializeAgentNode(FrontendNode): name: str = "AgentInitializer" @@ -168,9 +153,9 @@ class InitializeAgentNode(FrontendNode): type_name="initialize_agent", fields=[ TemplateField( - field_type="str", # pyright: ignore + field_type="str", # pyright: ignore required=True, - is_list=True, # pyright: ignore + is_list=True, # pyright: ignore show=True, multiline=False, options=list(NON_CHAT_AGENTS.keys()), @@ -179,22 +164,22 @@ class InitializeAgentNode(FrontendNode): advanced=False, ), TemplateField( - field_type="BaseChatMemory", # pyright: ignore + field_type="BaseChatMemory", # pyright: ignore required=False, show=True, name="memory", advanced=False, ), TemplateField( - field_type="Tool", # pyright: ignore + field_type="Tool", # pyright: ignore required=True, show=True, name="tools", - is_list=True, # pyright: ignore + is_list=True, # pyright: ignore advanced=False, ), TemplateField( - field_type="BaseLanguageModel", # pyright: ignore + field_type="BaseLanguageModel", # pyright: ignore required=True, show=True, name="llm", @@ -206,9 +191,6 @@ class InitializeAgentNode(FrontendNode): description: str = """Construct a zero shot agent from an LLM and tools.""" base_classes: list[str] = ["AgentExecutor", "Callable"] - def to_dict(self): - return super().to_dict() - @staticmethod def format_field(field: TemplateField, name: Optional[str] = None) -> None: # do nothing and don't return anything @@ -221,13 +203,13 @@ class JsonAgentNode(FrontendNode): type_name="json_agent", fields=[ TemplateField( - field_type="BaseToolkit", # pyright: ignore + field_type="BaseToolkit", # pyright: ignore required=True, show=True, name="toolkit", ), TemplateField( - field_type="BaseLanguageModel", # pyright: ignore + field_type="BaseLanguageModel", # pyright: ignore required=True, show=True, name="llm", @@ -237,6 +219,3 @@ class JsonAgentNode(FrontendNode): ) description: str = """Construct a json agent from an LLM and tools.""" base_classes: list[str] = ["AgentExecutor"] - - def to_dict(self): - return super().to_dict() diff --git a/src/backend/langflow/template/frontend_node/base.py b/src/backend/langflow/template/frontend_node/base.py index f05260f68..bcf93ad30 100644 --- a/src/backend/langflow/template/frontend_node/base.py +++ b/src/backend/langflow/template/frontend_node/base.py @@ -3,8 +3,7 @@ from collections import defaultdict from typing import ClassVar, Dict, List, Optional from langflow.template.field.base import TemplateField -from langflow.template.frontend_node.constants import (CLASSES_TO_REMOVE, - FORCE_SHOW_FIELDS) +from langflow.template.frontend_node.constants import CLASSES_TO_REMOVE, FORCE_SHOW_FIELDS from langflow.template.frontend_node.formatter import field_formatters from langflow.template.template.base import Template from langflow.utils import constants @@ -76,15 +75,16 @@ class FrontendNode(BaseModel): return display_name or self.name - @model_serializer(mode="wrap") def serialize(self, handler): result = handler(self) - result["template"] = self.template.to_dict(self.format_field) + if hasattr(self, "template") and hasattr(self.template, "to_dict"): + result["template"] = self.template.to_dict(self.format_field) name = result.pop("name") return {name: result} + # For backwards compatibility def to_dict(self) -> dict: """Returns a dict representation of the frontend node.""" diff --git a/src/backend/langflow/template/frontend_node/chains.py b/src/backend/langflow/template/frontend_node/chains.py index f6e1eacd1..4168bd9e1 100644 --- a/src/backend/langflow/template/frontend_node/chains.py +++ b/src/backend/langflow/template/frontend_node/chains.py @@ -247,9 +247,6 @@ class CombineDocsChainNode(FrontendNode): description: str = """Load question answering chain.""" base_classes: list[str] = ["BaseCombineDocumentsChain", "Callable"] - def to_dict(self): - return super().to_dict() - @staticmethod def format_field(field: TemplateField, name: Optional[str] = None) -> None: # do nothing and don't return anything diff --git a/src/backend/langflow/template/frontend_node/custom_components.py b/src/backend/langflow/template/frontend_node/custom_components.py index a0a5a9956..9c9f18d5d 100644 --- a/src/backend/langflow/template/frontend_node/custom_components.py +++ b/src/backend/langflow/template/frontend_node/custom_components.py @@ -66,6 +66,3 @@ class CustomComponentFrontendNode(FrontendNode): ) description: Optional[str] = None base_classes: list[str] = [] - - def to_dict(self): - return super().to_dict() diff --git a/src/backend/langflow/template/frontend_node/formatter/field_formatters.py b/src/backend/langflow/template/frontend_node/formatter/field_formatters.py index dc57ff04f..40f2b780c 100644 --- a/src/backend/langflow/template/frontend_node/formatter/field_formatters.py +++ b/src/backend/langflow/template/frontend_node/formatter/field_formatters.py @@ -129,7 +129,7 @@ class MultilineFieldFormatter(FieldFormatter): class DefaultValueFormatter(FieldFormatter): def format(self, field: TemplateField, name: Optional[str] = None) -> None: - value = field.to_dict() + value = field.model_dump(by_alias=True, exclude_none=True) if "default" in value: field.value = value["default"] @@ -144,7 +144,7 @@ class HeadersDefaultValueFormatter(FieldFormatter): class DictCodeFileFormatter(FieldFormatter): def format(self, field: TemplateField, name: Optional[str] = None) -> None: key = field.name - value = field.to_dict() + value = field.model_dump(by_alias=True, exclude_none=True) _type = value["type"] if "dict" in _type.lower() and key == "dict_": field.field_type = "file" diff --git a/src/backend/langflow/template/frontend_node/prompts.py b/src/backend/langflow/template/frontend_node/prompts.py index dccd66301..ffeef921c 100644 --- a/src/backend/langflow/template/frontend_node/prompts.py +++ b/src/backend/langflow/template/frontend_node/prompts.py @@ -1,14 +1,9 @@ from typing import Optional from langchain.agents.mrkl import prompt - -from langflow.template.frontend_node.constants import ( - DEFAULT_PROMPT, - HUMAN_PROMPT, - SYSTEM_PROMPT, -) from langflow.template.field.base import TemplateField from langflow.template.frontend_node.base import FrontendNode +from langflow.template.frontend_node.constants import DEFAULT_PROMPT, HUMAN_PROMPT, SYSTEM_PROMPT from langflow.template.template.base import Template @@ -50,9 +45,6 @@ class PromptTemplateNode(FrontendNode): description: str base_classes: list[str] = ["BasePromptTemplate"] - def to_dict(self): - return super().to_dict() - @staticmethod def format_field(field: TemplateField, name: Optional[str] = None) -> None: FrontendNode.format_field(field, name) @@ -66,9 +58,6 @@ class BasePromptFrontendNode(FrontendNode): description: str base_classes: list[str] - def to_dict(self): - return super().to_dict() - class ZeroShotPromptNode(BasePromptFrontendNode): name: str = "ZeroShotPrompt" @@ -110,9 +99,6 @@ class ZeroShotPromptNode(BasePromptFrontendNode): description: str = "Prompt template for Zero Shot Agent." base_classes: list[str] = ["BasePromptTemplate"] - def to_dict(self): - return super().to_dict() - @staticmethod def format_field(field: TemplateField, name: Optional[str] = None) -> None: PromptFrontendNode.format_field(field, name) diff --git a/src/backend/langflow/template/frontend_node/tools.py b/src/backend/langflow/template/frontend_node/tools.py index 315b216a4..5bed90c05 100644 --- a/src/backend/langflow/template/frontend_node/tools.py +++ b/src/backend/langflow/template/frontend_node/tools.py @@ -1,9 +1,7 @@ from langflow.template.field.base import TemplateField from langflow.template.frontend_node.base import FrontendNode from langflow.template.template.base import Template -from langflow.utils.constants import ( - DEFAULT_PYTHON_FUNCTION, -) +from langflow.utils.constants import DEFAULT_PYTHON_FUNCTION class ToolNode(FrontendNode): @@ -57,9 +55,6 @@ class ToolNode(FrontendNode): description: str = "Converts a chain, agent or function into a tool." base_classes: list[str] = ["Tool", "BaseTool"] - def to_dict(self): - return super().to_dict() - class PythonFunctionToolNode(FrontendNode): name: str = "PythonFunctionTool" @@ -113,9 +108,6 @@ class PythonFunctionToolNode(FrontendNode): description: str = "Python function to be executed." base_classes: list[str] = ["BaseTool", "Tool"] - def to_dict(self): - return super().to_dict() - class PythonFunctionNode(FrontendNode): name: str = "PythonFunction" @@ -136,6 +128,3 @@ class PythonFunctionNode(FrontendNode): ) description: str = "Python function to be executed." base_classes: list[str] = ["Callable"] - - def to_dict(self): - return super().to_dict() diff --git a/src/backend/langflow/template/template/base.py b/src/backend/langflow/template/template/base.py index 6b1fc1b8e..bb5d35375 100644 --- a/src/backend/langflow/template/template/base.py +++ b/src/backend/langflow/template/template/base.py @@ -27,15 +27,14 @@ class Template(BaseModel): def serialize_model(self, handler): result = handler(self) for field in self.fields: - result[field.name] = field.to_dict() + result[field.name] = field.model_dump(by_alias=True, exclude_none=True) result["_type"] = result.pop("type_name") return result + # For backwards compatibility def to_dict(self, format_field_func=None): self.process_fields(format_field_func) self.sort_fields() - # result = {field.name: field.to_dict() for field in self.fields} - # result["_type"] = self.type_name # type: ignore return self.model_dump(by_alias=True, exclude_none=True, exclude={"fields"}) def add_field(self, field: TemplateField) -> None: