diff --git a/src/backend/langflow/cache/__init__.py b/src/backend/langflow/cache/__init__.py deleted file mode 100644 index 723aa9e18..000000000 --- a/src/backend/langflow/cache/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -from langflow.cache.manager import cache_manager -from langflow.cache.flow import InMemoryCache - -__all__ = [ - "cache_manager", - "InMemoryCache", -] diff --git a/src/backend/langflow/cache/base.py b/src/backend/langflow/cache/base.py deleted file mode 100644 index 88cb3a1da..000000000 --- a/src/backend/langflow/cache/base.py +++ /dev/null @@ -1,84 +0,0 @@ -import abc - - -class BaseCache(abc.ABC): - """ - Abstract base class for a cache. - """ - - @abc.abstractmethod - def get(self, key): - """ - Retrieve an item from the cache. - - Args: - key: The key of the item to retrieve. - - Returns: - The value associated with the key, or None if the key is not found. - """ - - @abc.abstractmethod - def set(self, key, value): - """ - Add an item to the cache. - - Args: - key: The key of the item. - value: The value to cache. - """ - - @abc.abstractmethod - def delete(self, key): - """ - Remove an item from the cache. - - Args: - key: The key of the item to remove. - """ - - @abc.abstractmethod - def clear(self): - """ - Clear all items from the cache. - """ - - @abc.abstractmethod - def __contains__(self, key): - """ - Check if the key is in the cache. - - Args: - key: The key of the item to check. - - Returns: - True if the key is in the cache, False otherwise. - """ - - @abc.abstractmethod - def __getitem__(self, key): - """ - Retrieve an item from the cache using the square bracket notation. - - Args: - key: The key of the item to retrieve. - """ - - @abc.abstractmethod - def __setitem__(self, key, value): - """ - Add an item to the cache using the square bracket notation. - - Args: - key: The key of the item. - value: The value to cache. - """ - - @abc.abstractmethod - def __delitem__(self, key): - """ - Remove an item from the cache using the square bracket notation. - - Args: - key: The key of the item to remove. - """ diff --git a/src/backend/langflow/cache/flow.py b/src/backend/langflow/cache/flow.py deleted file mode 100644 index 6d8fee977..000000000 --- a/src/backend/langflow/cache/flow.py +++ /dev/null @@ -1,146 +0,0 @@ -import threading -import time -from collections import OrderedDict - -from langflow.cache.base import BaseCache - - -class InMemoryCache(BaseCache): - """ - A simple in-memory cache using an OrderedDict. - - This cache supports setting a maximum size and expiration time for cached items. - When the cache is full, it uses a Least Recently Used (LRU) eviction policy. - Thread-safe using a threading Lock. - - Attributes: - max_size (int, optional): Maximum number of items to store in the cache. - expiration_time (int, optional): Time in seconds after which a cached item expires. Default is 1 hour. - - Example: - - cache = InMemoryCache(max_size=3, expiration_time=5) - - # setting cache values - cache.set("a", 1) - cache.set("b", 2) - cache["c"] = 3 - - # getting cache values - a = cache.get("a") - b = cache["b"] - """ - - def __init__(self, max_size=None, expiration_time=60 * 60): - """ - Initialize a new InMemoryCache instance. - - Args: - max_size (int, optional): Maximum number of items to store in the cache. - expiration_time (int, optional): Time in seconds after which a cached item expires. Default is 1 hour. - """ - self._cache = OrderedDict() - self._lock = threading.Lock() - self.max_size = max_size - self.expiration_time = expiration_time - - def get(self, key): - """ - Retrieve an item from the cache. - - Args: - key: The key of the item to retrieve. - - Returns: - The value associated with the key, or None if the key is not found or the item has expired. - """ - with self._lock: - if key in self._cache: - item = self._cache.pop(key) - if ( - self.expiration_time is None - or time.time() - item["time"] < self.expiration_time - ): - # Move the key to the end to make it recently used - self._cache[key] = item - return item["value"] - else: - self.delete(key) - return None - - def set(self, key, value): - """ - Add an item to the cache. - - If the cache is full, the least recently used item is evicted. - - Args: - key: The key of the item. - value: The value to cache. - """ - with self._lock: - if key in self._cache: - # Remove existing key before re-inserting to update order - self.delete(key) - elif self.max_size and len(self._cache) >= self.max_size: - # Remove least recently used item - self._cache.popitem(last=False) - self._cache[key] = {"value": value, "time": time.time()} - - def get_or_set(self, key, value): - """ - Retrieve an item from the cache. If the item does not exist, set it with the provided value. - - Args: - key: The key of the item. - value: The value to cache if the item doesn't exist. - - Returns: - The cached value associated with the key. - """ - with self._lock: - if key in self._cache: - return self.get(key) - self.set(key, value) - return value - - def delete(self, key): - """ - Remove an item from the cache. - - Args: - key: The key of the item to remove. - """ - # with self._lock: - self._cache.pop(key, None) - - def clear(self): - """ - Clear all items from the cache. - """ - with self._lock: - self._cache.clear() - - def __contains__(self, key): - """Check if the key is in the cache.""" - return key in self._cache - - def __getitem__(self, key): - """Retrieve an item from the cache using the square bracket notation.""" - return self.get(key) - - def __setitem__(self, key, value): - """Add an item to the cache using the square bracket notation.""" - self.set(key, value) - - def __delitem__(self, key): - """Remove an item from the cache using the square bracket notation.""" - self.delete(key) - - def __len__(self): - """Return the number of items in the cache.""" - return len(self._cache) - - def __repr__(self): - """Return a string representation of the InMemoryCache instance.""" - return f"InMemoryCache(max_size={self.max_size}, expiration_time={self.expiration_time})" diff --git a/src/backend/langflow/cache/manager.py b/src/backend/langflow/cache/manager.py deleted file mode 100644 index 13b281008..000000000 --- a/src/backend/langflow/cache/manager.py +++ /dev/null @@ -1,150 +0,0 @@ -from contextlib import contextmanager -from typing import Any, Awaitable, Callable, List, Optional - -import pandas as pd -from PIL import Image - - -class Subject: - """Base class for implementing the observer pattern.""" - - def __init__(self): - self.observers: List[Callable[[], None]] = [] - - def attach(self, observer: Callable[[], None]): - """Attach an observer to the subject.""" - self.observers.append(observer) - - def detach(self, observer: Callable[[], None]): - """Detach an observer from the subject.""" - self.observers.remove(observer) - - def notify(self): - """Notify all observers about an event.""" - for observer in self.observers: - if observer is None: - continue - observer() - - -class AsyncSubject: - """Base class for implementing the async observer pattern.""" - - def __init__(self): - self.observers: List[Callable[[], Awaitable]] = [] - - def attach(self, observer: Callable[[], Awaitable]): - """Attach an observer to the subject.""" - self.observers.append(observer) - - def detach(self, observer: Callable[[], Awaitable]): - """Detach an observer from the subject.""" - self.observers.remove(observer) - - async def notify(self): - """Notify all observers about an event.""" - for observer in self.observers: - if observer is None: - continue - await observer() - - -class CacheManager(Subject): - """Manages cache for different clients and notifies observers on changes.""" - - def __init__(self): - super().__init__() - self._cache = {} - self.current_client_id = None - self.current_cache = {} - - @contextmanager - def set_client_id(self, client_id: str): - """ - Context manager to set the current client_id and associated cache. - - Args: - client_id (str): The client identifier. - """ - previous_client_id = self.current_client_id - self.current_client_id = client_id - self.current_cache = self._cache.setdefault(client_id, {}) - try: - yield - finally: - self.current_client_id = previous_client_id - self.current_cache = self._cache.get(self.current_client_id, {}) - - def add(self, name: str, obj: Any, obj_type: str, extension: Optional[str] = None): - """ - Add an object to the current client's cache. - - Args: - name (str): The cache key. - obj (Any): The object to cache. - obj_type (str): The type of the object. - """ - object_extensions = { - "image": "png", - "pandas": "csv", - } - if obj_type in object_extensions: - _extension = object_extensions[obj_type] - else: - _extension = type(obj).__name__.lower() - self.current_cache[name] = { - "obj": obj, - "type": obj_type, - "extension": extension or _extension, - } - self.notify() - - def add_pandas(self, name: str, obj: Any): - """ - Add a pandas DataFrame or Series to the current client's cache. - - Args: - name (str): The cache key. - obj (Any): The pandas DataFrame or Series object. - """ - if isinstance(obj, (pd.DataFrame, pd.Series)): - self.add(name, obj.to_csv(), "pandas", extension="csv") - else: - raise ValueError("Object is not a pandas DataFrame or Series") - - def add_image(self, name: str, obj: Any, extension: str = "png"): - """ - Add a PIL Image to the current client's cache. - - Args: - name (str): The cache key. - obj (Any): The PIL Image object. - """ - if isinstance(obj, Image.Image): - self.add(name, obj, "image", extension=extension) - else: - raise ValueError("Object is not a PIL Image") - - def get(self, name: str): - """ - Get an object from the current client's cache. - - Args: - name (str): The cache key. - - Returns: - The cached object associated with the given cache key. - """ - return self.current_cache[name] - - def get_last(self): - """ - Get the last added item in the current client's cache. - - Returns: - The last added item in the cache. - """ - return list(self.current_cache.values())[-1] - - -cache_manager = CacheManager() diff --git a/src/backend/langflow/cache/utils.py b/src/backend/langflow/cache/utils.py deleted file mode 100644 index 3deabe9f4..000000000 --- a/src/backend/langflow/cache/utils.py +++ /dev/null @@ -1,179 +0,0 @@ -import base64 -import contextlib -import functools -import hashlib -import json -import os -import tempfile -from collections import OrderedDict -from pathlib import Path -from typing import Any, Dict -from appdirs import user_cache_dir - -CACHE: Dict[str, Any] = {} - -CACHE_DIR = user_cache_dir("langflow", "langflow") - - -def create_cache_folder(func): - def wrapper(*args, **kwargs): - # Get the destination folder - cache_path = Path(CACHE_DIR) / PREFIX - - # Create the destination folder if it doesn't exist - os.makedirs(cache_path, exist_ok=True) - - return func(*args, **kwargs) - - return wrapper - - -def memoize_dict(maxsize=128): - cache = OrderedDict() - - def decorator(func): - @functools.wraps(func) - def wrapper(*args, **kwargs): - hashed = compute_dict_hash(args[0]) - key = (func.__name__, hashed, frozenset(kwargs.items())) - if key not in cache: - result = func(*args, **kwargs) - cache[key] = result - if len(cache) > maxsize: - cache.popitem(last=False) - else: - result = cache[key] - return result - - def clear_cache(): - cache.clear() - - wrapper.clear_cache = clear_cache # type: ignore - wrapper.cache = cache # type: ignore - return wrapper - - return decorator - - -PREFIX = "langflow_cache" - - -@create_cache_folder -def clear_old_cache_files(max_cache_size: int = 3): - cache_dir = Path(tempfile.gettempdir()) / PREFIX - cache_files = list(cache_dir.glob("*.dill")) - - if len(cache_files) > max_cache_size: - cache_files_sorted_by_mtime = sorted( - cache_files, key=lambda x: x.stat().st_mtime, reverse=True - ) - - for cache_file in cache_files_sorted_by_mtime[max_cache_size:]: - with contextlib.suppress(OSError): - os.remove(cache_file) - - -def compute_dict_hash(graph_data): - graph_data = filter_json(graph_data) - - cleaned_graph_json = json.dumps(graph_data, sort_keys=True) - return hashlib.sha256(cleaned_graph_json.encode("utf-8")).hexdigest() - - -def filter_json(json_data): - filtered_data = json_data.copy() - - # Remove 'viewport' and 'chatHistory' keys - if "viewport" in filtered_data: - del filtered_data["viewport"] - if "chatHistory" in filtered_data: - del filtered_data["chatHistory"] - - # Filter nodes - if "nodes" in filtered_data: - for node in filtered_data["nodes"]: - if "position" in node: - del node["position"] - if "positionAbsolute" in node: - del node["positionAbsolute"] - if "selected" in node: - del node["selected"] - if "dragging" in node: - del node["dragging"] - - return filtered_data - - -@create_cache_folder -def save_binary_file(content: str, file_name: str, accepted_types: list[str]) -> str: - """ - Save a binary file to the specified folder. - - Args: - content: The content of the file as a bytes object. - file_name: The name of the file, including its extension. - - Returns: - The path to the saved file. - """ - if not any(file_name.endswith(suffix) for suffix in accepted_types): - raise ValueError(f"File {file_name} is not accepted") - - # Get the destination folder - cache_path = Path(CACHE_DIR) / PREFIX - if not content: - raise ValueError("Please, reload the file in the loader.") - data = content.split(",")[1] - decoded_bytes = base64.b64decode(data) - - # Create the full file path - file_path = os.path.join(cache_path, file_name) - - # Save the binary content to the file - with open(file_path, "wb") as file: - file.write(decoded_bytes) - - return file_path - - -@create_cache_folder -def save_uploaded_file(file, folder_name): - """ - Save an uploaded file to the specified folder with a hash of its content as the file name. - - Args: - file: The uploaded file object. - folder_name: The name of the folder to save the file in. - - Returns: - The path to the saved file. - """ - cache_path = Path(CACHE_DIR) - folder_path = cache_path / folder_name - - # Create the folder if it doesn't exist - if not folder_path.exists(): - folder_path.mkdir() - - # Create a hash of the file content - sha256_hash = hashlib.sha256() - # Reset the file cursor to the beginning of the file - file.seek(0) - # Iterate over the uploaded file in small chunks to conserve memory - while chunk := file.read(8192): # Read 8KB at a time (adjust as needed) - sha256_hash.update(chunk) - - # Use the hex digest of the hash as the file name - hex_dig = sha256_hash.hexdigest() - file_name = hex_dig - - # Reset the file cursor to the beginning of the file - file.seek(0) - - # Save the file with the hash as its name - file_path = folder_path / file_name - with open(file_path, "wb") as new_file: - while chunk := file.read(8192): - new_file.write(chunk) - - return file_path diff --git a/src/backend/langflow/chat/__init__.py b/src/backend/langflow/chat/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/src/backend/langflow/chat/config.py b/src/backend/langflow/chat/config.py deleted file mode 100644 index 274f4d5bd..000000000 --- a/src/backend/langflow/chat/config.py +++ /dev/null @@ -1,2 +0,0 @@ -class ChatConfig: - streaming: bool = True diff --git a/src/backend/langflow/chat/manager.py b/src/backend/langflow/chat/manager.py deleted file mode 100644 index 1e93174e2..000000000 --- a/src/backend/langflow/chat/manager.py +++ /dev/null @@ -1,217 +0,0 @@ -from collections import defaultdict -from fastapi import WebSocket, status -from langflow.api.v1.schemas import ChatMessage, ChatResponse, FileResponse -from langflow.cache import cache_manager -from langflow.cache.manager import Subject -from langflow.chat.utils import process_graph -from langflow.interface.utils import pil_to_base64 -from langflow.utils.logger import logger - - -import asyncio -import json -from typing import Any, Dict, List - -from langflow.cache.flow import InMemoryCache - - -class ChatHistory(Subject): - def __init__(self): - super().__init__() - self.history: Dict[str, List[ChatMessage]] = defaultdict(list) - - def add_message(self, client_id: str, message: ChatMessage): - """Add a message to the chat history.""" - - self.history[client_id].append(message) - - if not isinstance(message, FileResponse): - self.notify() - - def get_history(self, client_id: str, filter_messages=True) -> List[ChatMessage]: - """Get the chat history for a client.""" - if history := self.history.get(client_id, []): - if filter_messages: - return [msg for msg in history if msg.type not in ["start", "stream"]] - return history - else: - return [] - - def empty_history(self, client_id: str): - """Empty the chat history for a client.""" - self.history[client_id] = [] - - -class ChatManager: - def __init__(self): - self.active_connections: Dict[str, WebSocket] = {} - self.chat_history = ChatHistory() - self.cache_manager = cache_manager - self.cache_manager.attach(self.update) - self.in_memory_cache = InMemoryCache() - - def on_chat_history_update(self): - """Send the last chat message to the client.""" - client_id = self.cache_manager.current_client_id - if client_id in self.active_connections: - chat_response = self.chat_history.get_history( - client_id, filter_messages=False - )[-1] - if chat_response.is_bot: - # Process FileResponse - if isinstance(chat_response, FileResponse): - # If data_type is pandas, convert to csv - if chat_response.data_type == "pandas": - chat_response.data = chat_response.data.to_csv() - elif chat_response.data_type == "image": - # Base64 encode the image - chat_response.data = pil_to_base64(chat_response.data) - # get event loop - loop = asyncio.get_event_loop() - - coroutine = self.send_json(client_id, chat_response) - asyncio.run_coroutine_threadsafe(coroutine, loop) - - def update(self): - if self.cache_manager.current_client_id in self.active_connections: - self.last_cached_object_dict = self.cache_manager.get_last() - # Add a new ChatResponse with the data - chat_response = FileResponse( - message=None, - type="file", - data=self.last_cached_object_dict["obj"], - data_type=self.last_cached_object_dict["type"], - ) - - self.chat_history.add_message( - self.cache_manager.current_client_id, chat_response - ) - - async def connect(self, client_id: str, websocket: WebSocket): - await websocket.accept() - self.active_connections[client_id] = websocket - - def disconnect(self, client_id: str): - self.active_connections.pop(client_id, None) - - async def send_message(self, client_id: str, message: str): - websocket = self.active_connections[client_id] - await websocket.send_text(message) - - async def send_json(self, client_id: str, message: ChatMessage): - websocket = self.active_connections[client_id] - await websocket.send_json(message.dict()) - - async def close_connection(self, client_id: str, code: int, reason: str): - if websocket := self.active_connections[client_id]: - try: - await websocket.close(code=code, reason=reason) - self.disconnect(client_id) - except RuntimeError as exc: - # This is to catch the following error: - # Unexpected ASGI message 'websocket.close', after sending 'websocket.close' - if "after sending" in str(exc): - logger.error(f"Error closing connection: {exc}") - - async def process_message( - self, client_id: str, payload: Dict, langchain_object: Any - ): - # Process the graph data and chat message - chat_inputs = payload.pop("inputs", "") - chat_inputs = ChatMessage(message=chat_inputs) - self.chat_history.add_message(client_id, chat_inputs) - - # graph_data = payload - start_resp = ChatResponse(message=None, type="start", intermediate_steps="") - await self.send_json(client_id, start_resp) - - # is_first_message = len(self.chat_history.get_history(client_id=client_id)) <= 1 - # Generate result and thought - try: - logger.debug("Generating result and thought") - - result, intermediate_steps = await process_graph( - langchain_object=langchain_object, - chat_inputs=chat_inputs, - websocket=self.active_connections[client_id], - ) - except Exception as e: - # Log stack trace - logger.exception(e) - self.chat_history.empty_history(client_id) - raise e - # Send a response back to the frontend, if needed - intermediate_steps = intermediate_steps or "" - history = self.chat_history.get_history(client_id, filter_messages=False) - file_responses = [] - if history: - # Iterate backwards through the history - for msg in reversed(history): - if isinstance(msg, FileResponse): - if msg.data_type == "image": - # Base64 encode the image - if isinstance(msg.data, str): - continue - msg.data = pil_to_base64(msg.data) - file_responses.append(msg) - if msg.type == "start": - break - - response = ChatResponse( - message=result, - intermediate_steps=intermediate_steps.strip(), - type="end", - files=file_responses, - ) - await self.send_json(client_id, response) - self.chat_history.add_message(client_id, response) - - def set_cache(self, client_id: str, langchain_object: Any) -> bool: - """ - Set the cache for a client. - """ - - self.in_memory_cache.set(client_id, langchain_object) - return client_id in self.in_memory_cache - - async def handle_websocket(self, client_id: str, websocket: WebSocket): - await self.connect(client_id, websocket) - - try: - chat_history = self.chat_history.get_history(client_id) - # iterate and make BaseModel into dict - chat_history = [chat.dict() for chat in chat_history] - await websocket.send_json(chat_history) - - while True: - json_payload = await websocket.receive_json() - try: - payload = json.loads(json_payload) - except TypeError: - payload = json_payload - if "clear_history" in payload: - self.chat_history.history[client_id] = [] - continue - - with self.cache_manager.set_client_id(client_id): - langchain_object = self.in_memory_cache.get(client_id) - await self.process_message(client_id, payload, langchain_object) - - except Exception as exc: - # Handle any exceptions that might occur - logger.error(f"Error handling websocket: {exc}") - await self.close_connection( - client_id=client_id, - code=status.WS_1011_INTERNAL_ERROR, - reason=str(exc)[:120], - ) - finally: - try: - await self.close_connection( - client_id=client_id, - code=status.WS_1000_NORMAL_CLOSURE, - reason="Client disconnected", - ) - except Exception as exc: - logger.error(f"Error closing connection: {exc}") - self.disconnect(client_id) diff --git a/src/backend/langflow/chat/utils.py b/src/backend/langflow/chat/utils.py deleted file mode 100644 index 7db65b8e3..000000000 --- a/src/backend/langflow/chat/utils.py +++ /dev/null @@ -1,37 +0,0 @@ -from fastapi import WebSocket -from langflow.api.v1.schemas import ChatMessage -from langflow.processing.base import get_result_and_steps -from langflow.interface.utils import try_setting_streaming_options -from langflow.utils.logger import logger - - -async def process_graph( - langchain_object, - chat_inputs: ChatMessage, - websocket: WebSocket, -): - langchain_object = try_setting_streaming_options(langchain_object, websocket) - logger.debug("Loaded langchain object") - - if langchain_object is None: - # Raise user facing error - raise ValueError( - "There was an error loading the langchain_object. Please, check all the nodes and try again." - ) - - # Generate result and thought - try: - if not chat_inputs.message: - logger.debug("No message provided") - raise ValueError("No message provided") - - logger.debug("Generating result and thought") - result, intermediate_steps = await get_result_and_steps( - langchain_object, chat_inputs.message, websocket=websocket - ) - logger.debug("Generated result and intermediate_steps") - return result, intermediate_steps - except Exception as e: - # Log stack trace - logger.exception(e) - raise e diff --git a/src/backend/langflow/database/__init__.py b/src/backend/langflow/database/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/src/backend/langflow/database/base.py b/src/backend/langflow/database/base.py deleted file mode 100644 index 518b95d15..000000000 --- a/src/backend/langflow/database/base.py +++ /dev/null @@ -1,133 +0,0 @@ -from contextlib import contextmanager -import os -from pathlib import Path -from langflow.database import models # noqa -from sqlmodel import SQLModel, Session, create_engine -from langflow.utils.logger import logger -from alembic.config import Config -from alembic import command - - -class Engine: - _instance = None - - @classmethod - def get(cls): - logger.debug("Getting database engine") - if cls._instance is None: - cls.create() - return cls._instance - - @classmethod - def create(cls): - logger.debug("Creating database engine") - from langflow.settings import settings - - if langflow_database_url := os.getenv("LANGFLOW_DATABASE_URL"): - settings.DATABASE_URL = langflow_database_url - logger.debug("Using LANGFLOW_DATABASE_URL") - - if settings.DATABASE_URL and settings.DATABASE_URL.startswith("sqlite"): - connect_args = {"check_same_thread": False} - else: - connect_args = {} - if not settings.DATABASE_URL: - raise RuntimeError("No database_url provided") - cls._instance = create_engine(settings.DATABASE_URL, connect_args=connect_args) - - @classmethod - def update(cls): - logger.debug("Updating database engine") - cls._instance = None - cls.create() - - -def create_db_and_tables(): - logger.debug("Creating database and tables") - try: - SQLModel.metadata.create_all(Engine.get()) - except Exception as exc: - logger.error(f"Error creating database and tables: {exc}") - raise RuntimeError("Error creating database and tables") from exc - # Now check if the table Flow exists, if not, something went wrong - # and we need to create the tables again. - from sqlalchemy import inspect - - inspector = inspect(Engine.get()) - if "flow" not in inspector.get_table_names(): - logger.error("Something went wrong creating the database and tables.") - logger.error("Please check your database settings.") - - raise RuntimeError("Something went wrong creating the database and tables.") - else: - logger.debug("Database and tables created successfully") - - -class DatabaseManager: - def __init__(self, database_url: str): - self.database_url = database_url - # This file is in langflow.database.base.py - # the ini is in langflow - self.script_location = Path(__file__).parent.parent / "alembic" - self.alembic_cfg_path = Path(__file__).parent.parent / "alembic.ini" - self.engine = create_engine(database_url) - - def __enter__(self): - self._session = Session(self.engine) - return self._session - - def __exit__(self, exc_type, exc_value, traceback): - if exc_type is not None: # If an exception has been raised - logger.error( - f"Session rollback because of exception: {exc_type.__name__} {exc_value}" - ) - self._session.rollback() - else: - self._session.commit() - self._session.close() - - def get_session(self): - with Session(self.engine) as session: - yield session - - def run_migrations(self): - logger.info( - f"Running DB migrations in {self.script_location} on {self.database_url}" - ) - alembic_cfg = Config() - alembic_cfg.set_main_option("script_location", str(self.script_location)) - alembic_cfg.set_main_option("sqlalchemy.url", self.database_url) - command.upgrade(alembic_cfg, "head") - - def create_db_and_tables(self): - logger.debug("Creating database and tables") - try: - SQLModel.metadata.create_all(self.engine) - except Exception as exc: - logger.error(f"Error creating database and tables: {exc}") - raise RuntimeError("Error creating database and tables") from exc - - # Now check if the table "flow" exists, if not, something went wrong - # and we need to create the tables again. - from sqlalchemy import inspect - - inspector = inspect(self.engine) - if "flow" not in inspector.get_table_names(): - logger.error("Something went wrong creating the database and tables.") - logger.error("Please check your database settings.") - raise RuntimeError("Something went wrong creating the database and tables.") - else: - logger.debug("Database and tables created successfully") - - -@contextmanager -def session_getter(db_manager: DatabaseManager): - try: - session = Session(DatabaseManager.engine) - yield session - except Exception as e: - print("Session rollback because of exception:", e) - session.rollback() - raise - finally: - session.close() diff --git a/src/backend/langflow/database/models/__init__.py b/src/backend/langflow/database/models/__init__.py deleted file mode 100644 index da47bc5fe..000000000 --- a/src/backend/langflow/database/models/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .flow import Flow - - -__all__ = ["Flow"] diff --git a/src/backend/langflow/database/models/base.py b/src/backend/langflow/database/models/base.py deleted file mode 100644 index e20895b93..000000000 --- a/src/backend/langflow/database/models/base.py +++ /dev/null @@ -1,14 +0,0 @@ -from sqlmodel import SQLModel -import orjson - - -def orjson_dumps(v, *, default): - # orjson.dumps returns bytes, to match standard json.dumps we need to decode - return orjson.dumps(v, default=default).decode() - - -class SQLModelSerializable(SQLModel): - class Config: - orm_mode = True - json_loads = orjson.loads - json_dumps = orjson_dumps diff --git a/src/backend/langflow/database/models/component.py b/src/backend/langflow/database/models/component.py deleted file mode 100644 index bb2408cdb..000000000 --- a/src/backend/langflow/database/models/component.py +++ /dev/null @@ -1,29 +0,0 @@ -from langflow.database.models.base import SQLModelSerializable, SQLModel -from sqlmodel import Field -from typing import Optional -from datetime import datetime -import uuid - - -class Component(SQLModelSerializable, table=True): - id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True) - frontend_node_id: uuid.UUID = Field(index=True) - name: str = Field(index=True) - description: Optional[str] = Field(default=None) - python_code: Optional[str] = Field(default=None) - return_type: Optional[str] = Field(default=None) - is_disabled: bool = Field(default=False) - is_read_only: bool = Field(default=False) - create_at: datetime = Field(default_factory=datetime.utcnow) - update_at: datetime = Field(default_factory=datetime.utcnow) - - -class ComponentModel(SQLModel): - id: uuid.UUID = Field(default_factory=uuid.uuid4) - frontend_node_id: uuid.UUID = Field(default=uuid.uuid4()) - name: str = Field(default="") - description: Optional[str] = None - python_code: Optional[str] = None - return_type: Optional[str] = None - is_disabled: bool = False - is_read_only: bool = False diff --git a/src/backend/langflow/database/models/flow.py b/src/backend/langflow/database/models/flow.py deleted file mode 100644 index f9e3aa249..000000000 --- a/src/backend/langflow/database/models/flow.py +++ /dev/null @@ -1,60 +0,0 @@ -# Path: src/backend/langflow/database/models/flow.py - -from langflow.database.models.base import SQLModelSerializable -from pydantic import validator -from sqlmodel import Field, Relationship, JSON, Column -from uuid import UUID, uuid4 -from typing import Dict, Optional - -# if TYPE_CHECKING: -from langflow.database.models.flow_style import FlowStyle, FlowStyleRead - - -class FlowBase(SQLModelSerializable): - name: str = Field(index=True) - description: Optional[str] = Field(index=True) - data: Optional[Dict] = Field(default=None) - - @validator("data") - def validate_json(v): - # dict_keys(['description', 'name', 'id', 'data']) - if not v: - return v - if not isinstance(v, dict): - raise ValueError("Flow must be a valid JSON") - - # data must contain nodes and edges - if "nodes" not in v.keys(): - raise ValueError("Flow must have nodes") - if "edges" not in v.keys(): - raise ValueError("Flow must have edges") - - return v - - -class Flow(FlowBase, table=True): - id: UUID = Field(default_factory=uuid4, primary_key=True, unique=True) - data: Optional[Dict] = Field(default=None, sa_column=Column(JSON)) - style: Optional["FlowStyle"] = Relationship( - back_populates="flow", - # use "uselist=False" to make it a one-to-one relationship - sa_relationship_kwargs={"uselist": False}, - ) - - -class FlowCreate(FlowBase): - pass - - -class FlowRead(FlowBase): - id: UUID - - -class FlowReadWithStyle(FlowRead): - style: Optional["FlowStyleRead"] = None - - -class FlowUpdate(SQLModelSerializable): - name: Optional[str] = None - description: Optional[str] = None - data: Optional[Dict] = None diff --git a/src/backend/langflow/database/models/flow_style.py b/src/backend/langflow/database/models/flow_style.py deleted file mode 100644 index fe53799fe..000000000 --- a/src/backend/langflow/database/models/flow_style.py +++ /dev/null @@ -1,33 +0,0 @@ -# Path: src/backend/langflow/database/models/flowstyle.py - -from langflow.database.models.base import SQLModelSerializable -from sqlmodel import Field, Relationship -from uuid import UUID, uuid4 -from typing import TYPE_CHECKING, Optional - -if TYPE_CHECKING: - from langflow.database.models.flow import Flow - - -class FlowStyleBase(SQLModelSerializable): - color: str - emoji: str - flow_id: UUID = Field(default=None, foreign_key="flow.id") - - -class FlowStyle(FlowStyleBase, table=True): - id: UUID = Field(default_factory=uuid4, primary_key=True, unique=True) - flow: "Flow" = Relationship(back_populates="style") - - -class FlowStyleUpdate(SQLModelSerializable): - color: Optional[str] = None - emoji: Optional[str] = None - - -class FlowStyleCreate(FlowStyleBase): - pass - - -class FlowStyleRead(FlowStyleBase): - id: UUID