Fix async issue in CustomComponent and

decrypt_api_key function
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-11-23 10:03:48 -03:00
commit 1da20fadef
2 changed files with 9 additions and 5 deletions

View file

@ -3,6 +3,7 @@ from uuid import UUID
import yaml import yaml
from fastapi import HTTPException from fastapi import HTTPException
from langflow.field_typing.constants import CUSTOM_COMPONENT_SUPPORTED_TYPES from langflow.field_typing.constants import CUSTOM_COMPONENT_SUPPORTED_TYPES
from langflow.interface.custom.component import Component from langflow.interface.custom.component import Component
from langflow.interface.custom.directory_reader import DirectoryReader from langflow.interface.custom.directory_reader import DirectoryReader
@ -212,7 +213,7 @@ class CustomComponent(Component):
except Exception as e: except Exception as e:
raise ValueError("Session is invalid") from e raise ValueError("Session is invalid") from e
def get_flow( async def get_flow(
self, self,
*, *,
flow_name: Optional[str] = None, flow_name: Optional[str] = None,
@ -232,7 +233,7 @@ class CustomComponent(Component):
if not flow: if not flow:
raise ValueError(f"Flow {flow_name or flow_id} not found") raise ValueError(f"Flow {flow_name or flow_id} not found")
return self.load_flow(flow.id, tweaks) return await self.load_flow(flow.id, tweaks)
def build(self, *args: Any, **kwargs: Any) -> Any: def build(self, *args: Any, **kwargs: Any) -> Any:
raise NotImplementedError raise NotImplementedError

View file

@ -6,12 +6,13 @@ from cryptography.fernet import Fernet
from fastapi import Depends, HTTPException, Security, status from fastapi import Depends, HTTPException, Security, status
from fastapi.security import APIKeyHeader, APIKeyQuery, OAuth2PasswordBearer from fastapi.security import APIKeyHeader, APIKeyQuery, OAuth2PasswordBearer
from jose import JWTError, jwt from jose import JWTError, jwt
from sqlmodel import Session
from langflow.services.database.models.api_key.api_key import ApiKey from langflow.services.database.models.api_key.api_key import ApiKey
from langflow.services.database.models.api_key.crud import check_key from langflow.services.database.models.api_key.crud import check_key
from langflow.services.database.models.user.crud import get_user_by_id, get_user_by_username, update_user_last_login_at from langflow.services.database.models.user.crud import get_user_by_id, get_user_by_username, update_user_last_login_at
from langflow.services.database.models.user.user import User from langflow.services.database.models.user.user import User
from langflow.services.deps import get_session, get_settings_service from langflow.services.deps import get_session, get_settings_service
from sqlmodel import Session
oauth2_login = OAuth2PasswordBearer(tokenUrl="api/v1/login", auto_error=False) oauth2_login = OAuth2PasswordBearer(tokenUrl="api/v1/login", auto_error=False)
@ -323,6 +324,8 @@ def decrypt_api_key(encrypted_api_key: str, settings_service=Depends(get_setting
fernet = get_fernet(settings_service) fernet = get_fernet(settings_service)
# Two-way decryption # Two-way decryption
if isinstance(encrypted_api_key, str): if isinstance(encrypted_api_key, str):
encrypted_api_key = encrypted_api_key.encode() encoded_bytes = encrypted_api_key.encode()
decrypted_key = fernet.decrypt(encrypted_api_key).decode() else:
encoded_bytes = encrypted_api_key
decrypted_key = fernet.decrypt(encoded_bytes).decode()
return decrypted_key return decrypted_key