diff --git a/src/backend/langflow/api/v1/chat.py b/src/backend/langflow/api/v1/chat.py index 0f277fbb7..a81d088d6 100644 --- a/src/backend/langflow/api/v1/chat.py +++ b/src/backend/langflow/api/v1/chat.py @@ -127,6 +127,7 @@ async def stream_build( flow_id: str, chat_service: "ChatService" = Depends(get_chat_service), cache_service: "BaseCacheService" = Depends(get_cache_service), + user=Depends(get_current_active_user), ): """Stream the build process based on stored flow data.""" @@ -167,9 +168,9 @@ async def stream_build( } yield str(StreamData(event="log", data=log_dict)) if vertex.is_task: - vertex = try_running_celery_task(vertex) + vertex = try_running_celery_task(vertex, user.id) else: - vertex.build() + vertex.build(user_id=user.id) params = vertex._built_object_repr() valid = True logger.debug(f"Building node {str(vertex.vertex_type)}") @@ -233,7 +234,7 @@ async def stream_build( raise HTTPException(status_code=500, detail=str(exc)) -def try_running_celery_task(vertex): +def try_running_celery_task(vertex, user_id): # Try running the task in celery # and set the task_id to the local vertex # if it fails, run the task locally @@ -245,5 +246,5 @@ def try_running_celery_task(vertex): except Exception as exc: logger.debug(f"Error running task in celery: {exc}") vertex.task_id = None - vertex.build() + vertex.build(user_id=user_id) return vertex diff --git a/src/backend/langflow/api/v1/endpoints.py b/src/backend/langflow/api/v1/endpoints.py index 9b1213c2c..fee1fc367 100644 --- a/src/backend/langflow/api/v1/endpoints.py +++ b/src/backend/langflow/api/v1/endpoints.py @@ -227,6 +227,7 @@ def get_version(): @router.post("/custom_component", status_code=HTTPStatus.OK) async def custom_component( raw_code: CustomComponentCode, + user: User = Depends(get_current_active_user), ): from langflow.interface.types import ( build_langchain_template_custom_component, @@ -235,4 +236,4 @@ async def custom_component( extractor = CustomComponent(code=raw_code.code) extractor.is_check_valid() - return build_langchain_template_custom_component(extractor) + return build_langchain_template_custom_component(extractor, user_id=user.id) diff --git a/src/backend/langflow/interface/custom/component.py b/src/backend/langflow/interface/custom/component.py index cd59ff8ff..2e88b1f26 100644 --- a/src/backend/langflow/interface/custom/component.py +++ b/src/backend/langflow/interface/custom/component.py @@ -1,6 +1,5 @@ import ast from typing import Any, ClassVar, Optional -from pydantic import BaseModel from fastapi import HTTPException from langflow.utils import validate @@ -15,7 +14,7 @@ class ComponentFunctionEntrypointNameNullError(HTTPException): pass -class Component(BaseModel): +class Component: ERROR_CODE_NULL: ClassVar[str] = "Python code must be provided." ERROR_FUNCTION_ENTRYPOINT_NAME_NULL: ClassVar[ str @@ -26,7 +25,8 @@ class Component(BaseModel): field_config: dict = {} def __init__(self, **data): - super().__init__(**data) + for key, value in data.items(): + setattr(self, key, value) def get_code_tree(self, code: str): parser = CodeParser(code) diff --git a/src/backend/langflow/interface/custom/custom_component.py b/src/backend/langflow/interface/custom/custom_component.py index 825d18e1b..440c93bb0 100644 --- a/src/backend/langflow/interface/custom/custom_component.py +++ b/src/backend/langflow/interface/custom/custom_component.py @@ -14,7 +14,7 @@ from langflow.services.database.models.flow import Flow import yaml -class CustomComponent(Component, extra="allow"): +class CustomComponent(Component): display_name: Optional[str] = "Custom Component" description: Optional[str] = "Custom Component" code: Optional[str] = None @@ -203,7 +203,7 @@ class CustomComponent(Component, extra="allow"): raise ValueError(f"Flow {flow_id} not found") if tweaks: graph_data = process_tweaks(graph_data=graph_data, tweaks=tweaks) - return build_sorted_vertices(graph_data) + return build_sorted_vertices(graph_data, self.user_id) def list_flows(self, *, get_session: Optional[Callable] = None) -> List[Flow]: if not self.user_id: diff --git a/src/backend/langflow/interface/run.py b/src/backend/langflow/interface/run.py index 63391204a..f57418a9a 100644 --- a/src/backend/langflow/interface/run.py +++ b/src/backend/langflow/interface/run.py @@ -3,7 +3,7 @@ from langflow.graph import Graph from loguru import logger -def build_sorted_vertices(data_graph) -> Tuple[Graph, Dict]: +def build_sorted_vertices(data_graph, user_id) -> Tuple[Graph, Dict]: """ Build langchain object from data_graph. """ @@ -13,7 +13,7 @@ def build_sorted_vertices(data_graph) -> Tuple[Graph, Dict]: sorted_vertices = graph.topological_sort() artifacts = {} for vertex in sorted_vertices: - vertex.build() + vertex.build(user_id=user_id) if vertex.artifacts: artifacts.update(vertex.artifacts) return graph, artifacts diff --git a/src/backend/langflow/interface/types.py b/src/backend/langflow/interface/types.py index 0046c84fb..ec603818a 100644 --- a/src/backend/langflow/interface/types.py +++ b/src/backend/langflow/interface/types.py @@ -208,7 +208,7 @@ def update_attributes(frontend_node, template_config): frontend_node[attribute] = template_config[attribute] -def build_field_config(custom_component: CustomComponent): +def build_field_config(custom_component: CustomComponent, user_id: str = None): """Build the field configuration for a custom component""" try: @@ -218,7 +218,7 @@ def build_field_config(custom_component: CustomComponent): return {} try: - return custom_class().build_config() + return custom_class(user_id=user_id).build_config() except Exception as exc: logger.error(f"Error while building field config: {str(exc)}") return {} @@ -306,7 +306,9 @@ def add_output_types(frontend_node, return_types: List[str]): frontend_node.get("output_types").append(return_type) -def build_langchain_template_custom_component(custom_component: CustomComponent): +def build_langchain_template_custom_component( + custom_component: CustomComponent, user_id: str = None +): """Build a custom component template for the langchain""" try: logger.debug("Building custom component template") @@ -319,7 +321,7 @@ def build_langchain_template_custom_component(custom_component: CustomComponent) update_attributes(frontend_node, template_config) logger.debug("Updated attributes") - field_config = build_field_config(custom_component) + field_config = build_field_config(custom_component, user_id=user_id) logger.debug("Built field config") entrypoint_args = custom_component.get_function_entrypoint_args