diff --git a/src/backend/langflow/graph/vertex/base.py b/src/backend/langflow/graph/vertex/base.py index 09c813a2e..d5c4beed9 100644 --- a/src/backend/langflow/graph/vertex/base.py +++ b/src/backend/langflow/graph/vertex/base.py @@ -258,9 +258,9 @@ class Vertex: raise ValueError(message) - def build(self, force: bool = False, user_id=None) -> Any: + def build(self, force: bool = False, user_id=None, *args, **kwargs) -> Any: if not self._built or force: - self._build(user_id) + self._build(user_id, *args, **kwargs) return self._built_object diff --git a/src/backend/langflow/graph/vertex/types.py b/src/backend/langflow/graph/vertex/types.py index a70a06e88..d5b7d2d58 100644 --- a/src/backend/langflow/graph/vertex/types.py +++ b/src/backend/langflow/graph/vertex/types.py @@ -21,7 +21,7 @@ class AgentVertex(Vertex): elif isinstance(source_node, ChainVertex): self.chains.append(source_node) - def build(self, force: bool = False, user_id=None) -> Any: + def build(self, force: bool = False, user_id=None, *args, **kwargs) -> Any: if not self._built or force: self._set_tools_and_chains() # First, build the tools @@ -49,7 +49,7 @@ class LLMVertex(Vertex): def __init__(self, data: Dict): super().__init__(data, base_type="llms") - def build(self, force: bool = False, user_id=None) -> Any: + def build(self, force: bool = False, user_id=None, *args, **kwargs) -> Any: # LLM is different because some models might take up too much memory # or time to load. So we only load them when we need them.ß if self.vertex_type == self.built_node_type: @@ -77,7 +77,7 @@ class WrapperVertex(Vertex): def __init__(self, data: Dict): super().__init__(data, base_type="wrappers") - def build(self, force: bool = False, user_id=None) -> Any: + def build(self, force: bool = False, user_id=None, *args, **kwargs) -> Any: if not self._built or force: if "headers" in self.params: self.params["headers"] = ast.literal_eval(self.params["headers"]) @@ -148,14 +148,16 @@ class ChainVertex(Vertex): def build( self, force: bool = False, - tools: Optional[List[Union[ToolkitVertex, ToolVertex]]] = None, user_id=None, + *args, + **kwargs, ) -> Any: if not self._built or force: # Check if the chain requires a PromptVertex for key, value in self.params.items(): if isinstance(value, PromptVertex): # Build the PromptVertex, passing the tools if available + tools = kwargs.get("tools", None) self.params[key] = value.build(tools=tools, force=force) self._build(user_id=user_id) @@ -170,8 +172,10 @@ class PromptVertex(Vertex): def build( self, force: bool = False, - tools: Optional[List[Union[ToolkitVertex, ToolVertex]]] = None, user_id=None, + tools: Optional[List[Union[ToolkitVertex, ToolVertex]]] = None, + *args, + **kwargs, ) -> Any: if not self._built or force: if ( diff --git a/src/backend/langflow/services/auth/utils.py b/src/backend/langflow/services/auth/utils.py index 333ba226b..360b8fc8d 100644 --- a/src/backend/langflow/services/auth/utils.py +++ b/src/backend/langflow/services/auth/utils.py @@ -2,7 +2,7 @@ from datetime import datetime, timedelta, timezone from fastapi import Depends, HTTPException, Security, status from fastapi.security import APIKeyHeader, APIKeyQuery, OAuth2PasswordBearer from jose import JWTError, jwt -from typing import Annotated, Coroutine, Optional +from typing import Annotated, Coroutine, Optional, Union from uuid import UUID from langflow.services.database.models.api_key.api_key import ApiKey from langflow.services.database.models.api_key.crud import check_key @@ -32,9 +32,9 @@ async def api_key_security( query_param: str = Security(api_key_query), header_param: str = Security(api_key_header), db: Session = Depends(get_session), -) -> Optional[ApiKey]: +) -> Optional[User]: settings_manager = get_settings_manager() - result = None + result: Optional[Union[ApiKey, User]] = None if settings_manager.auth_settings.AUTO_LOGIN: # Get the first user settings_manager.auth_settings.FIRST_SUPERUSER