From 865fc24293ee8edefe598d3a33f94bb7fb197b1c Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Thu, 1 Aug 2024 09:34:16 -0300 Subject: [PATCH] refactor: update schema.py to include Edge related types (#3117) * refactor: update code references to use _code instead of code * refactor: add backwards compatible attributes to Component class * refactor: update Component constructor to pass config params with underscore Refactored the `Component` class in `component.py` to handle inputs and outputs. Added a new method `map_outputs` to map a list of outputs to the component. Also updated the `__init__` method to properly initialize the inputs, outputs, and other attributes. This change improves the flexibility and extensibility of the `Component` class. Co-authored-by: Gabriel Luiz Freitas Almeida * refactor: change attribute to use underscore * refactor: update CustomComponent initialization parameters Refactored the `instantiate_class` function in `loading.py` to update the initialization parameters for the `CustomComponent` class. Changed the parameter names from `user_id`, `parameters`, `vertex`, and `tracing_service` to `_user_id`, `_parameters`, `_vertex`, and `_tracing_service` respectively. This change ensures consistency and improves code readability. Co-authored-by: Gabriel Luiz Freitas Almeida * refactor: update BaseComponent to accept UUID for _user_id Updated the `BaseComponent` class in `base_component.py` to accept a `UUID` type for the `_user_id` attribute. This change improves the type safety and ensures consistency with the usage of `_user_id` throughout the codebase. * refactor: import nanoid with type annotation The `nanoid` import in `component.py` has been updated to include a type annotation `# type: ignore`. This change ensures that the type checker ignores any errors related to the `nanoid` import. * fix(custom_component.py): convert _user_id to string before passing to functions to ensure compatibility with function signatures * feat(component.py): add method to set output types based on method return type to improve type checking and validation in custom components * refactor: extract method to get method return type in CustomComponent * refactor(utils.py): refactor code to use _user_id instead of user_id for consistency and clarity perf(utils.py): optimize code by reusing cc_instance instead of calling get_component_instance multiple times * refactor(utils.py, base.py): change parameter name 'add_name' to 'keep_name' for clarity and consistency in codebase * [autofix.ci] apply automated fixes * refactor: update schema.py to include Edge related typres The `schema.py` file in the `src/backend/base/langflow/graph/edge` directory has been updated to include the `TargetHandle` and `SourceHandle` models. These models define the structure and attributes of the target and source handles used in the edge data. This change improves the clarity and consistency of the codebase. * refactor: update _extract_return_type method in CustomComponent to accept Any type The _extract_return_type method in CustomComponent has been updated to accept the Any type as the return_type parameter. This change improves the flexibility and compatibility of the method, allowing it to handle a wider range of return types. * refactor: update BaseComponent to use get_template_config method Refactored the `BaseComponent` class in `base_component.py` to use the `get_template_config` method instead of duplicating the code. This change improves code readability and reduces redundancy. * feat: add BaseModel class with model_config attribute A new `BaseModel` class has been added to the `base_model.py` file. This class extends the `PydanticBaseModel` and includes a `model_config` attribute of type `ConfigDict`. This change improves the codebase by providing a base model with a configuration dictionary for models. Co-authored-by: Gabriel Luiz Freitas Almeida * refactor: update langflow.graph.edge.schema.py Refactor the `langflow.graph.edge.schema.py` file to include the `TargetHandle` and `SourceHandle` models. This change improves the clarity and consistency of the codebase. Co-authored-by: Gabriel Luiz Freitas Almeida * refactor: update build_custom_component_template to use add_name instead of keep_name Refactor the `build_custom_component_template` function in `utils.py` to use the `add_name` parameter instead of the deprecated `keep_name` parameter. This change ensures consistency with the updated method signature and improves code clarity. * feat(component.py): add method to set output types based on method return type to improve type checking and validation in custom components (#3115) * feat(component.py): add method to set output types based on method return type to improve type checking and validation in custom components * refactor: extract method to get method return type in CustomComponent * refactor: update _extract_return_type method in CustomComponent to accept Any type The _extract_return_type method in CustomComponent has been updated to accept the Any type as the return_type parameter. This change improves the flexibility and compatibility of the method, allowing it to handle a wider range of return types. * refactor: add _template_config property to BaseComponent Add a new `_template_config` property to the `BaseComponent` class in `base_component.py`. This property is used to store the template configuration for the custom component. If the `_template_config` property is empty, it is populated by calling the `build_template_config` method. This change improves the efficiency of accessing the template configuration and ensures that it is only built when needed. * refactor: add type checking for Output types in add_types method Improve type checking in the `add_types` method of the `Output` class in `base.py`. Check if the `type_` already exists in the `types` list before adding it. This change ensures that duplicate types are not added to the list. * update starter projects * refactor: optimize imports in base.py Optimize imports in the `base.py` file by removing unused imports and organizing the remaining imports. This change improves code readability and reduces unnecessary clutter. * fix(base.py): fix condition to check if self.types is not None before checking if type_ is in self.types * refactor: update build_custom_component_template to use add_name instead of keep_name --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- .../base/langflow/graph/edge/schema.py | 64 ++++++++++++++++++- .../base/langflow/helpers/base_model.py | 6 ++ 2 files changed, 68 insertions(+), 2 deletions(-) create mode 100644 src/backend/base/langflow/helpers/base_model.py diff --git a/src/backend/base/langflow/graph/edge/schema.py b/src/backend/base/langflow/graph/edge/schema.py index 628073d13..d8ae9963c 100644 --- a/src/backend/base/langflow/graph/edge/schema.py +++ b/src/backend/base/langflow/graph/edge/schema.py @@ -1,5 +1,9 @@ -from typing import Optional, Any, List -from pydantic import BaseModel +from typing import Any, List, Optional + +from pydantic import Field, field_validator +from typing_extensions import TypedDict + +from langflow.helpers.base_model import BaseModel class ResultPair(BaseModel): @@ -32,3 +36,59 @@ class Payload(BaseModel): for result_pair in self.result_pairs[:-1] ] ) + + +class TargetHandle(BaseModel): + fieldName: str = Field(..., alias="fieldName", description="Field name for the target handle.") + id: str = Field(..., description="Unique identifier for the target handle.") + input_types: List[str] = Field( + default_factory=list, alias="inputTypes", description="List of input types for the target handle." + ) + type: str = Field(..., description="Type of the target handle.") + + +class SourceHandle(BaseModel): + base_classes: list[str] = Field( + default_factory=list, alias="baseClasses", description="List of base classes for the source handle." + ) + data_type: str = Field(..., alias="dataType", description="Data type for the source handle.") + id: str = Field(..., description="Unique identifier for the source handle.") + name: Optional[str] = Field(None, description="Name of the source handle.") + output_types: List[str] = Field(default_factory=list, description="List of output types for the source handle.") + + @field_validator("name", mode="before") + @classmethod + def validate_name(cls, v, _info): + if _info.data["data_type"] == "GroupNode": + # 'OpenAIModel-u4iGV_text_output' + splits = v.split("_", 1) + if len(splits) != 2: + raise ValueError(f"Invalid source handle name {v}") + v = splits[1] + return v + + +class SourceHandleDict(TypedDict, total=False): + baseClasses: list[str] + dataType: str + id: str + name: Optional[str] + output_types: List[str] + + +class TargetHandleDict(TypedDict): + fieldName: str + id: str + inputTypes: Optional[List[str]] + type: str + + +class EdgeDataDetails(TypedDict): + sourceHandle: SourceHandleDict + targetHandle: TargetHandleDict + + +class EdgeData(TypedDict, total=False): + source: str + target: str + data: EdgeDataDetails diff --git a/src/backend/base/langflow/helpers/base_model.py b/src/backend/base/langflow/helpers/base_model.py new file mode 100644 index 000000000..c81fd99d2 --- /dev/null +++ b/src/backend/base/langflow/helpers/base_model.py @@ -0,0 +1,6 @@ +from pydantic import BaseModel as PydanticBaseModel +from pydantic import ConfigDict + + +class BaseModel(PydanticBaseModel): + model_config = ConfigDict(populate_by_name=True)