🔧 fix(env.py): update import statement for SQLModel in env.py to match new file structure
🔧 fix(chat.py): update import statements for ChatService and BaseCacheService in chat.py to match new file structure 🔧 fix(__init__.py): update import statement for InMemoryCache in __init__.py to match new file structure 🔧 fix(factory.py): update import statement for InMemoryCache, RedisCache, and BaseCacheService in factory.py to match new file structure 🔧 fix(service.py): create new file service.py and move InMemoryCache class from base.py to service.py ✨ feat(service.py): add InMemoryCache class to provide a simple in-memory cache implementation ✨ feat(service.py): add RedisCache class to provide a Redis-based cache implementation 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): update import statements to reflect changes in file structure 🔨 refactor(factory.py): 📦 chore(service.py): add new file `service.py` to implement DatabaseService class 📦 chore(utils.py): add import statement for DatabaseService in `utils.py` 📦 chore(getters.py): add import statement for DatabaseService in `getters.py` 📦 chore(factory.py): add import statement for SessionService in `factory.py` 📦 chore(service.py): add new file `service.py` to implement SessionService class 📦 chore(service.py): add SettingsService and TaskService classes to handle settings and tasks respectively 🔧 refactor(service.py): refactor load_settings_from_yaml method to load settings from a YAML file and convert keys to uppercase for consistency 🔧 refactor(service.py): refactor get_backend method to return the appropriate backend based on USE_CELERY flag 🔧 refactor(service.py): refactor launch_and_await_task method to handle launching and awaiting tasks based on USE_CELERY flag 🔧 refactor(service.py): refactor launch_task method to handle launching tasks based on USE_CELERY flag 🔧 refactor(service.py): refactor get_task method to get a task based on USE_CELERY flag 🔧 refactor(service.py): refactor check_celery_availability method to check if Celery is available and return its status 🔧 refactor(service.py): refactor USE_CELERY flag to check if Celery is available and set it accordingly
This commit is contained in:
Gabriel Luiz Freitas Almeida
parent
c1bc2de720
commit
75ac412662
15 changed files with 20 additions and 20 deletions
|
|
@ -5,7 +5,7 @@ from sqlalchemy import pool
|
|||
|
||||
from alembic import context
|
||||
|
||||
from langflow.services.database.manager import SQLModel
|
||||
from langflow.services.database.service import SQLModel
|
||||
|
||||
# this is the Alembic Config object, which provides
|
||||
# access to the values within the .ini file in use.
|
||||
|
|
|
|||
|
|
@ -17,8 +17,8 @@ from langflow.services.cache.utils import update_build_status
|
|||
from loguru import logger
|
||||
from langflow.services.getters import get_chat_service, get_session, get_cache_service
|
||||
from sqlmodel import Session
|
||||
from langflow.services.chat.manager import ChatService
|
||||
from langflow.services.cache.manager import BaseCacheService
|
||||
from langflow.services.chat.service import ChatService
|
||||
from langflow.services.cache.service import BaseCacheService
|
||||
|
||||
|
||||
router = APIRouter(tags=["Chat"])
|
||||
|
|
|
|||
|
|
@ -1,9 +1,9 @@
|
|||
from . import factory, manager
|
||||
from langflow.services.cache.manager import InMemoryCache
|
||||
from . import factory, service
|
||||
from langflow.services.cache.service import InMemoryCache
|
||||
|
||||
|
||||
__all__ = [
|
||||
"factory",
|
||||
"manager",
|
||||
"service",
|
||||
"InMemoryCache",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -1,10 +1,10 @@
|
|||
from langflow.services.cache.manager import InMemoryCache, RedisCache, BaseCacheService
|
||||
from langflow.services.cache.service import InMemoryCache, RedisCache, BaseCacheService
|
||||
from langflow.services.factory import ServiceFactory
|
||||
from langflow.utils.logger import logger
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langflow.services.settings.manager import SettingsService
|
||||
from langflow.services.settings.service import SettingsService
|
||||
|
||||
|
||||
class CacheServiceFactory(ServiceFactory):
|
||||
|
|
|
|||
336
src/backend/langflow/services/cache/service.py
vendored
Normal file
336
src/backend/langflow/services/cache/service.py
vendored
Normal file
|
|
@ -0,0 +1,336 @@
|
|||
import threading
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from langflow.services.base import Service
|
||||
|
||||
from langflow.services.cache.base import BaseCacheService
|
||||
|
||||
import pickle
|
||||
|
||||
from loguru import logger
|
||||
|
||||
|
||||
class InMemoryCache(BaseCacheService, Service):
|
||||
|
||||
"""
|
||||
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.RLock()
|
||||
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:
|
||||
return self._get_without_lock(key)
|
||||
|
||||
def _get_without_lock(self, key):
|
||||
"""
|
||||
Retrieve an item from the cache without acquiring the lock.
|
||||
"""
|
||||
if item := self._cache.get(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.move_to_end(key)
|
||||
# Check if the value is pickled
|
||||
if isinstance(item["value"], bytes):
|
||||
value = pickle.loads(item["value"])
|
||||
else:
|
||||
value = item["value"]
|
||||
return value
|
||||
else:
|
||||
self.delete(key)
|
||||
return None
|
||||
|
||||
def set(self, key, value, pickle=False):
|
||||
"""
|
||||
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)
|
||||
# pickle locally to mimic Redis
|
||||
if pickle:
|
||||
value = pickle.dumps(value)
|
||||
|
||||
self._cache[key] = {"value": value, "time": time.time()}
|
||||
|
||||
def upsert(self, key, value):
|
||||
"""
|
||||
Inserts or updates a value in the cache.
|
||||
If the existing value and the new value are both dictionaries, they are merged.
|
||||
|
||||
Args:
|
||||
key: The key of the item.
|
||||
value: The value to insert or update.
|
||||
"""
|
||||
with self._lock:
|
||||
existing_value = self._get_without_lock(key)
|
||||
if (
|
||||
existing_value is not None
|
||||
and isinstance(existing_value, dict)
|
||||
and isinstance(value, dict)
|
||||
):
|
||||
existing_value.update(value)
|
||||
value = existing_value
|
||||
|
||||
self.set(key, value)
|
||||
|
||||
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})"
|
||||
|
||||
|
||||
class RedisCache(BaseCacheService, Service):
|
||||
"""
|
||||
A Redis-based cache implementation.
|
||||
|
||||
This cache supports setting an expiration time for cached items.
|
||||
|
||||
Attributes:
|
||||
expiration_time (int, optional): Time in seconds after which a cached item expires. Default is 1 hour.
|
||||
|
||||
Example:
|
||||
|
||||
cache = RedisCache(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, host="localhost", port=6379, db=0, expiration_time=60 * 60):
|
||||
"""
|
||||
Initialize a new RedisCache instance.
|
||||
|
||||
Args:
|
||||
host (str, optional): Redis host.
|
||||
port (int, optional): Redis port.
|
||||
db (int, optional): Redis DB.
|
||||
expiration_time (int, optional): Time in seconds after which a
|
||||
ached item expires. Default is 1 hour.
|
||||
"""
|
||||
try:
|
||||
import redis
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"RedisCache requires the redis-py package."
|
||||
" Please install Langflow with the deploy extra: pip install langflow[deploy]"
|
||||
) from exc
|
||||
logger.warning(
|
||||
"RedisCache is an experimental feature and may not work as expected."
|
||||
" Please report any issues to our GitHub repository."
|
||||
)
|
||||
self._client = redis.StrictRedis(host=host, port=port, db=db)
|
||||
self.expiration_time = expiration_time
|
||||
|
||||
# check connection
|
||||
def is_connected(self):
|
||||
"""
|
||||
Check if the Redis client is connected.
|
||||
"""
|
||||
import redis
|
||||
|
||||
try:
|
||||
self._client.ping()
|
||||
return True
|
||||
except redis.exceptions.ConnectionError:
|
||||
return False
|
||||
|
||||
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.
|
||||
"""
|
||||
value = self._client.get(key)
|
||||
return pickle.loads(value) if value else None
|
||||
|
||||
def set(self, key, value):
|
||||
"""
|
||||
Add an item to the cache.
|
||||
|
||||
Args:
|
||||
key: The key of the item.
|
||||
value: The value to cache.
|
||||
"""
|
||||
try:
|
||||
if pickled := pickle.dumps(value):
|
||||
result = self._client.setex(key, self.expiration_time, pickled)
|
||||
if not result:
|
||||
raise ValueError("RedisCache could not set the value.")
|
||||
except TypeError as exc:
|
||||
raise TypeError(
|
||||
"RedisCache only accepts values that can be pickled. "
|
||||
) from exc
|
||||
|
||||
def upsert(self, key, value):
|
||||
"""
|
||||
Inserts or updates a value in the cache.
|
||||
If the existing value and the new value are both dictionaries, they are merged.
|
||||
|
||||
Args:
|
||||
key: The key of the item.
|
||||
value: The value to insert or update.
|
||||
"""
|
||||
existing_value = self.get(key)
|
||||
if (
|
||||
existing_value is not None
|
||||
and isinstance(existing_value, dict)
|
||||
and isinstance(value, dict)
|
||||
):
|
||||
existing_value.update(value)
|
||||
value = existing_value
|
||||
|
||||
self.set(key, value)
|
||||
|
||||
def delete(self, key):
|
||||
"""
|
||||
Remove an item from the cache.
|
||||
|
||||
Args:
|
||||
key: The key of the item to remove.
|
||||
"""
|
||||
self._client.delete(key)
|
||||
|
||||
def clear(self):
|
||||
"""
|
||||
Clear all items from the cache.
|
||||
"""
|
||||
self._client.flushdb()
|
||||
|
||||
def __contains__(self, key):
|
||||
"""Check if the key is in the cache."""
|
||||
return False if key is None else self._client.exists(key)
|
||||
|
||||
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 __repr__(self):
|
||||
"""Return a string representation of the RedisCache instance."""
|
||||
return f"RedisCache(expiration_time={self.expiration_time})"
|
||||
|
|
@ -1,4 +1,4 @@
|
|||
from langflow.services.chat.manager import ChatService
|
||||
from langflow.services.chat.service import ChatService
|
||||
from langflow.services.factory import ServiceFactory
|
||||
|
||||
|
||||
|
|
|
|||
246
src/backend/langflow/services/chat/service.py
Normal file
246
src/backend/langflow/services/chat/service.py
Normal file
|
|
@ -0,0 +1,246 @@
|
|||
from collections import defaultdict
|
||||
import uuid
|
||||
from fastapi import WebSocket, status
|
||||
from starlette.websockets import WebSocketState
|
||||
from langflow.api.v1.schemas import ChatMessage, ChatResponse, FileResponse
|
||||
from langflow.interface.utils import pil_to_base64
|
||||
from langflow.services.base import Service
|
||||
from langflow.services.chat.cache import Subject
|
||||
from langflow.services.chat.utils import process_graph
|
||||
from loguru import logger
|
||||
|
||||
from .cache import cache_service
|
||||
import asyncio
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from langflow.services import service_manager, ServiceType
|
||||
import orjson
|
||||
|
||||
|
||||
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 ChatService(Service):
|
||||
name = "chat_service"
|
||||
|
||||
def __init__(self):
|
||||
self.active_connections: Dict[str, WebSocket] = {}
|
||||
self.connection_ids: Dict[str, str] = {}
|
||||
self.chat_history = ChatHistory()
|
||||
self.chat_cache = cache_service
|
||||
self.chat_cache.attach(self.update)
|
||||
self.cache_service = service_manager.get(ServiceType.CACHE_SERVICE)
|
||||
|
||||
def on_chat_history_update(self):
|
||||
"""Send the last chat message to the client."""
|
||||
client_id = self.chat_cache.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.chat_cache.current_client_id in self.active_connections:
|
||||
self.last_cached_object_dict = self.chat_cache.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.chat_cache.current_client_id, chat_response
|
||||
)
|
||||
|
||||
async def connect(self, client_id: str, websocket: WebSocket):
|
||||
self.active_connections[client_id] = websocket
|
||||
# This is to avoid having multiple clients with the same id
|
||||
#! Temporary solution
|
||||
self.connection_ids[client_id] = f"{client_id}-{uuid.uuid4()}"
|
||||
|
||||
def disconnect(self, client_id: str):
|
||||
self.active_connections.pop(client_id, None)
|
||||
self.connection_ids.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", {})
|
||||
chatkey = payload.pop("chatKey", None)
|
||||
chat_inputs = ChatMessage(message=chat_inputs, chatKey=chatkey)
|
||||
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,
|
||||
client_id=client_id,
|
||||
session_id=self.connection_ids[client_id],
|
||||
)
|
||||
self.set_cache(client_id, langchain_object)
|
||||
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.
|
||||
"""
|
||||
# client_id is the flow id but that already exists in the cache
|
||||
# so we need to change it to something else
|
||||
|
||||
result_dict = {
|
||||
"result": langchain_object,
|
||||
"type": type(langchain_object),
|
||||
}
|
||||
self.cache_service.upsert(client_id, result_dict)
|
||||
return client_id in self.cache_service
|
||||
|
||||
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()
|
||||
if isinstance(json_payload, str):
|
||||
payload = orjson.loads(json_payload)
|
||||
elif isinstance(json_payload, dict):
|
||||
payload = json_payload
|
||||
if "clear_history" in payload and payload["clear_history"]:
|
||||
self.chat_history.history[client_id] = []
|
||||
continue
|
||||
|
||||
with self.chat_cache.set_client_id(client_id):
|
||||
if langchain_object := self.cache_service.get(client_id).get(
|
||||
"result"
|
||||
):
|
||||
await self.process_message(client_id, payload, langchain_object)
|
||||
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Could not find a build result for client_id {client_id}"
|
||||
)
|
||||
except Exception as exc:
|
||||
# Handle any exceptions that might occur
|
||||
logger.exception(f"Error handling websocket: {exc}")
|
||||
if websocket.client_state == WebSocketState.CONNECTED:
|
||||
await self.close_connection(
|
||||
client_id=client_id,
|
||||
code=status.WS_1011_INTERNAL_ERROR,
|
||||
reason=str(exc)[:120],
|
||||
)
|
||||
elif websocket.client_state == WebSocketState.DISCONNECTED:
|
||||
self.disconnect(client_id)
|
||||
|
||||
finally:
|
||||
try:
|
||||
# first check if the connection is still open
|
||||
if websocket.client_state == WebSocketState.CONNECTED:
|
||||
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)
|
||||
|
|
@ -1,9 +1,9 @@
|
|||
from typing import TYPE_CHECKING
|
||||
from langflow.services.database.manager import DatabaseService
|
||||
from langflow.services.database.service import DatabaseService
|
||||
from langflow.services.factory import ServiceFactory
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langflow.services.settings.manager import SettingsService
|
||||
from langflow.services.settings.service import SettingsService
|
||||
|
||||
|
||||
class DatabaseServiceFactory(ServiceFactory):
|
||||
|
|
|
|||
240
src/backend/langflow/services/database/service.py
Normal file
240
src/backend/langflow/services/database/service.py
Normal file
|
|
@ -0,0 +1,240 @@
|
|||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
from langflow.services.base import Service
|
||||
from langflow.services.database.models.user.crud import get_user_by_username
|
||||
from langflow.services.database.utils import Result, TableResults
|
||||
from langflow.services.getters import get_settings_service
|
||||
from sqlalchemy import inspect
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.exc import OperationalError
|
||||
from sqlmodel import SQLModel, Session, create_engine
|
||||
from loguru import logger
|
||||
from alembic.config import Config
|
||||
from alembic import command
|
||||
from langflow.services.database import models # noqa
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sqlalchemy.engine import Engine
|
||||
|
||||
|
||||
class DatabaseService(Service):
|
||||
name = "database_service"
|
||||
|
||||
def __init__(self, database_url: str):
|
||||
self.database_url = database_url
|
||||
# This file is in langflow.services.database.manager.py
|
||||
# the ini is in langflow
|
||||
langflow_dir = Path(__file__).parent.parent.parent
|
||||
self.script_location = langflow_dir / "alembic"
|
||||
self.alembic_cfg_path = langflow_dir / "alembic.ini"
|
||||
self.engine = self._create_engine()
|
||||
|
||||
def _create_engine(self) -> "Engine":
|
||||
"""Create the engine for the database."""
|
||||
settings_service = get_settings_service()
|
||||
if (
|
||||
settings_service.settings.DATABASE_URL
|
||||
and settings_service.settings.DATABASE_URL.startswith("sqlite")
|
||||
):
|
||||
connect_args = {"check_same_thread": False}
|
||||
else:
|
||||
connect_args = {}
|
||||
return create_engine(self.database_url, connect_args=connect_args)
|
||||
|
||||
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 migrate_flows_if_auto_login(self):
|
||||
# if auto_login is enabled, we need to migrate the flows
|
||||
# to the default superuser if they don't have a user id
|
||||
# associated with them
|
||||
settings_service = get_settings_service()
|
||||
if settings_service.auth_settings.AUTO_LOGIN:
|
||||
with Session(self.engine) as session:
|
||||
flows = (
|
||||
session.query(models.Flow)
|
||||
.filter(models.Flow.user_id == None) # noqa
|
||||
.all()
|
||||
)
|
||||
if flows:
|
||||
logger.debug("Migrating flows to default superuser")
|
||||
username = settings_service.auth_settings.SUPERUSER
|
||||
user = get_user_by_username(session, username)
|
||||
for flow in flows:
|
||||
flow.user_id = user.id
|
||||
session.commit()
|
||||
logger.debug("Flows migrated successfully")
|
||||
|
||||
def check_schema_health(self) -> bool:
|
||||
inspector = inspect(self.engine)
|
||||
|
||||
model_mapping = {
|
||||
"flow": models.Flow,
|
||||
"user": models.User,
|
||||
"apikey": models.ApiKey,
|
||||
# Add other SQLModel classes here
|
||||
}
|
||||
|
||||
# To account for tables that existed in older versions
|
||||
legacy_tables = ["flowstyle"]
|
||||
|
||||
for table, model in model_mapping.items():
|
||||
expected_columns = list(model.__fields__.keys())
|
||||
|
||||
try:
|
||||
available_columns = [
|
||||
col["name"] for col in inspector.get_columns(table)
|
||||
]
|
||||
except sa.exc.NoSuchTableError:
|
||||
logger.error(f"Missing table: {table}")
|
||||
return False
|
||||
|
||||
for column in expected_columns:
|
||||
if column not in available_columns:
|
||||
logger.error(f"Missing column: {column} in table {table}")
|
||||
return False
|
||||
|
||||
for table in legacy_tables:
|
||||
if table in inspector.get_table_names():
|
||||
logger.warning(f"Legacy table exists: {table}")
|
||||
|
||||
return True
|
||||
|
||||
def init_alembic(self):
|
||||
logger.info("Initializing alembic")
|
||||
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.stamp(alembic_cfg, "head")
|
||||
logger.info("Alembic initialized")
|
||||
|
||||
def run_migrations(self):
|
||||
# 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
|
||||
# We need to check if the alembic_version table exists
|
||||
# if not, we need to initialize alembic
|
||||
with Session(self.engine) as session:
|
||||
# If the table does not exist it throws an error
|
||||
# so we need to catch it
|
||||
try:
|
||||
session.execute("SELECT * FROM alembic_version")
|
||||
except Exception:
|
||||
logger.info("Alembic not initialized")
|
||||
try:
|
||||
self.init_alembic()
|
||||
except Exception as exc:
|
||||
logger.error(f"Error initializing alembic: {exc}")
|
||||
raise RuntimeError("Error initializing alembic") from exc
|
||||
else:
|
||||
logger.info("Alembic already initialized")
|
||||
|
||||
logger.info(f"Running DB migrations in {self.script_location}")
|
||||
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 run_migrations_test(self):
|
||||
# This method is used for testing purposes only
|
||||
# We will check that all models are in the database
|
||||
# and that the database is up to date with all columns
|
||||
sql_models = [models.Flow, models.User, models.ApiKey]
|
||||
return [
|
||||
TableResults(sql_model.__tablename__, self.check_table(sql_model))
|
||||
for sql_model in sql_models
|
||||
]
|
||||
|
||||
def check_table(self, model):
|
||||
results = []
|
||||
inspector = inspect(self.engine)
|
||||
table_name = model.__tablename__
|
||||
expected_columns = list(model.__fields__.keys())
|
||||
try:
|
||||
available_columns = [
|
||||
col["name"] for col in inspector.get_columns(table_name)
|
||||
]
|
||||
results.append(Result(name=table_name, type="table", success=True))
|
||||
except sa.exc.NoSuchTableError:
|
||||
logger.error(f"Missing table: {table_name}")
|
||||
results.append(Result(name=table_name, type="table", success=False))
|
||||
|
||||
for column in expected_columns:
|
||||
if column not in available_columns:
|
||||
logger.error(f"Missing column: {column} in table {table_name}")
|
||||
results.append(Result(name=column, type="column", success=False))
|
||||
else:
|
||||
results.append(Result(name=column, type="column", success=True))
|
||||
return results
|
||||
|
||||
def create_db_and_tables(self):
|
||||
from sqlalchemy import inspect
|
||||
|
||||
inspector = inspect(self.engine)
|
||||
table_names = inspector.get_table_names()
|
||||
current_tables = ["flow", "user", "apikey"]
|
||||
|
||||
if table_names and all(table in table_names for table in current_tables):
|
||||
logger.debug("Database and tables already exist")
|
||||
return
|
||||
|
||||
logger.debug("Creating database and tables")
|
||||
|
||||
for table in SQLModel.metadata.sorted_tables:
|
||||
try:
|
||||
table.create(self.engine, checkfirst=True)
|
||||
except OperationalError as oe:
|
||||
logger.warning(
|
||||
f"Table {table} already exists, skipping. Exception: {oe}"
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error(f"Error creating table {table}: {exc}")
|
||||
raise RuntimeError(f"Error creating table {table}") from exc
|
||||
|
||||
# Now check if the required tables exist, if not, something went wrong.
|
||||
inspector = inspect(self.engine)
|
||||
table_names = inspector.get_table_names()
|
||||
for table in current_tables:
|
||||
if table not in 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."
|
||||
)
|
||||
|
||||
logger.debug("Database and tables created successfully")
|
||||
|
||||
def teardown(self):
|
||||
logger.debug("Tearing down database")
|
||||
try:
|
||||
settings_service = get_settings_service()
|
||||
# remove the default superuser if auto_login is enabled
|
||||
# using the SUPERUSER to get the user
|
||||
if settings_service.auth_settings.AUTO_LOGIN:
|
||||
logger.debug("Removing default superuser")
|
||||
username = settings_service.auth_settings.SUPERUSER
|
||||
with Session(self.engine) as session:
|
||||
user = get_user_by_username(session, username)
|
||||
session.delete(user)
|
||||
session.commit()
|
||||
logger.debug("Default superuser removed")
|
||||
|
||||
except Exception as exc:
|
||||
logger.error(f"Error tearing down database: {exc}")
|
||||
|
||||
self.engine.dispose()
|
||||
|
|
@ -6,7 +6,7 @@ from alembic.util.exc import CommandError
|
|||
from sqlmodel import Session
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langflow.services.database.manager import DatabaseService
|
||||
from langflow.services.database.service import DatabaseService
|
||||
|
||||
|
||||
def initialize_database():
|
||||
|
|
|
|||
|
|
@ -3,12 +3,12 @@ from typing import TYPE_CHECKING, Generator
|
|||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langflow.services.database.manager import DatabaseService
|
||||
from langflow.services.settings.manager import SettingsService
|
||||
from langflow.services.cache.manager import BaseCacheService
|
||||
from langflow.services.session.manager import SessionService
|
||||
from langflow.services.task.manager import TaskService
|
||||
from langflow.services.chat.manager import ChatService
|
||||
from langflow.services.database.service import DatabaseService
|
||||
from langflow.services.settings.service import SettingsService
|
||||
from langflow.services.cache.service import BaseCacheService
|
||||
from langflow.services.session.service import SessionService
|
||||
from langflow.services.task.service import TaskService
|
||||
from langflow.services.chat.service import ChatService
|
||||
from sqlmodel import Session
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,9 +1,9 @@
|
|||
from typing import TYPE_CHECKING
|
||||
from langflow.services.session.manager import SessionService
|
||||
from langflow.services.session.service import SessionService
|
||||
from langflow.services.factory import ServiceFactory
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langflow.services.cache.manager import BaseCacheService
|
||||
from langflow.services.cache.service import BaseCacheService
|
||||
|
||||
|
||||
class SessionServiceFactory(ServiceFactory):
|
||||
|
|
|
|||
47
src/backend/langflow/services/session/service.py
Normal file
47
src/backend/langflow/services/session/service.py
Normal file
|
|
@ -0,0 +1,47 @@
|
|||
from typing import TYPE_CHECKING
|
||||
from langflow.interface.run import build_sorted_vertices
|
||||
from langflow.services.base import Service
|
||||
from langflow.services.cache.utils import compute_dict_hash
|
||||
from langflow.services.session.utils import session_id_generator
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langflow.services.cache.base import BaseCacheService
|
||||
|
||||
|
||||
class SessionService(Service):
|
||||
name = "session_service"
|
||||
|
||||
def __init__(self, cache_service):
|
||||
self.cache_service: "BaseCacheService" = cache_service
|
||||
|
||||
def load_session(self, key, data_graph):
|
||||
# Check if the data is cached
|
||||
if key in self.cache_service:
|
||||
return self.cache_service.get(key)
|
||||
|
||||
if key is None:
|
||||
key = self.generate_key(session_id=None, data_graph=data_graph)
|
||||
|
||||
# If not cached, build the graph and cache it
|
||||
graph, artifacts = build_sorted_vertices(data_graph)
|
||||
|
||||
self.cache_service.set(key, (graph, artifacts))
|
||||
|
||||
return graph, artifacts
|
||||
|
||||
def build_key(self, session_id, data_graph):
|
||||
json_hash = compute_dict_hash(data_graph)
|
||||
return f"{session_id}{':' if session_id else ''}{json_hash}"
|
||||
|
||||
def generate_key(self, session_id, data_graph):
|
||||
# Hash the JSON and combine it with the session_id to create a unique key
|
||||
if session_id is None:
|
||||
# generate a 5 char session_id to concatenate with the json_hash
|
||||
session_id = session_id_generator()
|
||||
return self.build_key(session_id, data_graph=data_graph)
|
||||
|
||||
def update_session(self, session_id, value):
|
||||
self.cache_service.set(session_id, value)
|
||||
|
||||
def clear_session(self, session_id):
|
||||
self.cache_service.delete(session_id)
|
||||
44
src/backend/langflow/services/settings/service.py
Normal file
44
src/backend/langflow/services/settings/service.py
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
from langflow.services.base import Service
|
||||
from langflow.services.settings.auth import AuthSettings
|
||||
from langflow.services.settings.base import Settings
|
||||
from loguru import logger
|
||||
import os
|
||||
import yaml
|
||||
|
||||
|
||||
class SettingsService(Service):
|
||||
name = "settings_service"
|
||||
|
||||
def __init__(self, settings: Settings, auth_settings: AuthSettings):
|
||||
super().__init__()
|
||||
self.settings = settings
|
||||
self.auth_settings = auth_settings
|
||||
|
||||
@classmethod
|
||||
def load_settings_from_yaml(cls, file_path: str) -> "SettingsService":
|
||||
# Check if a string is a valid path or a file name
|
||||
if "/" not in file_path:
|
||||
# Get current path
|
||||
current_path = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
file_path = os.path.join(current_path, file_path)
|
||||
|
||||
with open(file_path, "r") as f:
|
||||
settings_dict = yaml.safe_load(f)
|
||||
settings_dict = {k.upper(): v for k, v in settings_dict.items()}
|
||||
|
||||
for key in settings_dict:
|
||||
if key not in Settings.__fields__.keys():
|
||||
raise KeyError(f"Key {key} not found in settings")
|
||||
logger.debug(
|
||||
f"Loading {len(settings_dict[key])} {key} from {file_path}"
|
||||
)
|
||||
|
||||
settings = Settings(**settings_dict)
|
||||
if not settings.CONFIG_DIR:
|
||||
raise ValueError("CONFIG_DIR must be set in settings")
|
||||
|
||||
auth_settings = AuthSettings(
|
||||
CONFIG_DIR=settings.CONFIG_DIR,
|
||||
)
|
||||
return cls(settings, auth_settings)
|
||||
75
src/backend/langflow/services/task/service.py
Normal file
75
src/backend/langflow/services/task/service.py
Normal file
|
|
@ -0,0 +1,75 @@
|
|||
from typing import Any, Callable, Coroutine, Union
|
||||
from langflow.utils.logger import configure
|
||||
from loguru import logger
|
||||
from langflow.services.base import Service
|
||||
from langflow.services.task.backends.anyio import AnyIOBackend
|
||||
from langflow.services.task.backends.base import TaskBackend
|
||||
from langflow.services.task.utils import get_celery_worker_status
|
||||
|
||||
|
||||
def check_celery_availability():
|
||||
try:
|
||||
from langflow.worker import celery_app
|
||||
|
||||
status = get_celery_worker_status(celery_app)
|
||||
logger.debug(f"Celery status: {status}")
|
||||
except Exception as exc:
|
||||
logger.debug(f"Celery not available: {exc}")
|
||||
status = {"availability": None}
|
||||
return status
|
||||
|
||||
|
||||
try:
|
||||
configure()
|
||||
status = check_celery_availability()
|
||||
|
||||
USE_CELERY = status.get("availability") is not None
|
||||
except ImportError:
|
||||
USE_CELERY = False
|
||||
|
||||
|
||||
class TaskService(Service):
|
||||
name = "task_service"
|
||||
|
||||
def __init__(self):
|
||||
self.backend = self.get_backend()
|
||||
self.use_celery = USE_CELERY
|
||||
|
||||
@property
|
||||
def backend_name(self) -> str:
|
||||
return self.backend.name
|
||||
|
||||
def get_backend(self) -> TaskBackend:
|
||||
if USE_CELERY:
|
||||
from langflow.services.task.backends.celery import CeleryBackend
|
||||
|
||||
logger.debug("Using Celery backend")
|
||||
return CeleryBackend()
|
||||
logger.debug("Using AnyIO backend")
|
||||
return AnyIOBackend()
|
||||
|
||||
# In your TaskService class
|
||||
async def launch_and_await_task(
|
||||
self,
|
||||
task_func: Callable[..., Any],
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
if not self.use_celery:
|
||||
return None, await task_func(*args, **kwargs)
|
||||
if not hasattr(task_func, "apply"):
|
||||
raise ValueError(f"Task function {task_func} does not have an apply method")
|
||||
task = task_func.apply(args=args, kwargs=kwargs)
|
||||
result = task.get()
|
||||
return task.id, result
|
||||
|
||||
async def launch_task(
|
||||
self, task_func: Callable[..., Any], *args: Any, **kwargs: Any
|
||||
) -> Any:
|
||||
logger.debug(f"Launching task {task_func} with args {args} and kwargs {kwargs}")
|
||||
logger.debug(f"Using backend {self.backend}")
|
||||
task = self.backend.launch_task(task_func, *args, **kwargs)
|
||||
return await task if isinstance(task, Coroutine) else task
|
||||
|
||||
def get_task(self, task_id: Union[int, str]) -> Any:
|
||||
return self.backend.get_task(task_id)
|
||||
Loading…
Reference in a new issue