🐛 fix(base.py): add *args and **kwargs to build method in Vertex class to allow for additional arguments to be passed

🐛 fix(types.py): add *args and **kwargs to build method in AgentVertex, LLMVertex, WrapperVertex, ChainVertex, and PromptVertex classes to allow for additional arguments to be passed
🐛 fix(utils.py): change return type annotation in api_key_security function from Optional[ApiKey] to Optional[User] to reflect the actual return type

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
 

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 (

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