From 9804d58d329c791bd98e915e91e2298aaf1db744 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Mon, 28 Aug 2023 18:05:43 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(base.py):=20add=20*args=20an?= =?UTF-8?q?d=20**kwargs=20to=20build=20method=20in=20Vertex=20class=20to?= =?UTF-8?q?=20allow=20for=20additional=20arguments=20to=20be=20passed=20?= =?UTF-8?q?=F0=9F=90=9B=20fix(types.py):=20add=20*args=20and=20**kwargs=20?= =?UTF-8?q?to=20build=20method=20in=20AgentVertex,=20LLMVertex,=20WrapperV?= =?UTF-8?q?ertex,=20ChainVertex,=20and=20PromptVertex=20classes=20to=20all?= =?UTF-8?q?ow=20for=20additional=20arguments=20to=20be=20passed=20?= =?UTF-8?q?=F0=9F=90=9B=20fix(utils.py):=20change=20return=20type=20annota?= =?UTF-8?q?tion=20in=20api=5Fkey=5Fsecurity=20function=20from=20Optional[A?= =?UTF-8?q?piKey]=20to=20Optional[User]=20to=20reflect=20the=20actual=20re?= =?UTF-8?q?turn=20type?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/backend/langflow/graph/vertex/base.py | 4 ++-- src/backend/langflow/graph/vertex/types.py | 14 +++++++++----- src/backend/langflow/services/auth/utils.py | 6 +++--- 3 files changed, 14 insertions(+), 10 deletions(-) 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