Adding API improvements, HomePage and validation improvements (#493)

This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-06-16 19:17:13 -03:00 • committed by GitHub
commit 0f59382704
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
132 changed files with 16331 additions and 11993 deletions

1
.gitignore vendored
View file

@ -241,3 +241,4 @@ dmypy.json
# Poetry
.testenv/*
langflow.db

2
.vscode/launch.json vendored
View file

@ -35,7 +35,7 @@
"name": "Debug Frontend",
"type": "chrome",
"request": "launch",
"url": "http://localhost:3000/*",
"url": "http://localhost:3000/",
"webRoot": "${workspaceRoot}/src/frontend"
}
]

629
poetry.lock generated

File diff suppressed because it is too large Load diff

View file

@ -29,10 +29,9 @@ google-search-results = "^2.4.1"
google-api-python-client = "^2.79.0"
typer = "^0.7.0"
gunicorn = "^20.1.0"
langchain = "^0.0.194"
openai = "^0.27.7"
langchain = "^0.0.200"
openai = "^0.27.8"
types-pyyaml = "^6.0.12.8"
dill = "^0.3.6"
pandas = "^1.5.3"
chromadb = "^0.3.21"
huggingface-hub = "^0.13.3"
@ -57,9 +56,13 @@ jina = "3.15.2"
sentence-transformers = "^2.2.2"
ctransformers = "^0.2.2"
cohere = "^4.6.0"
sqlmodel = "^0.0.8"
faiss-cpu = "^1.7.4"
anthropic = "^0.2.9"
orjson = "^3.9.0"
multiprocess = "^0.70.14"
cachetools = "^5.3.1"
types-cachetools = "^5.3.0.5"
[tool.poetry.group.dev.dependencies]

View file

@ -1,16 +1,18 @@
import sys
import time
from fastapi import FastAPI
import httpx
from multiprocess import Process, cpu_count # type: ignore
import platform
from pathlib import Path
from typing import Optional
import socket
from rich.panel import Panel
from rich import box
from rich import print as rprint
import typer
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from langflow.main import create_app
from langflow.settings import settings
from langflow.utils.logger import configure, logger
@ -26,10 +28,19 @@ def get_number_of_workers(workers=None):
return workers
def update_settings(config: str, dev: bool = False):
def update_settings(
config: str,
dev: bool = False,
database_url: Optional[str] = None,
remove_api_keys: bool = False,
):
"""Update the settings from a config file."""
if config:
settings.update_from_yaml(config, dev=dev)
if database_url:
settings.update_settings(database_url=database_url)
if remove_api_keys:
settings.update_settings(remove_api_keys=remove_api_keys)
def serve_on_jcloud():
@ -93,6 +104,10 @@ def serve(
log_file: Path = typer.Option("logs/langflow.log", help="Path to the log file."),
jcloud: bool = typer.Option(False, help="Deploy on Jina AI Cloud"),
dev: bool = typer.Option(False, help="Run in development mode (may contain bugs)"),
database_url: str = typer.Option(
None,
help="Database URL to connect to. If not provided, a local SQLite database will be used.",
),
path: str = typer.Option(
None,
help="Path to the frontend directory containing build files. This is for development purposes only.",
@ -100,23 +115,12 @@ def serve(
open_browser: bool = typer.Option(
True, help="Open the browser after starting the server."
),
remove_api_keys: bool = typer.Option(
False, help="Remove API keys from the projects saved in the database."
),
):
"""
Run the Langflow server.
Args:
host (str): Host to bind the server to.
workers (int): Number of worker processes.
timeout (int): Worker timeout in seconds.
port (int): Port to listen on.
config (str): Path to the configuration file.
env_file (Path): Path to the .env file containing environment variables.
log_level (str): Logging level.
log_file (Path): Path to the log file.
jcloud (bool): Deploy on Jina AI Cloud.
dev (bool): Run in development mode (may contain bugs).
path (str): Path to the frontend directory containing build files. This is for development purposes only.
open_browser (bool): Open the browser after starting the server.
"""
if jcloud:
@ -125,19 +129,22 @@ def serve(
load_dotenv(env_file)
configure(log_level=log_level, log_file=log_file)
update_settings(config, dev=dev)
app = create_app()
update_settings(
config, dev=dev, database_url=database_url, remove_api_keys=remove_api_keys
)
# get the directory of the current file
if not path:
frontend_path = Path(__file__).parent
static_files_dir = frontend_path / "frontend"
else:
static_files_dir = Path(path)
app.mount(
"/",
StaticFiles(directory=static_files_dir, html=True),
name="static",
)
app = create_app()
setup_static_files(app, static_files_dir)
# check if port is being used
if is_port_in_use(port, host):
port = get_free_port(port)
options = {
"bind": f"{host}:{port}",
"workers": get_number_of_workers(workers),
@ -152,7 +159,8 @@ def serve(
status_code = 0
while status_code != 200:
try:
status_code = httpx.get(f"http://{host}:{port}").status_code
status_code = httpx.get(f"http://{host}:{port}/health").status_code
except Exception:
time.sleep(1)
@ -161,11 +169,64 @@ def serve(
webbrowser.open(f"http://{host}:{port}")
def setup_static_files(app: FastAPI, static_files_dir: Path):
"""
Setup the static files directory.
Args:
app (FastAPI): FastAPI app.
path (str): Path to the static files directory.
"""
app.mount(
"/",
StaticFiles(directory=static_files_dir, html=True),
name="static",
)
@app.exception_handler(404)
async def custom_404_handler(request, __):
path = static_files_dir / "index.html"
if not path.exists():
raise RuntimeError(f"File at path {path} does not exist.")
return FileResponse(path)
def is_port_in_use(port, host="localhost"):
"""
Check if a port is in use.
Args:
port (int): The port number to check.
host (str): The host to check the port on. Defaults to 'localhost'.
Returns:
bool: True if the port is in use, False otherwise.
"""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
return s.connect_ex((host, port)) == 0
def get_free_port(port):
"""
Given a used port, find a free port.
Args:
port (int): The port number to check.
Returns:
int: A free port number.
"""
while is_port_in_use(port):
port += 1
return port
def print_banner(host, port):
# console = Console()
word = "LangFlow"
colors = ["#690080", "#660099", "#4d00b3", "#3300cc", "#1a00e6", "#0000ff"]
colors = ["#3300cc"]
styled_word = ""

View file

@ -1,8 +1,18 @@
# Router for base api
from fastapi import APIRouter
from langflow.api.v1 import chat_router, endpoints_router, validate_router
from langflow.api.v1 import (
chat_router,
endpoints_router,
validate_router,
flows_router,
flow_styles_router,
)
router = APIRouter(prefix="/api/v1", tags=["api"])
router = APIRouter(
prefix="/api/v1",
)
router.include_router(chat_router)
router.include_router(endpoints_router)
router.include_router(validate_router)
router.include_router(flows_router)
router.include_router(flow_styles_router)

View file

@ -0,0 +1,24 @@
API_WORDS = ["api", "key", "token"]
def has_api_terms(word: str):
return "api" in word and (
"key" in word or ("token" in word and "tokens" not in word)
)
def remove_api_keys(flow: dict):
"""Remove api keys from flow data."""
if flow.get("data") and flow["data"].get("nodes"):
for node in flow["data"]["nodes"]:
node_data = node.get("data").get("node")
template = node_data.get("template")
for value in template.values():
if (
isinstance(value, dict)
and has_api_terms(value["name"])
and value.get("password")
):
value["value"] = None
return flow

View file

@ -1,5 +1,13 @@
from langflow.api.v1.endpoints import router as endpoints_router
from langflow.api.v1.validate import router as validate_router
from langflow.api.v1.chat import router as chat_router
from langflow.api.v1.flows import router as flows_router
from langflow.api.v1.flow_styles import router as flow_styles_router
__all__ = ["chat_router", "endpoints_router", "validate_router"]
__all__ = [
"chat_router",
"endpoints_router",
"validate_router",
"flows_router",
"flow_styles_router",
]

View file

@ -1,26 +1,122 @@
import json
from fastapi import (
APIRouter,
HTTPException,
WebSocket,
WebSocketDisconnect,
WebSocketException,
status,
)
from fastapi.responses import StreamingResponse
from langflow.api.v1.schemas import BuiltResponse, InitResponse
from langflow.chat.manager import ChatManager
from langflow.graph.graph.base import Graph
from langflow.utils.logger import logger
from cachetools import LRUCache
router = APIRouter()
router = APIRouter(tags=["Chat"])
chat_manager = ChatManager()
flow_data_store: LRUCache = LRUCache(maxsize=10)
@router.websocket("/chat/{client_id}")
async def websocket_endpoint(client_id: str, websocket: WebSocket):
async def chat(client_id: str, websocket: WebSocket):
"""Websocket endpoint for chat."""
try:
await chat_manager.handle_websocket(client_id, websocket)
if client_id in chat_manager.in_memory_cache:
await chat_manager.handle_websocket(client_id, websocket)
else:
message = "Please, build the flow before sending messages"
await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason=message)
except WebSocketException as exc:
logger.error(exc)
await websocket.close(code=status.WS_1011_INTERNAL_ERROR, reason=str(exc))
except WebSocketDisconnect as exc:
@router.post("/build/init", response_model=InitResponse, status_code=201)
async def init_build(graph_data: dict):
"""Initialize the build by storing graph data and returning a unique session ID."""
try:
flow_id = graph_data.get("id")
if flow_id is None:
raise ValueError("No ID provided")
flow_data_store[flow_id] = graph_data
return InitResponse(flowId=flow_id)
except Exception as exc:
logger.error(exc)
await websocket.close(code=status.WS_1000_NORMAL_CLOSURE, reason=str(exc))
return HTTPException(status_code=500, detail=str(exc))
@router.get("/build/{flow_id}/status", response_model=BuiltResponse)
async def build_status(flow_id: str):
"""Check the flow_id is in the flow_data_store."""
try:
built = flow_id in flow_data_store and not isinstance(
flow_data_store[flow_id], dict
)
return BuiltResponse(
built=built,
)
except Exception as exc:
logger.error(exc)
return HTTPException(status_code=500, detail=str(exc))
@router.get("/build/stream/{flow_id}", response_class=StreamingResponse)
async def stream_build(flow_id: str):
"""Stream the build process based on stored flow data."""
async def event_stream(flow_id):
final_response = json.dumps({"end_of_stream": True})
try:
if flow_id not in flow_data_store:
error_message = "Invalid session ID"
yield f"data: {json.dumps({'error': error_message})}\n\n"
return
graph_data = flow_data_store[flow_id].get("data")
if not graph_data:
error_message = "No data provided"
yield f"data: {json.dumps({'error': error_message})}\n\n"
return
logger.debug("Building langchain object")
graph = Graph.from_payload(graph_data)
for node in graph.generator_build():
try:
node.build()
params = node._built_object_repr()
valid = True
logger.debug(
f"Building node {params[:50]}{'...' if len(params) > 50 else ''}"
)
except Exception as exc:
params = str(exc)
valid = False
response = json.dumps(
{
"valid": valid,
"params": params,
"id": node.id,
}
)
yield f"data: {response}\n\n"
chat_manager.set_cache(flow_id, graph.build())
except Exception as exc:
logger.error("Error while building the flow: %s", exc)
yield f"error: {json.dumps({'error': str(exc)})}\n\n"
finally:
yield f"data: {final_response}\n\n"
try:
return StreamingResponse(event_stream(flow_id), media_type="text/event-stream")
except Exception as exc:
logger.error(exc)
raise HTTPException(status_code=500, detail=str(exc))

View file

@ -1,20 +1,20 @@
import logging
from importlib.metadata import version
from langflow.database.models.flow import Flow
from langflow.processing.process import process_graph_cached, process_tweaks
from langflow.utils.logger import logger
from fastapi import APIRouter, HTTPException
from fastapi import APIRouter, Depends, HTTPException
from langflow.api.v1.schemas import (
ExportedFlow,
GraphData,
PredictRequest,
PredictResponse,
)
from langflow.interface.types import build_langchain_types_dict
from langflow.database.base import get_session
from sqlmodel import Session
# build router
router = APIRouter()
logger = logging.getLogger(__name__)
router = APIRouter(tags=["Base"])
@router.get("/all")
@ -22,16 +22,34 @@ def get_all():
return build_langchain_types_dict()
@router.post("/predict", response_model=PredictResponse)
async def get_load(predict_request: PredictRequest):
try:
from langflow.processing.process import process_graph_cached
@router.post("/predict/{flow_id}", response_model=PredictResponse)
async def predict_flow(
predict_request: PredictRequest,
flow_id: str,
session: Session = Depends(get_session),
):
"""
Endpoint to process a message using the flow passed in the bearer token.
"""
exported_flow: ExportedFlow = predict_request.exported_flow
graph_data: GraphData = exported_flow.data
data = graph_data.dict()
response = process_graph_cached(data, predict_request.message)
return PredictResponse(result=response.get("result", ""))
try:
flow = session.get(Flow, flow_id)
if flow is None:
raise ValueError(f"Flow {flow_id} not found")
if flow.data is None:
raise ValueError(f"Flow {flow_id} has no data")
graph_data = flow.data
if predict_request.tweaks:
try:
graph_data = process_tweaks(graph_data, predict_request.tweaks)
except Exception as exc:
logger.error(f"Error processing tweaks: {exc}")
response = process_graph_cached(graph_data, predict_request.message)
return PredictResponse(
result=response.get("result", ""),
intermediate_steps=response.get("thought", ""),
)
except Exception as e:
# Log stack trace
logger.exception(e)

View file

@ -0,0 +1,83 @@
from uuid import UUID
from langflow.database.models.flow_style import (
FlowStyle,
FlowStyleCreate,
FlowStyleRead,
FlowStyleUpdate,
)
from langflow.database.base import get_session
from sqlmodel import Session, select
from fastapi import APIRouter, Depends, HTTPException
# build router
router = APIRouter(prefix="/flow_styles", tags=["FlowStyles"])
# FlowStyleCreate:
# class FlowStyleBase(SQLModel):
# color: str = Field(index=True)
# emoji: str = Field(index=False)
# flow_id: UUID = Field(default=None, foreign_key="flow.id")
@router.post("/", response_model=FlowStyleRead)
def create_flow_style(
*, session: Session = Depends(get_session), flow_style: FlowStyleCreate
):
"""Create a new flow_style."""
db_flow_style = FlowStyle.from_orm(flow_style)
session.add(db_flow_style)
session.commit()
session.refresh(db_flow_style)
return db_flow_style
@router.get("/", response_model=list[FlowStyleRead])
def read_flow_styles(*, session: Session = Depends(get_session)):
"""Read all flows."""
try:
flows = session.exec(select(FlowStyle)).all()
except Exception as e:
raise HTTPException(status_code=500, detail=str(e)) from e
return flows
@router.get("/{flow_styles_id}", response_model=FlowStyleRead)
def read_flow_style(*, session: Session = Depends(get_session), flow_styles_id: UUID):
"""Read a flow_style."""
if flow_style := session.get(FlowStyle, flow_styles_id):
return flow_style
else:
raise HTTPException(status_code=404, detail="FlowStyle not found")
@router.patch("/{flow_style_id}", response_model=FlowStyleRead)
def update_flow_style(
*,
session: Session = Depends(get_session),
flow_style_id: UUID,
flow_style: FlowStyleUpdate,
):
"""Update a flow_style."""
db_flow_style = session.get(FlowStyle, flow_style_id)
if not db_flow_style:
raise HTTPException(status_code=404, detail="FlowStyle not found")
flow_data = flow_style.dict(exclude_unset=True)
for key, value in flow_data.items():
if hasattr(db_flow_style, key) and value is not None:
setattr(db_flow_style, key, value)
session.add(db_flow_style)
session.commit()
session.refresh(db_flow_style)
return db_flow_style
@router.delete("/{flow_id}")
def delete_flow_style(*, session: Session = Depends(get_session), flow_id: UUID):
"""Delete a flow_style."""
flow_style = session.get(FlowStyle, flow_id)
if not flow_style:
raise HTTPException(status_code=404, detail="FlowStyle not found")
session.delete(flow_style)
session.commit()
return {"message": "FlowStyle deleted successfully"}

View file

@ -0,0 +1,120 @@
from typing import List
from uuid import UUID
from langflow.settings import settings
from langflow.api.utils import remove_api_keys
from langflow.api.v1.schemas import FlowListCreate, FlowListRead
from langflow.database.models.flow import (
Flow,
FlowCreate,
FlowRead,
FlowReadWithStyle,
FlowUpdate,
)
from langflow.database.base import get_session
from sqlmodel import Session, select
from fastapi import APIRouter, Depends, HTTPException
from fastapi.encoders import jsonable_encoder
from fastapi import File, UploadFile
import json
# build router
router = APIRouter(prefix="/flows", tags=["Flows"])
@router.post("/", response_model=FlowRead, status_code=201)
def create_flow(*, session: Session = Depends(get_session), flow: FlowCreate):
"""Create a new flow."""
db_flow = Flow.from_orm(flow)
session.add(db_flow)
session.commit()
session.refresh(db_flow)
return db_flow
@router.get("/", response_model=list[FlowReadWithStyle], status_code=200)
def read_flows(*, session: Session = Depends(get_session)):
"""Read all flows."""
try:
flows = session.exec(select(Flow)).all()
except Exception as e:
raise HTTPException(status_code=500, detail=str(e)) from e
return [jsonable_encoder(flow) for flow in flows]
@router.get("/{flow_id}", response_model=FlowReadWithStyle, status_code=200)
def read_flow(*, session: Session = Depends(get_session), flow_id: UUID):
"""Read a flow."""
if flow := session.get(Flow, flow_id):
return flow
else:
raise HTTPException(status_code=404, detail="Flow not found")
@router.patch("/{flow_id}", response_model=FlowRead, status_code=200)
def update_flow(
*, session: Session = Depends(get_session), flow_id: UUID, flow: FlowUpdate
):
"""Update a flow."""
db_flow = session.get(Flow, flow_id)
if not db_flow:
raise HTTPException(status_code=404, detail="Flow not found")
flow_data = flow.dict(exclude_unset=True)
if not settings.remove_api_keys:
flow_data = remove_api_keys(flow_data)
for key, value in flow_data.items():
setattr(db_flow, key, value)
session.add(db_flow)
session.commit()
session.refresh(db_flow)
return db_flow
@router.delete("/{flow_id}", status_code=200)
def delete_flow(*, session: Session = Depends(get_session), flow_id: UUID):
"""Delete a flow."""
flow = session.get(Flow, flow_id)
if not flow:
raise HTTPException(status_code=404, detail="Flow not found")
session.delete(flow)
session.commit()
return {"message": "Flow deleted successfully"}
# Define a new model to handle multiple flows
@router.post("/batch/", response_model=List[FlowRead], status_code=201)
def create_flows(*, session: Session = Depends(get_session), flow_list: FlowListCreate):
"""Create multiple new flows."""
db_flows = []
for flow in flow_list.flows:
db_flow = Flow.from_orm(flow)
session.add(db_flow)
db_flows.append(db_flow)
session.commit()
for db_flow in db_flows:
session.refresh(db_flow)
return db_flows
@router.post("/upload/", response_model=List[FlowRead], status_code=201)
async def upload_file(
*, session: Session = Depends(get_session), file: UploadFile = File(...)
):
"""Upload flows from a file."""
contents = await file.read()
data = json.loads(contents)
if "flows" in data:
flow_list = FlowListCreate(**data)
else:
flow_list = FlowListCreate(flows=[FlowCreate(**flow) for flow in data])
return create_flows(session=session, flow_list=flow_list)
@router.get("/download/", response_model=FlowListRead, status_code=200)
async def download_file(*, session: Session = Depends(get_session)):
"""Download all flows as a file."""
flows = read_flows(session=session)
return FlowListRead(flows=flows)

View file

@ -1,6 +1,6 @@
from typing import Any, Dict, List, Union
from pydantic import BaseModel, validator
from typing import Any, Dict, List, Optional, Union
from langflow.database.models.flow import FlowCreate, FlowRead
from pydantic import BaseModel, Field, validator
class GraphData(BaseModel):
@ -23,13 +23,30 @@ class PredictRequest(BaseModel):
"""Predict request schema."""
message: str
exported_flow: ExportedFlow
tweaks: Optional[Dict[str, Dict[str, str]]] = Field(default_factory=dict)
class Config:
schema_extra = {
"example": {
"message": "Hello, how are you?",
"tweaks": {
"dndnode_986363f0-4677-4035-9f38-74b94af5dd78": {
"name": "A tool name",
"description": "A tool description",
},
"dndnode_986363f0-4677-4035-9f38-74b94af57378": {
"template": "A {template}",
},
},
}
}
class PredictResponse(BaseModel):
"""Predict response schema."""
result: str
intermediate_steps: str = ""
class ChatMessage(BaseModel):
@ -68,3 +85,19 @@ class FileResponse(ChatMessage):
if v not in ["image", "csv"]:
raise ValueError("data_type must be image or csv")
return v
class FlowListCreate(BaseModel):
flows: List[FlowCreate]
class FlowListRead(BaseModel):
flows: List[FlowRead]
class InitResponse(BaseModel):
flowId: str
class BuiltResponse(BaseModel):
built: bool

View file

@ -15,7 +15,7 @@ from langflow.utils.logger import logger
from langflow.utils.validate import validate_code
# build router
router = APIRouter(prefix="/validate", tags=["validate"])
router = APIRouter(prefix="/validate", tags=["Validate"])
@router.post("/code", status_code=200, response_model=CodeValidationResponse)

View file

@ -1 +1,7 @@
from langflow.cache.manager import cache_manager # noqa
from langflow.cache.manager import cache_manager
from langflow.cache.flow import InMemoryCache
__all__ = [
"cache_manager",
"InMemoryCache",
]

View file

@ -1,154 +1,87 @@
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
import dill # type: ignore
CACHE: Dict[str, Any] = {}
import abc
def create_cache_folder(func):
def wrapper(*args, **kwargs):
# Get the destination folder
cache_path = Path(tempfile.gettempdir()) / 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:
class BaseCache(abc.ABC):
"""
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.
Abstract base class for a cache.
"""
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(tempfile.gettempdir()) / PREFIX
if not content:
raise ValueError("Please, reload the file in the loader.")
data = content.split(",")[1]
decoded_bytes = base64.b64decode(data)
@abc.abstractmethod
def get(self, key):
"""
Retrieve an item from the cache.
# Create the full file path
file_path = os.path.join(cache_path, file_name)
Args:
key: The key of the item to retrieve.
# Save the binary content to the file
with open(file_path, "wb") as file:
file.write(decoded_bytes)
Returns:
The value associated with the key, or None if the key is not found.
"""
return file_path
@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.
"""
@create_cache_folder
def save_cache(hash_val: str, chat_data, clean_old_cache_files: bool):
cache_path = Path(tempfile.gettempdir()) / PREFIX / f"{hash_val}.dill"
with cache_path.open("wb") as cache_file:
dill.dump(chat_data, cache_file)
@abc.abstractmethod
def delete(self, key):
"""
Remove an item from the cache.
if clean_old_cache_files:
clear_old_cache_files()
Args:
key: The key of the item to remove.
"""
@abc.abstractmethod
def clear(self):
"""
Clear all items from the cache.
"""
@create_cache_folder
def load_cache(hash_val):
cache_path = Path(tempfile.gettempdir()) / PREFIX / f"{hash_val}.dill"
if cache_path.exists():
with cache_path.open("rb") as cache_file:
return dill.load(cache_file)
return None
@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.
Returns:
The value associated with the key, or None if the key is not found.
"""
@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.
"""

146
src/backend/langflow/cache/flow.py vendored Normal file
View file

@ -0,0 +1,146 @@
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})"

View file

@ -54,7 +54,7 @@ class CacheManager(Subject):
def __init__(self):
super().__init__()
self.CACHE = {}
self._cache = {}
self.current_client_id = None
self.current_cache = {}
@ -68,12 +68,12 @@ class CacheManager(Subject):
"""
previous_client_id = self.current_client_id
self.current_client_id = client_id
self.current_cache = self.CACHE.setdefault(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, {})
self.current_cache = self._cache.get(self.current_client_id, {})
def add(self, name: str, obj: Any, obj_type: str, extension: Optional[str] = None):
"""

134
src/backend/langflow/cache/utils.py vendored Normal file
View file

@ -0,0 +1,134 @@
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
CACHE: Dict[str, Any] = {}
def create_cache_folder(func):
def wrapper(*args, **kwargs):
# Get the destination folder
cache_path = Path(tempfile.gettempdir()) / 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(tempfile.gettempdir()) / 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

View file

@ -10,7 +10,9 @@ from langflow.utils.logger import logger
import asyncio
import json
from typing import Dict, List
from typing import Any, Dict, List
from langflow.cache.flow import InMemoryCache
class ChatHistory(Subject):
@ -46,6 +48,7 @@ class ChatManager:
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."""
@ -99,24 +102,30 @@ class ChatManager:
websocket = self.active_connections[client_id]
await websocket.send_json(message.dict())
async def process_message(self, client_id: str, payload: Dict):
async def close_connection(self, client_id: str, code: int, reason: str):
if websocket := self.active_connections[client_id]:
await websocket.close(code=code, reason=reason)
self.disconnect(client_id)
async def process_message(
self, client_id: str, payload: Dict, langchain_object: Any
):
# Process the graph data and chat message
chat_message = payload.pop("message", "")
chat_message = ChatMessage(message=chat_message)
self.chat_history.add_message(client_id, chat_message)
graph_data = payload
# 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
# 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(
graph_data=graph_data,
is_first_message=is_first_message,
langchain_object=langchain_object,
chat_message=chat_message,
websocket=self.active_connections[client_id],
)
@ -149,6 +158,14 @@ class ChatManager:
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)
@ -169,22 +186,24 @@ class ChatManager:
continue
with self.cache_manager.set_client_id(client_id):
await self.process_message(client_id, payload)
langchain_object = self.in_memory_cache.get(client_id)
await self.process_message(client_id, payload, langchain_object)
except Exception as e:
# Handle any exceptions that might occur
logger.exception(e)
# send a message to the client
await self.active_connections[client_id].close(
code=status.WS_1011_INTERNAL_ERROR, reason=str(e)[:120]
logger.error(e)
await self.close_connection(
client_id=client_id,
code=status.WS_1011_INTERNAL_ERROR,
reason=str(e)[:120],
)
self.disconnect(client_id)
finally:
try:
connection = self.active_connections.get(client_id)
if connection:
await connection.close(code=1000, reason="Client disconnected")
self.disconnect(client_id)
await self.close_connection(
client_id=client_id,
code=status.WS_1000_NORMAL_CLOSURE,
reason="Client disconnected",
)
except Exception as e:
logger.exception(e)
logger.error(e)
self.disconnect(client_id)

Some files were not shown because too many files have changed in this diff Show more