From edc96718d7ce728085d1d7a0e967d44a2deef128 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Sun, 23 Jun 2024 17:28:16 -0300 Subject: [PATCH] refactor: Add CustomComponentResponse model and update custom_component endpoint This commit adds the CustomComponentResponse model to the schemas.py file, allowing the custom_component endpoint to return a structured response. The custom_component endpoint in the endpoints.py file has been updated to include the response_model parameter, specifying the CustomComponentResponse model. This ensures that the endpoint returns the built_frontend_node along with the type of the component_instance. These changes improve the consistency and clarity of the API response. --- src/backend/base/langflow/api/v1/endpoints.py | 5 +++-- src/backend/base/langflow/api/v1/schemas.py | 5 +++++ 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/src/backend/base/langflow/api/v1/endpoints.py b/src/backend/base/langflow/api/v1/endpoints.py index b4054fd3d..965c29292 100644 --- a/src/backend/base/langflow/api/v1/endpoints.py +++ b/src/backend/base/langflow/api/v1/endpoints.py @@ -11,6 +11,7 @@ from sqlmodel import Session, select from langflow.api.v1.schemas import ( ConfigResponse, CustomComponentRequest, + CustomComponentResponse, InputValueRequest, ProcessResponse, RunResponse, @@ -458,7 +459,7 @@ def get_version(): return {"version": version, "package": package} -@router.post("/custom_component", status_code=HTTPStatus.OK) +@router.post("/custom_component", status_code=HTTPStatus.OK, response_model=CustomComponentResponse) async def custom_component( raw_code: CustomComponentRequest, user: User = Depends(get_current_active_user), @@ -468,7 +469,7 @@ async def custom_component( built_frontend_node, component_instance = build_custom_component_template(component, user_id=user.id) if raw_code.frontend_node is not None: built_frontend_node = component_instance.post_code_processing(built_frontend_node, raw_code.frontend_node) - return built_frontend_node + return CustomComponentResponse(data=built_frontend_node, type=component_instance.__class__.__name__) @router.post("/custom_component/update", status_code=HTTPStatus.OK) diff --git a/src/backend/base/langflow/api/v1/schemas.py b/src/backend/base/langflow/api/v1/schemas.py index 7124f29e6..d8cf98d6d 100644 --- a/src/backend/base/langflow/api/v1/schemas.py +++ b/src/backend/base/langflow/api/v1/schemas.py @@ -182,6 +182,11 @@ class CustomComponentRequest(BaseModel): frontend_node: Optional[dict] = None +class CustomComponentResponse(BaseModel): + data: dict + type: str + + class UpdateCustomComponentRequest(CustomComponentRequest): field: str field_value: Optional[Union[str, int, float, bool, dict, list]] = None