diff --git a/src/backend/langflow/api/utils.py b/src/backend/langflow/api/utils.py index d88d2525f..b7b893dda 100644 --- a/src/backend/langflow/api/utils.py +++ b/src/backend/langflow/api/utils.py @@ -1,11 +1,12 @@ +import warnings from pathlib import Path from typing import TYPE_CHECKING, List + from fastapi import HTTPException +from platformdirs import user_cache_dir + from langflow.services.store.schema import StoreComponentCreate from langflow.services.store.utils import get_lf_version_from_pypi -import warnings - -from platformdirs import user_cache_dir if TYPE_CHECKING: from langflow.services.database.models.flow.model import Flow @@ -62,7 +63,7 @@ def build_input_keys_response(langchain_object, artifacts): return input_keys_response -def update_frontend_node_with_template_values(frontend_node, raw_template_data): +def update_frontend_node_with_template_values(frontend_node, raw_frontend_node): """ Updates the given frontend node with values from the raw template data. @@ -70,19 +71,28 @@ def update_frontend_node_with_template_values(frontend_node, raw_template_data): :param raw_template_data: A dict representing raw template data. :return: Updated frontend node. """ - if not is_valid_data(frontend_node, raw_template_data): + if not is_valid_data(frontend_node, raw_frontend_node): return frontend_node - update_template_values(frontend_node["template"], raw_template_data.template) + # Check if the display_name is different than "CustomComponent" + # if so, update the display_name in the frontend_node + if raw_frontend_node["display_name"] != "CustomComponent": + frontend_node["display_name"] = raw_frontend_node["display_name"] + + update_template_values(frontend_node["template"], raw_frontend_node["template"]) return frontend_node -def is_valid_data(frontend_node, raw_template_data): +def raw_frontend_data_is_valid(raw_frontend_data): + """Check if the raw frontend data is valid for processing.""" + return "template" in raw_frontend_data and "display_name" in raw_frontend_data + + +def is_valid_data(frontend_node, raw_frontend_data): """Check if the data is valid for processing.""" - return ( - frontend_node and "template" in frontend_node and raw_template_data and hasattr(raw_template_data, "template") - ) + + return frontend_node and "template" in frontend_node and raw_frontend_data_is_valid(raw_frontend_data) def update_template_values(frontend_template, raw_template): diff --git a/src/backend/langflow/api/v1/endpoints.py b/src/backend/langflow/api/v1/endpoints.py index d4bb083ef..6ebb04182 100644 --- a/src/backend/langflow/api/v1/endpoints.py +++ b/src/backend/langflow/api/v1/endpoints.py @@ -3,6 +3,9 @@ from typing import Annotated, Optional, Union import sqlalchemy as sa from fastapi import APIRouter, Body, Depends, HTTPException, UploadFile, status +from loguru import logger +from sqlmodel import select + from langflow.api.utils import update_frontend_node_with_template_values from langflow.api.v1.schemas import ( CustomComponentCode, @@ -20,8 +23,6 @@ from langflow.services.cache.utils import save_uploaded_file from langflow.services.database.models.flow import Flow from langflow.services.database.models.user.model import User from langflow.services.deps import get_session, get_session_service, get_settings_service, get_task_service -from loguru import logger -from sqlmodel import select try: from langflow.worker import process_graph_cached_task @@ -31,9 +32,10 @@ except ImportError: raise NotImplementedError("Celery is not installed") -from langflow.services.task.service import TaskService from sqlmodel import Session +from langflow.services.task.service import TaskService + # build router router = APIRouter(tags=["Base"]) @@ -218,7 +220,7 @@ async def custom_component( built_frontend_node = build_custom_component_template(component, user_id=user.id) - built_frontend_node = update_frontend_node_with_template_values(built_frontend_node, raw_code) + built_frontend_node = update_frontend_node_with_template_values(built_frontend_node, raw_code.frontend_node) return built_frontend_node diff --git a/src/backend/langflow/api/v1/schemas.py b/src/backend/langflow/api/v1/schemas.py index 2d3da1d21..8c8c4efd7 100644 --- a/src/backend/langflow/api/v1/schemas.py +++ b/src/backend/langflow/api/v1/schemas.py @@ -3,11 +3,12 @@ from pathlib import Path from typing import Any, Dict, List, Optional, Union from uuid import UUID +from pydantic import BaseModel, Field, field_validator + from langflow.services.database.models.api_key.model import ApiKeyRead from langflow.services.database.models.base import orjson_dumps from langflow.services.database.models.flow import FlowCreate, FlowRead from langflow.services.database.models.user import UserRead -from pydantic import BaseModel, Field, field_validator class BuildStatus(Enum): @@ -157,7 +158,7 @@ class StreamData(BaseModel): class CustomComponentCode(BaseModel): code: str field: Optional[str] = None - template: Optional[dict] = None + frontend_node: Optional[dict] = None class CustomComponentResponseError(BaseModel):