From 4583d30f4c29c8de83556a197e0221da79677013 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 1 Dec 2023 13:12:19 -0300 Subject: [PATCH 01/17] =?UTF-8?q?=E2=9C=A8=20feat(code=5Fparser.py):=20add?= =?UTF-8?q?=20support=20for=20parsing=20return=20statement=20in=20function?= =?UTF-8?q?=20or=20method=20nodes=20=F0=9F=90=9B=20fix(code=5Fparser.py):?= =?UTF-8?q?=20fix=20import=20formatting=20to=20improve=20readability=20?= =?UTF-8?q?=E2=9C=A8=20feat(code=5Fparser.py):=20add=20support=20for=20par?= =?UTF-8?q?sing=20function=20arguments=20and=20body=20in=20function=20or?= =?UTF-8?q?=20method=20nodes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../langflow/interface/custom/code_parser.py | 25 ++++++++++++++++--- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/src/backend/langflow/interface/custom/code_parser.py b/src/backend/langflow/interface/custom/code_parser.py index cb42b8142..c24653482 100644 --- a/src/backend/langflow/interface/custom/code_parser.py +++ b/src/backend/langflow/interface/custom/code_parser.py @@ -6,7 +6,8 @@ from typing import Any, Dict, List, Type, Union from cachetools import TTLCache, cachedmethod, keys from fastapi import HTTPException -from langflow.interface.custom.schema import CallableCodeDetails, ClassCodeDetails +from langflow.interface.custom.schema import (CallableCodeDetails, + ClassCodeDetails) class CodeSyntaxError(HTTPException): @@ -56,6 +57,9 @@ class CodeParser: ast.Assign: self.parse_global_vars, } + + + def __get_tree(self): """ Parses the provided code to validate its syntax. @@ -79,6 +83,7 @@ class CodeParser: if handler := self.handlers.get(type(node)): # type: ignore handler(node) # type: ignore + def parse_imports(self, node: Union[ast.Import, ast.ImportFrom]) -> None: """ Extracts "imports" from the code, including aliases. @@ -149,12 +154,16 @@ class CodeParser: # Handle cases where the type is not found in the constructed environment pass + func = CallableCodeDetails( - name=node.name, doc=ast.get_docstring(node), args=[], body=[], return_type=return_type or get_data_type() + name=node.name, + doc=ast.get_docstring(node), + args= self.parse_function_args(node), + body= self.parse_function_body(node), + return_type=return_type or get_data_type(), + has_return=self.parse_return_statement(node), ) - func.args = self.parse_function_args(node) - func.body = self.parse_function_body(node) return func.model_dump() @@ -230,6 +239,14 @@ class CodeParser: """ return [ast.unparse(line) for line in node.body] + def parse_return_statement(self, node: ast.FunctionDef) -> bool: + """ + Parses the return statement of a function or method node. + """ + + return any(isinstance(n, ast.Return) for n in node.body) + + def parse_assign(self, stmt): """ Parses an Assign statement and returns a dictionary From 7a2ba97b4c5089d910fe8caccffd0f8f4d9de1bc Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 1 Dec 2023 13:12:35 -0300 Subject: [PATCH 02/17] Add has_return attribute to CallableCodeDetails class --- src/backend/langflow/interface/custom/schema.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/backend/langflow/interface/custom/schema.py b/src/backend/langflow/interface/custom/schema.py index b055dd386..7c5975150 100644 --- a/src/backend/langflow/interface/custom/schema.py +++ b/src/backend/langflow/interface/custom/schema.py @@ -26,3 +26,4 @@ class CallableCodeDetails(BaseModel): args: list body: list return_type: Optional[Any] = None + has_return: bool = False From b1224ddb63c9a81a0340c4eb74438b8ecd10a9e5 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 1 Dec 2023 13:12:48 -0300 Subject: [PATCH 03/17] =?UTF-8?q?=E2=9C=A8=20feat(custom=5Fcomponent.py):?= =?UTF-8?q?=20add=20=5Ftree=20attribute=20to=20store=20code=20tree=20for?= =?UTF-8?q?=20better=20performance=20=F0=9F=94=A7=20refactor(custom=5Fcomp?= =?UTF-8?q?onent.py):=20refactor=20get=5Fcode=5Ftree=20method=20to=20use?= =?UTF-8?q?=20the=20=5Ftree=20attribute=20=F0=9F=94=A7=20refactor(custom?= =?UTF-8?q?=5Fcomponent.py):=20refactor=20get=5Fbuild=5Fmethod=20method=20?= =?UTF-8?q?to=20use=20the=20tree=20attribute=20=F0=9F=94=A7=20refactor(cus?= =?UTF-8?q?tom=5Fcomponent.py):=20refactor=20get=5Fmain=5Fclass=5Fname=20m?= =?UTF-8?q?ethod=20to=20use=20the=20tree=20attribute=20=F0=9F=94=A7=20refa?= =?UTF-8?q?ctor(custom=5Fcomponent.py):=20refactor=20build=5Ftemplate=5Fco?= =?UTF-8?q?nfig=20method=20to=20use=20the=20tree=20attribute=20?= =?UTF-8?q?=F0=9F=94=A7=20refactor(custom=5Fcomponent.py):=20refactor=20lo?= =?UTF-8?q?ad=5Fflow=20method=20to=20use=20import=20statement=20on=20separ?= =?UTF-8?q?ate=20lines=20for=20better=20readability?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../interface/custom/custom_component.py | 36 +++++++++++-------- 1 file changed, 22 insertions(+), 14 deletions(-) diff --git a/src/backend/langflow/interface/custom/custom_component.py b/src/backend/langflow/interface/custom/custom_component.py index 1ea2a1c78..14a0246b7 100644 --- a/src/backend/langflow/interface/custom/custom_component.py +++ b/src/backend/langflow/interface/custom/custom_component.py @@ -10,8 +10,7 @@ from langflow.interface.custom.component import Component from langflow.interface.custom.directory_reader import DirectoryReader from langflow.interface.custom.utils import ( extract_inner_type_from_generic_alias, - extract_union_types_from_generic_alias, -) + extract_union_types_from_generic_alias) from langflow.services.database.models.flow import Flow from langflow.services.database.utils import session_getter from langflow.services.deps import get_credential_service, get_db_service @@ -29,6 +28,7 @@ class CustomComponent(Component): repr_value: Optional[Any] = "" user_id: Optional[Union[UUID, str]] = None status: Optional[Any] = None + _tree: Optional[dict] = None def __init__(self, **data): self.cache = TTLCache(maxsize=1024, ttl=60) @@ -78,8 +78,11 @@ class CustomComponent(Component): def validate(self) -> bool: return self._class_template_validation(self.code) if self.code else False - def get_code_tree(self, code: str): - return super().get_code_tree(code) + + + @property + def tree(self): + return self.get_code_tree(self.code) @property def get_function_entrypoint_args(self) -> list: @@ -108,9 +111,10 @@ class CustomComponent(Component): def get_build_method(self): if not self.code: return [] - tree = self.get_code_tree(self.code) - component_classes = [cls for cls in tree["classes"] if self.code_class_base_inheritance in cls["bases"]] + + + component_classes = [cls for cls in self.tree["classes"] if self.code_class_base_inheritance in cls["bases"]] if not component_classes: return [] @@ -123,16 +127,19 @@ class CustomComponent(Component): if not build_methods: return [] + return build_methods[0] @property def get_function_entrypoint_return_type(self) -> List[Any]: build_method = self.get_build_method() if not build_method: - return build_method - return_type = build_method["return_type"] - if not return_type: return [] + elif not build_method["has_return"]: + return [] + + return_type = build_method["return_type"] + # If list or List is in the return type, then we remove it and return the inner type if hasattr(return_type, "__origin__") and return_type.__origin__ in [list, List]: return_type = extract_inner_type_from_generic_alias(return_type) @@ -152,13 +159,13 @@ class CustomComponent(Component): def get_main_class_name(self): if not self.code: return "" - tree = self.get_code_tree(self.code) + base_name = self.code_class_base_inheritance method_name = self.function_entrypoint_name classes = [] - for item in tree.get("classes", []): + for item in self.tree.get("classes", []): if base_name in item["bases"]: method_names = [method["name"] for method in item["methods"]] if method_name in method_names: @@ -171,11 +178,11 @@ class CustomComponent(Component): def build_template_config(self): if not self.code: return {} - tree = self.get_code_tree(self.code) + attributes = [ main_class["attributes"] - for main_class in tree.get("classes", []) + for main_class in self.tree.get("classes", []) if main_class["name"] == self.get_main_class_name ] # Get just the first item @@ -219,7 +226,8 @@ class CustomComponent(Component): return validate.create_function(self.code, self.function_entrypoint_name) async def load_flow(self, flow_id: str, tweaks: Optional[dict] = None) -> Any: - from langflow.processing.process import build_sorted_vertices, process_tweaks + from langflow.processing.process import (build_sorted_vertices, + process_tweaks) db_service = get_db_service() with session_getter(db_service) as session: From ead79e802daee25dff52ac6d3fd58a6690a60104 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 1 Dec 2023 13:14:17 -0300 Subject: [PATCH 04/17] Fix import statements and add missing Prompt class --- src/backend/langflow/field_typing/__init__.py | 31 +++++-------------- .../langflow/field_typing/constants.py | 4 +++ 2 files changed, 11 insertions(+), 24 deletions(-) diff --git a/src/backend/langflow/field_typing/__init__.py b/src/backend/langflow/field_typing/__init__.py index fb21a238e..c6680662f 100644 --- a/src/backend/langflow/field_typing/__init__.py +++ b/src/backend/langflow/field_typing/__init__.py @@ -1,27 +1,9 @@ -from .constants import ( - AgentExecutor, - BaseChatMemory, - BaseLanguageModel, - BaseLLM, - BaseLoader, - BaseMemory, - BaseOutputParser, - BasePromptTemplate, - BaseRetriever, - Callable, - Chain, - ChatPromptTemplate, - Data, - Document, - Embeddings, - NestedDict, - Object, - PromptTemplate, - TextSplitter, - Tool, - VectorStore, -) - +from .constants import (AgentExecutor, BaseChatMemory, BaseLanguageModel, + BaseLLM, BaseLoader, BaseMemory, BaseOutputParser, + BasePromptTemplate, BaseRetriever, Callable, Chain, + ChatPromptTemplate, Data, Document, Embeddings, + NestedDict, Object, Prompt, PromptTemplate, + TextSplitter, Tool, VectorStore) __all__ = [ "NestedDict", @@ -45,4 +27,5 @@ __all__ = [ "Callable", "BasePromptTemplate", "ChatPromptTemplate", + "Prompt" ] diff --git a/src/backend/langflow/field_typing/constants.py b/src/backend/langflow/field_typing/constants.py index 76ef3aa5b..028ab6aff 100644 --- a/src/backend/langflow/field_typing/constants.py +++ b/src/backend/langflow/field_typing/constants.py @@ -25,6 +25,8 @@ class Object: class Data: pass +class Prompt + LANGCHAIN_BASE_TYPES = { "Chain": Chain, @@ -44,6 +46,7 @@ LANGCHAIN_BASE_TYPES = { "BaseOutputParser": BaseOutputParser, "BaseMemory": BaseMemory, "BaseChatMemory": BaseChatMemory, + } # Langchain base types plus Python base types CUSTOM_COMPONENT_SUPPORTED_TYPES = { @@ -58,4 +61,5 @@ CUSTOM_COMPONENT_SUPPORTED_TYPES = { "Data": Data, "Object": Object, "Callable": Callable, + "Prompt": Prompt, } From a62d0bbb6b18a8815a86d65415b3caa32935ebb6 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 1 Dec 2023 14:35:22 -0300 Subject: [PATCH 05/17] Refactor import statements in constants.py --- src/backend/langflow/field_typing/constants.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/backend/langflow/field_typing/constants.py b/src/backend/langflow/field_typing/constants.py index 028ab6aff..d5f2c4321 100644 --- a/src/backend/langflow/field_typing/constants.py +++ b/src/backend/langflow/field_typing/constants.py @@ -5,7 +5,8 @@ from langchain.chains.base import Chain from langchain.document_loaders.base import BaseLoader from langchain.llms.base import BaseLLM from langchain.memory.chat_memory import BaseChatMemory -from langchain.prompts import BasePromptTemplate, ChatPromptTemplate, PromptTemplate +from langchain.prompts import (BasePromptTemplate, ChatPromptTemplate, + PromptTemplate) from langchain.schema import BaseOutputParser, BaseRetriever, Document from langchain.schema.embeddings import Embeddings from langchain.schema.language_model import BaseLanguageModel @@ -25,7 +26,8 @@ class Object: class Data: pass -class Prompt +class Prompt: + pass LANGCHAIN_BASE_TYPES = { From 0bff0407b8708c82eae2b938d0a57df71591c323 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 1 Dec 2023 15:10:16 -0300 Subject: [PATCH 06/17] Remove unused import and commented code --- src/backend/langflow/interface/custom/custom_component.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/src/backend/langflow/interface/custom/custom_component.py b/src/backend/langflow/interface/custom/custom_component.py index 14a0246b7..12536a20d 100644 --- a/src/backend/langflow/interface/custom/custom_component.py +++ b/src/backend/langflow/interface/custom/custom_component.py @@ -5,7 +5,6 @@ from uuid import UUID import yaml from cachetools import TTLCache, cachedmethod from fastapi import HTTPException -from langflow.field_typing.constants import CUSTOM_COMPONENT_SUPPORTED_TYPES from langflow.interface.custom.component import Component from langflow.interface.custom.directory_reader import DirectoryReader from langflow.interface.custom.utils import ( @@ -34,9 +33,7 @@ class CustomComponent(Component): self.cache = TTLCache(maxsize=1024, ttl=60) super().__init__(**data) - @property - def return_type_valid_list(self): - return list(CUSTOM_COMPONENT_SUPPORTED_TYPES.keys()) + def custom_repr(self): if self.repr_value == "": @@ -148,11 +145,10 @@ class CustomComponent(Component): if not hasattr(return_type, "__origin__") or return_type.__origin__ != Union: if isinstance(return_type, list): return return_type - return [return_type] # if return_type in self.return_type_valid_list else [] + return [return_type] # If the return type is a Union, then we need to parse itx return_type = extract_union_types_from_generic_alias(return_type) - # return [item for item in return_type if item in self.return_type_valid_list] return return_type @property From 4d6973ad19d7b4b240ef55d6e7e084a0a23508cf Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 1 Dec 2023 15:10:24 -0300 Subject: [PATCH 07/17] Remove redundant base types from CUSTOM_COMPONENT_SUPPORTED_TYPES --- src/backend/langflow/field_typing/constants.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/src/backend/langflow/field_typing/constants.py b/src/backend/langflow/field_typing/constants.py index d5f2c4321..fdfdf44a0 100644 --- a/src/backend/langflow/field_typing/constants.py +++ b/src/backend/langflow/field_typing/constants.py @@ -53,12 +53,6 @@ LANGCHAIN_BASE_TYPES = { # Langchain base types plus Python base types CUSTOM_COMPONENT_SUPPORTED_TYPES = { **LANGCHAIN_BASE_TYPES, - "str": str, - "int": int, - "float": float, - "bool": bool, - "list": list, - "dict": dict, "NestedDict": NestedDict, "Data": Data, "Object": Object, From 13135d3948f5785356bb76be996d80c405951a7d Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 1 Dec 2023 15:10:34 -0300 Subject: [PATCH 08/17] Update langflow imports in validate.py --- src/backend/langflow/utils/validate.py | 26 +++----------------------- 1 file changed, 3 insertions(+), 23 deletions(-) diff --git a/src/backend/langflow/utils/validate.py b/src/backend/langflow/utils/validate.py index b9236fe2d..855dd23e6 100644 --- a/src/backend/langflow/utils/validate.py +++ b/src/backend/langflow/utils/validate.py @@ -4,6 +4,8 @@ import importlib from types import FunctionType from typing import Dict +from langflow.field_typing.constants import CUSTOM_COMPONENT_SUPPORTED_TYPES + def add_type_ignores(): if not hasattr(ast, "TypeIgnore"): @@ -266,29 +268,7 @@ def get_default_imports(code_string): # Add more imports from the typing module as needed } - langflow_imports = [ - "AgentExecutor", - "BaseChatMemory", - "BaseLanguageModel", - "BaseLLM", - "BaseLoader", - "BaseMemory", - "BaseOutputParser", - "BasePromptTemplate", - "BaseRetriever", - "Callable", - "Chain", - "ChatPromptTemplate", - "Data", - "Document", - "Embeddings", - "NestedDict", - "Object", - "PromptTemplate", - "TextSplitter", - "Tool", - "VectorStore", - ] + langflow_imports = list(CUSTOM_COMPONENT_SUPPORTED_TYPES.keys()) necessary_imports = find_names_in_code(code_string, langflow_imports) langflow_module = importlib.import_module("langflow.field_typing") default_imports.update({name: getattr(langflow_module, name) for name in necessary_imports}) From 9a5b57a65a3afbd951c7eba0be9fae32d71d727a Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 1 Dec 2023 15:12:19 -0300 Subject: [PATCH 09/17] =?UTF-8?q?=F0=9F=94=A5=20refactor(types.py):=20remo?= =?UTF-8?q?ve=20unused=20import=20and=20variable=20'custom=5Fcomponent=5Fc?= =?UTF-8?q?reator'=20to=20improve=20code=20cleanliness=20and=20maintainabi?= =?UTF-8?q?lity?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/backend/langflow/interface/types.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/backend/langflow/interface/types.py b/src/backend/langflow/interface/types.py index 1d9c46f46..a2cda3450 100644 --- a/src/backend/langflow/interface/types.py +++ b/src/backend/langflow/interface/types.py @@ -10,7 +10,6 @@ from cachetools import LRUCache, cached from fastapi import HTTPException from langflow.interface.agents.base import agent_creator from langflow.interface.chains.base import chain_creator -from langflow.interface.custom.base import custom_component_creator from langflow.interface.custom.custom_component import CustomComponent from langflow.interface.custom.directory_reader import DirectoryReader from langflow.interface.custom.utils import extract_inner_type @@ -30,7 +29,8 @@ from langflow.interface.vector_store.base import vectorstore_creator from langflow.interface.wrappers.base import wrapper_creator from langflow.template.field.base import TemplateField from langflow.template.frontend_node.constants import CLASSES_TO_REMOVE -from langflow.template.frontend_node.custom_components import CustomComponentFrontendNode +from langflow.template.frontend_node.custom_components import \ + CustomComponentFrontendNode from langflow.utils.util import get_base_classes from loguru import logger @@ -69,7 +69,6 @@ def build_langchain_types_dict(): # sourcery skip: dict-assign-update-to-union utility_creator, output_parser_creator, retriever_creator, - custom_component_creator, ] all_types = {} From 1bdbbd12c025d021f2f41e5e7e2d478f0a2b4a58 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 1 Dec 2023 15:46:16 -0300 Subject: [PATCH 10/17] Fix store_api_key column type --- .../alembic/versions/7843803a87b5_store_updates.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/backend/langflow/alembic/versions/7843803a87b5_store_updates.py b/src/backend/langflow/alembic/versions/7843803a87b5_store_updates.py index ca82479bc..3d5f0c56b 100644 --- a/src/backend/langflow/alembic/versions/7843803a87b5_store_updates.py +++ b/src/backend/langflow/alembic/versions/7843803a87b5_store_updates.py @@ -7,10 +7,10 @@ Create Date: 2023-10-18 23:08:57.744906 """ from typing import Sequence, Union -from alembic import op import sqlalchemy as sa import sqlmodel - +from alembic import op +from loguru import logger # revision identifiers, used by Alembic. revision: str = "7843803a87b5" @@ -28,11 +28,11 @@ def upgrade() -> None: with op.batch_alter_table("user", schema=None) as batch_op: batch_op.add_column( sa.Column( - "store_api_key", sqlmodel.sql.sqltypes.AutoString(), nullable=True + "store_api_key", sqlmodel.AutoString(), nullable=True ) ) - except Exception: - pass + except Exception as e: + logger.exception(e) # ### end Alembic commands ### From 6aa1698383c5f4d274ae5e5af33886bffc9582b3 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 1 Dec 2023 16:27:15 -0300 Subject: [PATCH 11/17] Fix ValueError in Vertex class --- src/backend/langflow/graph/vertex/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/backend/langflow/graph/vertex/base.py b/src/backend/langflow/graph/vertex/base.py index 1079d29a6..b89e36ca4 100644 --- a/src/backend/langflow/graph/vertex/base.py +++ b/src/backend/langflow/graph/vertex/base.py @@ -345,7 +345,7 @@ class Vertex: if self.base_type == "custom_components": message += " Make sure your build method returns a component." - raise ValueError(message) + logger.warning(message) async def build(self, force: bool = False, user_id=None, *args, **kwargs) -> Any: if not self._built or force: From 163ef6973bf3996475815354f498718937d5c7d6 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 1 Dec 2023 16:30:26 -0300 Subject: [PATCH 12/17] Fix database initialization and migration --- src/backend/langflow/services/database/utils.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/backend/langflow/services/database/utils.py b/src/backend/langflow/services/database/utils.py index aae567e91..826b8d3a5 100644 --- a/src/backend/langflow/services/database/utils.py +++ b/src/backend/langflow/services/database/utils.py @@ -1,17 +1,18 @@ +from contextlib import contextmanager from dataclasses import dataclass from typing import TYPE_CHECKING -from loguru import logger -from contextlib import contextmanager + from alembic.util.exc import CommandError +from loguru import logger from sqlmodel import Session if TYPE_CHECKING: from langflow.services.database.service import DatabaseService -def initialize_database(): +def initialize_database(fix_migration: bool = False): logger.debug("Initializing database") - from langflow.services import service_manager, ServiceType + from langflow.services import ServiceType, service_manager database_service: "DatabaseService" = service_manager.get(ServiceType.DATABASE_SERVICE) try: @@ -28,7 +29,7 @@ def initialize_database(): logger.error(f"Error checking schema health: {exc}") raise RuntimeError("Error checking schema health") from exc try: - database_service.run_migrations() + database_service.run_migrations(fix=fix_migration) except CommandError as exc: # if "overlaps with other requested revisions" or "Can't locate revision identified by" # are not in the exception, we can't handle it From cb6228cea117fae9dbe0d4b1109ca8b90c501aac Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 1 Dec 2023 16:30:35 -0300 Subject: [PATCH 13/17] Fix imports and initialize services with fix_migration option --- src/backend/langflow/services/utils.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/src/backend/langflow/services/utils.py b/src/backend/langflow/services/utils.py index 40d2d7ea4..76d6f8b13 100644 --- a/src/backend/langflow/services/utils.py +++ b/src/backend/langflow/services/utils.py @@ -2,7 +2,8 @@ from langflow.services.auth.utils import create_super_user, verify_password from langflow.services.database.utils import initialize_database from langflow.services.manager import service_manager from langflow.services.schema import ServiceType -from langflow.services.settings.constants import DEFAULT_SUPERUSER, DEFAULT_SUPERUSER_PASSWORD +from langflow.services.settings.constants import (DEFAULT_SUPERUSER, + DEFAULT_SUPERUSER_PASSWORD) from loguru import logger from sqlmodel import Session @@ -13,12 +14,13 @@ def get_factories_and_deps(): from langflow.services.auth import factory as auth_factory from langflow.services.cache import factory as cache_factory from langflow.services.chat import factory as chat_factory + from langflow.services.credentials import factory as credentials_factory from langflow.services.database import factory as database_factory - from langflow.services.session import factory as session_service_factory # type: ignore + from langflow.services.session import \ + factory as session_service_factory # type: ignore from langflow.services.settings import factory as settings_factory from langflow.services.store import factory as store_factory from langflow.services.task import factory as task_factory - from langflow.services.credentials import factory as credentials_factory return [ (settings_factory.SettingsServiceFactory(), []), @@ -171,7 +173,8 @@ def initialize_session_service(): Initialize the session manager. """ from langflow.services.cache import factory as cache_factory - from langflow.services.session import factory as session_service_factory # type: ignore + from langflow.services.session import \ + factory as session_service_factory # type: ignore initialize_settings_service() @@ -183,7 +186,7 @@ def initialize_session_service(): ) -def initialize_services(): +def initialize_services(fix_migration: bool = False): """ Initialize all the services needed. """ @@ -197,7 +200,7 @@ def initialize_services(): # Test cache connection service_manager.get(ServiceType.CACHE_SERVICE) # Setup the superuser - initialize_database() + initialize_database(fix_migration=fix_migration) setup_superuser(service_manager.get(ServiceType.SETTINGS_SERVICE), next(get_session())) try: get_db_service().migrate_flows_if_auto_login() From be5cdc9c47560086ef9a62006de9dc240f28a71b Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 1 Dec 2023 16:30:43 -0300 Subject: [PATCH 14/17] Remove store_api_key and is_component columns from user and flow tables --- .../alembic/versions/7843803a87b5_store_updates.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/backend/langflow/alembic/versions/7843803a87b5_store_updates.py b/src/backend/langflow/alembic/versions/7843803a87b5_store_updates.py index 3d5f0c56b..54c418943 100644 --- a/src/backend/langflow/alembic/versions/7843803a87b5_store_updates.py +++ b/src/backend/langflow/alembic/versions/7843803a87b5_store_updates.py @@ -39,10 +39,12 @@ def upgrade() -> None: def downgrade() -> None: # ### commands auto generated by Alembic - please adjust! ### - with op.batch_alter_table("user", schema=None) as batch_op: - batch_op.drop_column("store_api_key") - - with op.batch_alter_table("flow", schema=None) as batch_op: - batch_op.drop_column("is_component") + try: + with op.batch_alter_table("user", schema=None) as batch_op: + batch_op.drop_column("store_api_key") + with op.batch_alter_table("flow", schema=None) as batch_op: + batch_op.drop_column("is_component") + except Exception: + pass # ### end Alembic commands ### From ab212a046f5e4a208e29b872869af4e695906bd2 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 1 Dec 2023 16:30:54 -0300 Subject: [PATCH 15/17] Fix migration issue with autogenerate diffs --- .../langflow/services/database/service.py | 44 +++++++++++++------ 1 file changed, 30 insertions(+), 14 deletions(-) diff --git a/src/backend/langflow/services/database/service.py b/src/backend/langflow/services/database/service.py index f41a2484d..3ecbfd971 100644 --- a/src/backend/langflow/services/database/service.py +++ b/src/backend/langflow/services/database/service.py @@ -5,17 +5,16 @@ from typing import TYPE_CHECKING import sqlalchemy as sa from alembic import command, util from alembic.config import Config -from loguru import logger -from sqlalchemy import inspect -from sqlalchemy.exc import OperationalError -from sqlmodel import Session, SQLModel, create_engine - from langflow.services.base import Service from langflow.services.database import models # noqa from langflow.services.database.models.user.crud import get_user_by_username from langflow.services.database.utils import Result, TableResults from langflow.services.deps import get_settings_service from langflow.services.utils import teardown_superuser +from loguru import logger +from sqlalchemy import inspect +from sqlalchemy.exc import OperationalError +from sqlmodel import Session, SQLModel, create_engine if TYPE_CHECKING: from sqlalchemy.engine import Engine @@ -118,9 +117,10 @@ class DatabaseService(Service): alembic_cfg.set_main_option("script_location", str(self.script_location)) alembic_cfg.set_main_option("sqlalchemy.url", self.database_url) command.stamp(alembic_cfg, "head") + # command.upgrade(alembic_cfg, "head") logger.info("Alembic initialized") - def run_migrations(self): + def run_migrations(self, fix=False): # First we need to check if alembic has been initialized # If not, we need to initialize it # if not self.script_location.exists(): # this is not the correct way to check if alembic has been initialized @@ -151,16 +151,32 @@ class DatabaseService(Service): if isinstance(exc, util.exc.CommandError) or isinstance(exc, util.exc.AutogenerateDiffsDetected): command.upgrade(alembic_cfg, "head") - # We should check the schema health after running migrations try: command.check(alembic_cfg) - except util.exc.AutogenerateDiffsDetected: - # downgrade to base and upgrade again - logger.warning("Autogenerate diffs detected, downgrading and upgrading") - command.downgrade(alembic_cfg, "-1") - # wait for the database to be ready - time.sleep(5) - command.upgrade(alembic_cfg, "head") + except util.exc.AutogenerateDiffsDetected as exc: + logger.exception("AutogenerateDiffsDetected: {exc}") + logger.warning("Something went wrong running migrations. Please, run `langflow migration --fix`") + + if fix: + self.try_downgrade_upgrade_until_success(alembic_cfg) + + def try_downgrade_upgrade_until_success(self, alembic_cfg, retries=5): + # Try -1 then head, if it fails, try -2 then head, etc. + # until we reach the number of retries + for i in range(1, retries + 1): + try: + command.check(alembic_cfg) + break + except util.exc.AutogenerateDiffsDetected as exc: + + # downgrade to base and upgrade again + logger.warning(f"AutogenerateDiffsDetected: {exc}") + command.downgrade(alembic_cfg, f"-{i}") + # wait for the database to be ready + time.sleep(3) + command.upgrade(alembic_cfg, "head") + + def run_migrations_test(self): # This method is used for testing purposes only From 8f04139ddd07121fca5a4257bf1354d54f47c73d Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 1 Dec 2023 16:30:59 -0300 Subject: [PATCH 16/17] Fix migration and add confirmation prompt --- src/backend/langflow/__main__.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/src/backend/langflow/__main__.py b/src/backend/langflow/__main__.py index 5eb92d133..d3c0ee8a2 100644 --- a/src/backend/langflow/__main__.py +++ b/src/backend/langflow/__main__.py @@ -12,7 +12,8 @@ from dotenv import load_dotenv from langflow.main import setup_app from langflow.services.database.utils import session_getter from langflow.services.deps import get_db_service, get_settings_service -from langflow.services.utils import initialize_services, initialize_settings_service +from langflow.services.utils import (initialize_services, + initialize_settings_service) from langflow.utils.logger import configure, logger from multiprocess import Process, cpu_count # type: ignore from rich import box @@ -327,11 +328,19 @@ def superuser( @app.command() -def migration(test: bool = typer.Option(True, help="Run migrations in test mode.")): +def migration(test: bool = typer.Option(True, help="Run migrations in test mode."), + fix: bool = typer.Option(False, help="Fix migrations. This is a destructive operation, and should only be used if you know what you are doing.") +): """ Run or test migrations. """ - initialize_services() + if fix: + if not typer.confirm("This will delete all data necessary to fix migrations. Are you sure you want to continue?"): + raise typer.Abort() + + + + initialize_services(fix_migration=fix) db_service = get_db_service() if not test: db_service.run_migrations() From dcefb8363ceee461dbc24c4cd5cede086c4729e4 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 1 Dec 2023 16:37:11 -0300 Subject: [PATCH 17/17] Fix migration error and initialize database in utils.py --- src/backend/langflow/services/database/service.py | 2 +- src/backend/langflow/services/database/utils.py | 5 +++-- src/backend/langflow/services/utils.py | 6 +++++- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/src/backend/langflow/services/database/service.py b/src/backend/langflow/services/database/service.py index 3ecbfd971..86793f779 100644 --- a/src/backend/langflow/services/database/service.py +++ b/src/backend/langflow/services/database/service.py @@ -155,7 +155,7 @@ class DatabaseService(Service): command.check(alembic_cfg) except util.exc.AutogenerateDiffsDetected as exc: logger.exception("AutogenerateDiffsDetected: {exc}") - logger.warning("Something went wrong running migrations. Please, run `langflow migration --fix`") + raise RuntimeError("Something went wrong running migrations. Please, run `langflow migration --fix`") if fix: self.try_downgrade_upgrade_until_success(alembic_cfg) diff --git a/src/backend/langflow/services/database/utils.py b/src/backend/langflow/services/database/utils.py index 826b8d3a5..afabfb5e2 100644 --- a/src/backend/langflow/services/database/utils.py +++ b/src/backend/langflow/services/database/utils.py @@ -48,8 +48,9 @@ def initialize_database(fix_migration: bool = False): # if the exception involves tables already existing # we can ignore it if "already exists" not in str(exc): - logger.error(f"Error running migrations: {exc}") - raise RuntimeError("Error running migrations") from exc + logger.error(exc) + + raise exc logger.debug("Database initialized") diff --git a/src/backend/langflow/services/utils.py b/src/backend/langflow/services/utils.py index 76d6f8b13..fa14d59a1 100644 --- a/src/backend/langflow/services/utils.py +++ b/src/backend/langflow/services/utils.py @@ -200,7 +200,11 @@ def initialize_services(fix_migration: bool = False): # Test cache connection service_manager.get(ServiceType.CACHE_SERVICE) # Setup the superuser - initialize_database(fix_migration=fix_migration) + try: + initialize_database(fix_migration=fix_migration) + except Exception as exc: + logger.exception(exc) + raise exc setup_superuser(service_manager.get(ServiceType.SETTINGS_SERVICE), next(get_session())) try: get_db_service().migrate_flows_if_auto_login()