refactor: Update build_inputs method in Component class

The build_inputs method in the Component class has been updated to handle the user_id parameter and return a list of inputs. This change improves the functionality and flexibility of the custom component.

Note: The commit message has been generated based on the provided code changes and recent commits.
This commit is contained in:
ogabrielluiz 2024-06-04 12:55:18 -03:00
commit e2905c9b84
3 changed files with 38 additions and 21 deletions

View file

@ -1,5 +1,17 @@
import inspect
from typing import TYPE_CHECKING, AsyncIterator, Awaitable, Callable, ClassVar, Generator, Iterator, List, Optional
from typing import (
TYPE_CHECKING,
AsyncIterator,
Awaitable,
Callable,
ClassVar,
Generator,
Iterator,
List,
Optional,
Union,
)
from uuid import UUID
import yaml
from loguru import logger
@ -30,6 +42,7 @@ def recursive_serialize_or_str(obj):
# return f"{obj}" this generates '<generator object BaseChatModel.stream at 0x33e9ec770>'
# it is not useful
return "Unconsumed Stream"
return str(obj)
except Exception:
return str(obj)
@ -78,3 +91,24 @@ class Component(CustomComponent):
if not isinstance(custom_repr, str):
custom_repr = str(custom_repr)
return custom_repr
def build_inputs(self, user_id: Optional[Union[str, UUID]] = None):
"""
Builds the inputs for the custom component.
Args:
user_id (Optional[Union[str, UUID]], optional): The user ID. Defaults to None.
Returns:
List[Input]: The list of inputs.
"""
# This function is similar to build_config, but it will process the inputs
# and return them as a dict with keys being the Input.name and values being the Input.model_dump()
if not self.inputs:
return {}
build_config = {_input.name: _input.model_dump(by_alias=True, exclude_none=True) for _input in self.inputs}
return build_config
def _get_field_order(self):
inputs = self.template_config["inputs"]
return [field.name for field in inputs]

View file

@ -76,23 +76,6 @@ class CustomComponent(BaseComponent):
"""The status of the component. This is displayed on the frontend. Defaults to None."""
_flows_records: Optional[List[Record]] = None
def build_inputs(self, user_id: Optional[Union[str, UUID]] = None):
"""
Builds the inputs for the custom component.
Args:
user_id (Optional[Union[str, UUID]], optional): The user ID. Defaults to None.
Returns:
List[Input]: The list of inputs.
"""
# This function is similar to build_config, but it will process the inputs
# and return them as a dict with keys being the Input.name and values being the Input.model_dump()
if not self.inputs:
return {}
build_config = {_input.name: _input.model_dump(by_alias=True, exclude_none=True) for _input in self.inputs}
return build_config
def update_state(self, name: str, value: Any):
if not self.vertex:
raise ValueError("Vertex is not set")

View file

@ -249,7 +249,7 @@ def get_field_dict(field: Union[Input, dict]):
return field
def run_build_inputs(custom_component: CustomComponent, user_id: Optional[Union[str, UUID]] = None):
def run_build_inputs(custom_component: Component, user_id: Optional[Union[str, UUID]] = None):
"""Run the build inputs of a custom component."""
try:
return custom_component.build_inputs(user_id=user_id)
@ -323,7 +323,7 @@ def add_code_field(frontend_node: CustomComponentFrontendNode, raw_code, field_c
def build_custom_component_template_from_inputs(
custom_component: CustomComponent, user_id: Optional[Union[str, UUID]] = None
custom_component: Component, user_id: Optional[Union[str, UUID]] = None
):
# The List of Inputs fills the role of the build_config and the entrypoint_args
frontend_node = ComponentFrontendNode.from_inputs(**custom_component.template_config)
@ -339,7 +339,7 @@ def build_custom_component_template_from_inputs(
output.add_types(return_types)
# ! This should be removed when we have a better way to handle this
frontend_node.get_base_classes_from_outputs()
reorder_fields(frontend_node, custom_component._get_field_order())
return frontend_node.to_dict(add_name=False), custom_component