diff --git a/.dockerignore b/.dockerignore index e28863bf6..d7e716758 100644 --- a/.dockerignore +++ b/.dockerignore @@ -10,7 +10,8 @@ node_modules vite.config.js.timestamp-* vite.config.ts.timestamp-* __pycache__ -.env +.idea +venv _old uploads .ipynb_checkpoints diff --git a/.github/workflows/integration-test.yml b/.github/workflows/integration-test.yml index 85810c2ed..639ea789f 100644 --- a/.github/workflows/integration-test.yml +++ b/.github/workflows/integration-test.yml @@ -35,6 +35,10 @@ jobs: done echo "Service is up!" + - name: Delete Docker build cache + run: | + docker builder prune --all --force + - name: Preload Ollama model run: | docker exec ollama ollama pull qwen:0.5b-chat-v1.5-q2_K @@ -43,7 +47,7 @@ jobs: uses: cypress-io/github-action@v6 with: browser: chrome - wait-on: "http://localhost:3000" + wait-on: 'http://localhost:3000' config: baseUrl=http://localhost:3000 - uses: actions/upload-artifact@v4 @@ -67,6 +71,28 @@ jobs: path: compose-logs.txt if-no-files-found: ignore + # pytest: + # name: Run Backend Tests + # runs-on: ubuntu-latest + # steps: + # - uses: actions/checkout@v4 + + # - name: Set up Python + # uses: actions/setup-python@v4 + # with: + # python-version: ${{ matrix.python-version }} + + # - name: Install dependencies + # run: | + # python -m pip install --upgrade pip + # pip install -r backend/requirements.txt + + # - name: pytest run + # run: | + # ls -al + # cd backend + # PYTHONPATH=. pytest . -o log_cli=true -o log_cli_level=INFO + migration_test: name: Run Migration Tests runs-on: ubuntu-latest @@ -126,11 +152,11 @@ jobs: cd backend uvicorn main:app --port "8080" --forwarded-allow-ips '*' & UVICORN_PID=$! - # Wait up to 20 seconds for the server to start - for i in {1..20}; do + # Wait up to 40 seconds for the server to start + for i in {1..40}; do curl -s http://localhost:8080/api/config > /dev/null && break sleep 1 - if [ $i -eq 20 ]; then + if [ $i -eq 40 ]; then echo "Server failed to start" kill -9 $UVICORN_PID exit 1 @@ -171,7 +197,7 @@ jobs: fi # Check that service will reconnect to postgres when connection will be closed - status_code=$(curl --write-out %{http_code} -s --output /dev/null http://localhost:8081/health) + status_code=$(curl --write-out %{http_code} -s --output /dev/null http://localhost:8081/health/db) if [[ "$status_code" -ne 200 ]] ; then echo "Server has failed before postgres reconnect check" exit 1 @@ -183,7 +209,7 @@ jobs: cur = conn.cursor(); \ cur.execute('SELECT pg_terminate_backend(psa.pid) FROM pg_stat_activity psa WHERE datname = current_database() AND pid <> pg_backend_pid();')" - status_code=$(curl --write-out %{http_code} -s --output /dev/null http://localhost:8081/health) + status_code=$(curl --write-out %{http_code} -s --output /dev/null http://localhost:8081/health/db) if [[ "$status_code" -ne 200 ]] ; then echo "Server has not reconnected to postgres after connection was closed: returned status $status_code" exit 1 diff --git a/.gitignore b/.gitignore index a54fef595..32271f808 100644 --- a/.gitignore +++ b/.gitignore @@ -306,3 +306,4 @@ dist # cypress artifacts cypress/videos cypress/screenshots +.vscode/settings.json diff --git a/CHANGELOG.md b/CHANGELOG.md index be9b8ec7a..c81aed300 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,33 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [0.3.8] - 2024-07-09 + +### Added + +- **💬 Chat Controls**: Easily adjust parameters for each chat session, offering more precise control over your interactions. +- **📌 Pinned Chats**: Support for pinned chats, allowing you to keep important conversations easily accessible. +- **📄 Apache Tika Integration**: Added support for using Apache Tika as a document loader, enhancing document processing capabilities. +- **🛠️ Custom Environment for OpenID Claims**: Allows setting custom claims for OpenID, providing more flexibility in user authentication. +- **🔧 Enhanced Tools & Functions API**: Introduced 'event_emitter' and 'event_call', now you can also add citations for better documentation and tracking. Detailed documentation will be provided on our documentation website. +- **↔️ Sideways Scrolling in Settings**: Settings tabs container now supports horizontal scrolling for easier navigation. +- **🌑 Darker OLED Theme**: Includes a new, darker OLED theme and improved styling for the light theme, enhancing visual appeal. +- **🌐 Language Updates**: Updated translations for Indonesian, German, French, and Catalan languages, expanding accessibility. + +### Fixed + +- **⏰ OpenAI Streaming Timeout**: Resolved issues with OpenAI streaming response using the 'AIOHTTP_CLIENT_TIMEOUT' setting, ensuring reliable performance. +- **💡 User Valves**: Fixed malfunctioning user valves, ensuring proper functionality. +- **🔄 Collapsible Components**: Addressed issues with collapsible components not working, restoring expected behavior. + +### Changed + +- **🗃️ Database Backend**: Switched from Peewee to SQLAlchemy for improved concurrency support, enhancing database performance. +- **🔤 Primary Font Styling**: Updated primary font to Archivo for better visual consistency. +- **🔄 Font Change for Windows**: Replaced Arimo with Inter font for Windows users, improving readability. +- **🚀 Lazy Loading**: Implemented lazy loading for 'faster_whisper' and 'sentence_transformers' to reduce startup memory usage. +- **📋 Task Generation Payload**: Task generations now include only the "task" field in the body instead of "title". + ## [0.3.7] - 2024-06-29 ### Added diff --git a/backend/alembic.ini b/backend/alembic.ini new file mode 100644 index 000000000..4eff85f0c --- /dev/null +++ b/backend/alembic.ini @@ -0,0 +1,114 @@ +# A generic, single database configuration. + +[alembic] +# path to migration scripts +script_location = migrations + +# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s +# Uncomment the line below if you want the files to be prepended with date and time +# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s + +# sys.path path, will be prepended to sys.path if present. +# defaults to the current working directory. +prepend_sys_path = . + +# timezone to use when rendering the date within the migration file +# as well as the filename. +# If specified, requires the python>=3.9 or backports.zoneinfo library. +# Any required deps can installed by adding `alembic[tz]` to the pip requirements +# string value is passed to ZoneInfo() +# leave blank for localtime +# timezone = + +# max length of characters to apply to the +# "slug" field +# truncate_slug_length = 40 + +# set to 'true' to run the environment during +# the 'revision' command, regardless of autogenerate +# revision_environment = false + +# set to 'true' to allow .pyc and .pyo files without +# a source .py file to be detected as revisions in the +# versions/ directory +# sourceless = false + +# version location specification; This defaults +# to migrations/versions. When using multiple version +# directories, initial revisions must be specified with --version-path. +# The path separator used here should be the separator specified by "version_path_separator" below. +# version_locations = %(here)s/bar:%(here)s/bat:migrations/versions + +# version path separator; As mentioned above, this is the character used to split +# version_locations. The default within new alembic.ini files is "os", which uses os.pathsep. +# If this key is omitted entirely, it falls back to the legacy behavior of splitting on spaces and/or commas. +# Valid values for version_path_separator are: +# +# version_path_separator = : +# version_path_separator = ; +# version_path_separator = space +version_path_separator = os # Use os.pathsep. Default configuration used for new projects. + +# set to 'true' to search source files recursively +# in each "version_locations" directory +# new in Alembic version 1.10 +# recursive_version_locations = false + +# the output encoding used when revision files +# are written from script.py.mako +# output_encoding = utf-8 + +# sqlalchemy.url = REPLACE_WITH_DATABASE_URL + + +[post_write_hooks] +# post_write_hooks defines scripts or Python functions that are run +# on newly generated revision scripts. See the documentation for further +# detail and examples + +# format using "black" - use the console_scripts runner, against the "black" entrypoint +# hooks = black +# black.type = console_scripts +# black.entrypoint = black +# black.options = -l 79 REVISION_SCRIPT_FILENAME + +# lint with attempts to fix using "ruff" - use the exec runner, execute a binary +# hooks = ruff +# ruff.type = exec +# ruff.executable = %(here)s/.venv/bin/ruff +# ruff.options = --fix REVISION_SCRIPT_FILENAME + +# Logging configuration +[loggers] +keys = root,sqlalchemy,alembic + +[handlers] +keys = console + +[formatters] +keys = generic + +[logger_root] +level = WARN +handlers = console +qualname = + +[logger_sqlalchemy] +level = WARN +handlers = +qualname = sqlalchemy.engine + +[logger_alembic] +level = INFO +handlers = +qualname = alembic + +[handler_console] +class = StreamHandler +args = (sys.stderr,) +level = NOTSET +formatter = generic + +[formatter_generic] +format = %(levelname)-5.5s [%(name)s] %(message)s +datefmt = %H:%M:%S diff --git a/backend/apps/audio/main.py b/backend/apps/audio/main.py index 8843f376f..f866d867f 100644 --- a/backend/apps/audio/main.py +++ b/backend/apps/audio/main.py @@ -14,7 +14,6 @@ from fastapi import ( from fastapi.responses import StreamingResponse, JSONResponse, FileResponse from fastapi.middleware.cors import CORSMiddleware -from faster_whisper import WhisperModel from pydantic import BaseModel import uuid @@ -277,6 +276,8 @@ def transcribe( f.close() if app.state.config.STT_ENGINE == "": + from faster_whisper import WhisperModel + whisper_kwargs = { "model_size_or_path": WHISPER_MODEL, "device": whisper_device_type, diff --git a/backend/apps/images/main.py b/backend/apps/images/main.py index 8f1a08e04..24542ee93 100644 --- a/backend/apps/images/main.py +++ b/backend/apps/images/main.py @@ -12,7 +12,6 @@ from fastapi import ( Form, ) from fastapi.middleware.cors import CORSMiddleware -from faster_whisper import WhisperModel from constants import ERROR_MESSAGES from utils.utils import ( diff --git a/backend/apps/openai/main.py b/backend/apps/openai/main.py index 31dd48741..7c67c40ae 100644 --- a/backend/apps/openai/main.py +++ b/backend/apps/openai/main.py @@ -25,6 +25,7 @@ from utils.task import prompt_template from config import ( SRC_LOG_LEVELS, ENABLE_OPENAI_API, + AIOHTTP_CLIENT_TIMEOUT, OPENAI_API_BASE_URLS, OPENAI_API_KEYS, CACHE_DIR, @@ -463,7 +464,9 @@ async def generate_chat_completion( streaming = False try: - session = aiohttp.ClientSession(trust_env=True) + session = aiohttp.ClientSession( + trust_env=True, timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT) + ) r = await session.request( method="POST", url=f"{url}/chat/completions", diff --git a/backend/apps/rag/main.py b/backend/apps/rag/main.py index 7c6974535..c0f8a09ed 100644 --- a/backend/apps/rag/main.py +++ b/backend/apps/rag/main.py @@ -48,8 +48,6 @@ import mimetypes import uuid import json -import sentence_transformers - from apps.webui.models.documents import ( Documents, DocumentForm, @@ -93,6 +91,8 @@ from config import ( SRC_LOG_LEVELS, UPLOAD_DIR, DOCS_DIR, + CONTENT_EXTRACTION_ENGINE, + TIKA_SERVER_URL, RAG_TOP_K, RAG_RELEVANCE_THRESHOLD, RAG_EMBEDDING_ENGINE, @@ -148,6 +148,9 @@ app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION = ( ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION ) +app.state.config.CONTENT_EXTRACTION_ENGINE = CONTENT_EXTRACTION_ENGINE +app.state.config.TIKA_SERVER_URL = TIKA_SERVER_URL + app.state.config.CHUNK_SIZE = CHUNK_SIZE app.state.config.CHUNK_OVERLAP = CHUNK_OVERLAP @@ -190,6 +193,8 @@ def update_embedding_model( update_model: bool = False, ): if embedding_model and app.state.config.RAG_EMBEDDING_ENGINE == "": + import sentence_transformers + app.state.sentence_transformer_ef = sentence_transformers.SentenceTransformer( get_model_path(embedding_model, update_model), device=DEVICE_TYPE, @@ -204,6 +209,8 @@ def update_reranking_model( update_model: bool = False, ): if reranking_model: + import sentence_transformers + app.state.sentence_transformer_rf = sentence_transformers.CrossEncoder( get_model_path(reranking_model, update_model), device=DEVICE_TYPE, @@ -388,6 +395,10 @@ async def get_rag_config(user=Depends(get_admin_user)): return { "status": True, "pdf_extract_images": app.state.config.PDF_EXTRACT_IMAGES, + "content_extraction": { + "engine": app.state.config.CONTENT_EXTRACTION_ENGINE, + "tika_server_url": app.state.config.TIKA_SERVER_URL, + }, "chunk": { "chunk_size": app.state.config.CHUNK_SIZE, "chunk_overlap": app.state.config.CHUNK_OVERLAP, @@ -417,6 +428,11 @@ async def get_rag_config(user=Depends(get_admin_user)): } +class ContentExtractionConfig(BaseModel): + engine: str = "" + tika_server_url: Optional[str] = None + + class ChunkParamUpdateForm(BaseModel): chunk_size: int chunk_overlap: int @@ -450,6 +466,7 @@ class WebConfig(BaseModel): class ConfigUpdateForm(BaseModel): pdf_extract_images: Optional[bool] = None + content_extraction: Optional[ContentExtractionConfig] = None chunk: Optional[ChunkParamUpdateForm] = None youtube: Optional[YoutubeLoaderConfig] = None web: Optional[WebConfig] = None @@ -463,6 +480,11 @@ async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_ else app.state.config.PDF_EXTRACT_IMAGES ) + if form_data.content_extraction is not None: + log.info(f"Updating text settings: {form_data.content_extraction}") + app.state.config.CONTENT_EXTRACTION_ENGINE = form_data.content_extraction.engine + app.state.config.TIKA_SERVER_URL = form_data.content_extraction.tika_server_url + if form_data.chunk is not None: app.state.config.CHUNK_SIZE = form_data.chunk.chunk_size app.state.config.CHUNK_OVERLAP = form_data.chunk.chunk_overlap @@ -499,6 +521,10 @@ async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_ return { "status": True, "pdf_extract_images": app.state.config.PDF_EXTRACT_IMAGES, + "content_extraction": { + "engine": app.state.config.CONTENT_EXTRACTION_ENGINE, + "tika_server_url": app.state.config.TIKA_SERVER_URL, + }, "chunk": { "chunk_size": app.state.config.CHUNK_SIZE, "chunk_overlap": app.state.config.CHUNK_OVERLAP, @@ -978,13 +1004,49 @@ def store_docs_in_vector_db(docs, collection_name, overwrite: bool = False) -> b return True except Exception as e: - log.exception(e) if e.__class__.__name__ == "UniqueConstraintError": return True + log.exception(e) + return False +class TikaLoader: + def __init__(self, file_path, mime_type=None): + self.file_path = file_path + self.mime_type = mime_type + + def load(self) -> List[Document]: + with open(self.file_path, "rb") as f: + data = f.read() + + if self.mime_type is not None: + headers = {"Content-Type": self.mime_type} + else: + headers = {} + + endpoint = app.state.config.TIKA_SERVER_URL + if not endpoint.endswith("/"): + endpoint += "/" + endpoint += "tika/text" + + r = requests.put(endpoint, data=data, headers=headers) + + if r.ok: + raw_metadata = r.json() + text = raw_metadata.get("X-TIKA:content", "") + + if "Content-Type" in raw_metadata: + headers["Content-Type"] = raw_metadata["Content-Type"] + + log.info("Tika extracted text: %s", text) + + return [Document(page_content=text, metadata=headers)] + else: + raise Exception(f"Error calling Tika: {r.reason}") + + def get_loader(filename: str, file_content_type: str, file_path: str): file_ext = filename.split(".")[-1].lower() known_type = True @@ -1035,47 +1097,58 @@ def get_loader(filename: str, file_content_type: str, file_path: str): "msg", ] - if file_ext == "pdf": - loader = PyPDFLoader( - file_path, extract_images=app.state.config.PDF_EXTRACT_IMAGES - ) - elif file_ext == "csv": - loader = CSVLoader(file_path) - elif file_ext == "rst": - loader = UnstructuredRSTLoader(file_path, mode="elements") - elif file_ext == "xml": - loader = UnstructuredXMLLoader(file_path) - elif file_ext in ["htm", "html"]: - loader = BSHTMLLoader(file_path, open_encoding="unicode_escape") - elif file_ext == "md": - loader = UnstructuredMarkdownLoader(file_path) - elif file_content_type == "application/epub+zip": - loader = UnstructuredEPubLoader(file_path) - elif ( - file_content_type - == "application/vnd.openxmlformats-officedocument.wordprocessingml.document" - or file_ext in ["doc", "docx"] + if ( + app.state.config.CONTENT_EXTRACTION_ENGINE == "tika" + and app.state.config.TIKA_SERVER_URL ): - loader = Docx2txtLoader(file_path) - elif file_content_type in [ - "application/vnd.ms-excel", - "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", - ] or file_ext in ["xls", "xlsx"]: - loader = UnstructuredExcelLoader(file_path) - elif file_content_type in [ - "application/vnd.ms-powerpoint", - "application/vnd.openxmlformats-officedocument.presentationml.presentation", - ] or file_ext in ["ppt", "pptx"]: - loader = UnstructuredPowerPointLoader(file_path) - elif file_ext == "msg": - loader = OutlookMessageLoader(file_path) - elif file_ext in known_source_ext or ( - file_content_type and file_content_type.find("text/") >= 0 - ): - loader = TextLoader(file_path, autodetect_encoding=True) + if file_ext in known_source_ext or ( + file_content_type and file_content_type.find("text/") >= 0 + ): + loader = TextLoader(file_path, autodetect_encoding=True) + else: + loader = TikaLoader(file_path, file_content_type) else: - loader = TextLoader(file_path, autodetect_encoding=True) - known_type = False + if file_ext == "pdf": + loader = PyPDFLoader( + file_path, extract_images=app.state.config.PDF_EXTRACT_IMAGES + ) + elif file_ext == "csv": + loader = CSVLoader(file_path) + elif file_ext == "rst": + loader = UnstructuredRSTLoader(file_path, mode="elements") + elif file_ext == "xml": + loader = UnstructuredXMLLoader(file_path) + elif file_ext in ["htm", "html"]: + loader = BSHTMLLoader(file_path, open_encoding="unicode_escape") + elif file_ext == "md": + loader = UnstructuredMarkdownLoader(file_path) + elif file_content_type == "application/epub+zip": + loader = UnstructuredEPubLoader(file_path) + elif ( + file_content_type + == "application/vnd.openxmlformats-officedocument.wordprocessingml.document" + or file_ext in ["doc", "docx"] + ): + loader = Docx2txtLoader(file_path) + elif file_content_type in [ + "application/vnd.ms-excel", + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + ] or file_ext in ["xls", "xlsx"]: + loader = UnstructuredExcelLoader(file_path) + elif file_content_type in [ + "application/vnd.ms-powerpoint", + "application/vnd.openxmlformats-officedocument.presentationml.presentation", + ] or file_ext in ["ppt", "pptx"]: + loader = UnstructuredPowerPointLoader(file_path) + elif file_ext == "msg": + loader = OutlookMessageLoader(file_path) + elif file_ext in known_source_ext or ( + file_content_type and file_content_type.find("text/") >= 0 + ): + loader = TextLoader(file_path, autodetect_encoding=True) + else: + loader = TextLoader(file_path, autodetect_encoding=True) + known_type = False return loader, known_type diff --git a/backend/apps/rag/utils.py b/backend/apps/rag/utils.py index 7b4324d9a..fde89b069 100644 --- a/backend/apps/rag/utils.py +++ b/backend/apps/rag/utils.py @@ -294,14 +294,16 @@ def get_rag_context( extracted_collections.extend(collection_names) - context_string = "" - + contexts = [] citations = [] + for context in relevant_contexts: try: if "documents" in context: - context_string += "\n\n".join( - [text for text in context["documents"][0] if text is not None] + contexts.append( + "\n\n".join( + [text for text in context["documents"][0] if text is not None] + ) ) if "metadatas" in context: @@ -315,9 +317,7 @@ def get_rag_context( except Exception as e: log.exception(e) - context_string = context_string.strip() - - return context_string, citations + return contexts, citations def get_model_path(model: str, update_model: bool = False): @@ -442,8 +442,6 @@ from langchain_core.documents import BaseDocumentCompressor, Document from langchain_core.callbacks import Callbacks from langchain_core.pydantic_v1 import Extra -from sentence_transformers import util - class RerankCompressor(BaseDocumentCompressor): embedding_function: Any @@ -468,6 +466,8 @@ class RerankCompressor(BaseDocumentCompressor): [(query, doc.page_content) for doc in documents] ) else: + from sentence_transformers import util + query_embedding = self.embedding_function(query) document_embedding = self.embedding_function( [doc.page_content for doc in documents] diff --git a/backend/apps/webui/internal/db.py b/backend/apps/webui/internal/db.py index 80c30d652..fbe287e18 100644 --- a/backend/apps/webui/internal/db.py +++ b/backend/apps/webui/internal/db.py @@ -1,18 +1,39 @@ import os import logging import json +from contextlib import contextmanager -from peewee import * from peewee_migrate import Router - from apps.webui.internal.wrappers import register_connection + +from typing import Optional, Any +from typing_extensions import Self + +from sqlalchemy import create_engine, types, Dialect +from sqlalchemy.ext.declarative import declarative_base +from sqlalchemy.orm import sessionmaker, scoped_session +from sqlalchemy.sql.type_api import _T + from config import SRC_LOG_LEVELS, DATA_DIR, DATABASE_URL, BACKEND_DIR log = logging.getLogger(__name__) log.setLevel(SRC_LOG_LEVELS["DB"]) -class JSONField(TextField): +class JSONField(types.TypeDecorator): + impl = types.Text + cache_ok = True + + def process_bind_param(self, value: Optional[_T], dialect: Dialect) -> Any: + return json.dumps(value) + + def process_result_value(self, value: Optional[_T], dialect: Dialect) -> Any: + if value is not None: + return json.loads(value) + + def copy(self, **kw: Any) -> Self: + return JSONField(self.impl.length) + def db_value(self, value): return json.dumps(value) @@ -30,25 +51,60 @@ else: pass -# The `register_connection` function encapsulates the logic for setting up -# the database connection based on the connection string, while `connect` -# is a Peewee-specific method to manage the connection state and avoid errors -# when a connection is already open. -try: - DB = register_connection(DATABASE_URL) - log.info(f"Connected to a {DB.__class__.__name__} database.") -except Exception as e: - log.error(f"Failed to initialize the database connection: {e}") - raise +# Workaround to handle the peewee migration +# This is required to ensure the peewee migration is handled before the alembic migration +def handle_peewee_migration(DATABASE_URL): + try: + # Replace the postgresql:// with postgres:// and %40 with @ in the DATABASE_URL + db = register_connection( + DATABASE_URL.replace("postgresql://", "postgres://").replace("%40", "@") + ) + migrate_dir = BACKEND_DIR / "apps" / "webui" / "internal" / "migrations" + router = Router(db, logger=log, migrate_dir=migrate_dir) + router.run() + db.close() -router = Router( - DB, - migrate_dir=BACKEND_DIR / "apps" / "webui" / "internal" / "migrations", - logger=log, + # check if db connection has been closed + + except Exception as e: + log.error(f"Failed to initialize the database connection: {e}") + raise + + finally: + # Properly closing the database connection + if db and not db.is_closed(): + db.close() + + # Assert if db connection has been closed + assert db.is_closed(), "Database connection is still open." + + +handle_peewee_migration(DATABASE_URL) + + +SQLALCHEMY_DATABASE_URL = DATABASE_URL +if "sqlite" in SQLALCHEMY_DATABASE_URL: + engine = create_engine( + SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False} + ) +else: + engine = create_engine(SQLALCHEMY_DATABASE_URL, pool_pre_ping=True) + + +SessionLocal = sessionmaker( + autocommit=False, autoflush=False, bind=engine, expire_on_commit=False ) -router.run() -try: - DB.connect(reuse_if_open=True) -except OperationalError as e: - log.info(f"Failed to connect to database again due to: {e}") - pass +Base = declarative_base() +Session = scoped_session(SessionLocal) + + +# Dependency +def get_session(): + db = SessionLocal() + try: + yield db + finally: + db.close() + + +get_db = contextmanager(get_session) diff --git a/backend/apps/webui/internal/migrations/017_add_user_oauth_sub.py b/backend/apps/webui/internal/migrations/017_add_user_oauth_sub.py index fd1d9b560..eaa3fa5fe 100644 --- a/backend/apps/webui/internal/migrations/017_add_user_oauth_sub.py +++ b/backend/apps/webui/internal/migrations/017_add_user_oauth_sub.py @@ -1,10 +1,7 @@ """Peewee migrations -- 017_add_user_oauth_sub.py. - Some examples (model - class or model name):: - > Model = migrator.orm['table_name'] # Return model in current state by name > Model = migrator.ModelClass # Return model in current state by name - > migrator.sql(sql) # Run custom SQL > migrator.run(func, *args, **kwargs) # Run python function with the given args > migrator.create_model(Model) # Create a model (could be used as decorator) @@ -21,7 +18,6 @@ Some examples (model - class or model name):: > migrator.drop_index(model, *col_names) > migrator.drop_not_null(model, *field_names) > migrator.drop_constraints(model, *constraints) - """ from contextlib import suppress diff --git a/backend/apps/webui/internal/migrations/README.md b/backend/apps/webui/internal/migrations/README.md deleted file mode 100644 index 260214113..000000000 --- a/backend/apps/webui/internal/migrations/README.md +++ /dev/null @@ -1,21 +0,0 @@ -# Database Migrations - -This directory contains all the database migrations for the web app. -Migrations are done using the [`peewee-migrate`](https://github.com/klen/peewee_migrate) library. - -Migrations are automatically ran at app startup. - -## Creating a migration - -Have you made a change to the schema of an existing model? -You will need to create a migration file to ensure that existing databases are updated for backwards compatibility. - -1. Have a database file (`webui.db`) that has the old schema prior to any of your changes. -2. Make your changes to the models. -3. From the `backend` directory, run the following command: - ```bash - pw_migrate create --auto --auto-source apps.webui.models --database sqlite:///${SQLITE_DB} --directory apps/web/internal/migrations ${MIGRATION_NAME} - ``` - - `$SQLITE_DB` should be the path to the database file. - - `$MIGRATION_NAME` should be a descriptive name for the migration. -4. The migration file will be created in the `apps/web/internal/migrations` directory. diff --git a/backend/apps/webui/main.py b/backend/apps/webui/main.py index 28b1b4aac..ab28868ae 100644 --- a/backend/apps/webui/main.py +++ b/backend/apps/webui/main.py @@ -3,7 +3,7 @@ from fastapi.routing import APIRoute from fastapi.responses import StreamingResponse from fastapi.middleware.cors import CORSMiddleware from starlette.middleware.sessions import SessionMiddleware - +from sqlalchemy.orm import Session from apps.webui.routers import ( auths, users, @@ -19,8 +19,13 @@ from apps.webui.routers import ( functions, ) from apps.webui.models.functions import Functions +from apps.webui.models.models import Models + from apps.webui.utils import load_function_module_by_id + from utils.misc import stream_message_template +from utils.task import prompt_template + from config import ( WEBUI_BUILD_HASH, @@ -39,6 +44,8 @@ from config import ( WEBUI_BANNERS, ENABLE_COMMUNITY_SHARING, AppConfig, + OAUTH_USERNAME_CLAIM, + OAUTH_PICTURE_CLAIM, ) import inspect @@ -74,6 +81,9 @@ app.state.config.BANNERS = WEBUI_BANNERS app.state.config.ENABLE_COMMUNITY_SHARING = ENABLE_COMMUNITY_SHARING +app.state.config.OAUTH_USERNAME_CLAIM = OAUTH_USERNAME_CLAIM +app.state.config.OAUTH_PICTURE_CLAIM = OAUTH_PICTURE_CLAIM + app.state.MODELS = {} app.state.TOOLS = {} app.state.FUNCTIONS = {} @@ -129,7 +139,6 @@ async def get_pipe_models(): function_module = app.state.FUNCTIONS[pipe.id] if hasattr(function_module, "valves") and hasattr(function_module, "Valves"): - print(f"Getting valves for {pipe.id}") valves = Functions.get_function_valves_by_id(pipe.id) function_module.valves = function_module.Valves( **(valves if valves else {}) @@ -181,6 +190,77 @@ async def get_pipe_models(): async def generate_function_chat_completion(form_data, user): + model_id = form_data.get("model") + model_info = Models.get_model_by_id(model_id) + + if model_info: + if model_info.base_model_id: + form_data["model"] = model_info.base_model_id + + model_info.params = model_info.params.model_dump() + + if model_info.params: + if model_info.params.get("temperature", None) is not None: + form_data["temperature"] = float(model_info.params.get("temperature")) + + if model_info.params.get("top_p", None): + form_data["top_p"] = int(model_info.params.get("top_p", None)) + + if model_info.params.get("max_tokens", None): + form_data["max_tokens"] = int(model_info.params.get("max_tokens", None)) + + if model_info.params.get("frequency_penalty", None): + form_data["frequency_penalty"] = int( + model_info.params.get("frequency_penalty", None) + ) + + if model_info.params.get("seed", None): + form_data["seed"] = model_info.params.get("seed", None) + + if model_info.params.get("stop", None): + form_data["stop"] = ( + [ + bytes(stop, "utf-8").decode("unicode_escape") + for stop in model_info.params["stop"] + ] + if model_info.params.get("stop", None) + else None + ) + + system = model_info.params.get("system", None) + if system: + system = prompt_template( + system, + **( + { + "user_name": user.name, + "user_location": ( + user.info.get("location") if user.info else None + ), + } + if user + else {} + ), + ) + # Check if the payload already has a system message + # If not, add a system message to the payload + if form_data.get("messages"): + for message in form_data["messages"]: + if message.get("role") == "system": + message["content"] = system + message["content"] + break + else: + form_data["messages"].insert( + 0, + { + "role": "system", + "content": system, + }, + ) + + else: + pass + async def job(): pipe_id = form_data["model"] if "." in pipe_id: @@ -259,6 +339,9 @@ async def generate_function_chat_completion(form_data, user): if isinstance(line, BaseModel): line = line.model_dump_json() line = f"data: {line}" + if isinstance(line, dict): + line = f"data: {json.dumps(line)}" + try: line = line.decode("utf-8") except: diff --git a/backend/apps/webui/models/auths.py b/backend/apps/webui/models/auths.py index 9ea38abcb..a17d92ff1 100644 --- a/backend/apps/webui/models/auths.py +++ b/backend/apps/webui/models/auths.py @@ -1,14 +1,13 @@ from pydantic import BaseModel -from typing import List, Union, Optional -import time +from typing import Optional import uuid import logging -from peewee import * +from sqlalchemy import String, Column, Boolean, Text from apps.webui.models.users import UserModel, Users from utils.utils import verify_password -from apps.webui.internal.db import DB +from apps.webui.internal.db import Base, get_db from config import SRC_LOG_LEVELS @@ -20,14 +19,13 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"]) #################### -class Auth(Model): - id = CharField(unique=True) - email = CharField() - password = TextField() - active = BooleanField() +class Auth(Base): + __tablename__ = "auth" - class Meta: - database = DB + id = Column(String, primary_key=True) + email = Column(String) + password = Column(Text) + active = Column(Boolean) class AuthModel(BaseModel): @@ -94,9 +92,6 @@ class AddUserForm(SignupForm): class AuthsTable: - def __init__(self, db): - self.db = db - self.db.create_tables([Auth]) def insert_new_auth( self, @@ -107,36 +102,44 @@ class AuthsTable: role: str = "pending", oauth_sub: Optional[str] = None, ) -> Optional[UserModel]: - log.info("insert_new_auth") + with get_db() as db: - id = str(uuid.uuid4()) + log.info("insert_new_auth") - auth = AuthModel( - **{"id": id, "email": email, "password": password, "active": True} - ) - result = Auth.create(**auth.model_dump()) + id = str(uuid.uuid4()) - user = Users.insert_new_user( - id, name, email, profile_image_url, role, oauth_sub - ) + auth = AuthModel( + **{"id": id, "email": email, "password": password, "active": True} + ) + result = Auth(**auth.model_dump()) + db.add(result) - if result and user: - return user - else: - return None + user = Users.insert_new_user( + id, name, email, profile_image_url, role, oauth_sub + ) + + db.commit() + db.refresh(result) + + if result and user: + return user + else: + return None def authenticate_user(self, email: str, password: str) -> Optional[UserModel]: log.info(f"authenticate_user: {email}") try: - auth = Auth.get(Auth.email == email, Auth.active == True) - if auth: - if verify_password(password, auth.password): - user = Users.get_user_by_id(auth.id) - return user + with get_db() as db: + + auth = db.query(Auth).filter_by(email=email, active=True).first() + if auth: + if verify_password(password, auth.password): + user = Users.get_user_by_id(auth.id) + return user + else: + return None else: return None - else: - return None except: return None @@ -155,46 +158,50 @@ class AuthsTable: def authenticate_user_by_trusted_header(self, email: str) -> Optional[UserModel]: log.info(f"authenticate_user_by_trusted_header: {email}") try: - auth = Auth.get(Auth.email == email, Auth.active == True) - if auth: - user = Users.get_user_by_id(auth.id) - return user + with get_db() as db: + auth = db.query(Auth).filter(email=email, active=True).first() + if auth: + user = Users.get_user_by_id(auth.id) + return user except: return None def update_user_password_by_id(self, id: str, new_password: str) -> bool: try: - query = Auth.update(password=new_password).where(Auth.id == id) - result = query.execute() - - return True if result == 1 else False + with get_db() as db: + result = ( + db.query(Auth).filter_by(id=id).update({"password": new_password}) + ) + db.commit() + return True if result == 1 else False except: return False def update_email_by_id(self, id: str, email: str) -> bool: try: - query = Auth.update(email=email).where(Auth.id == id) - result = query.execute() - - return True if result == 1 else False + with get_db() as db: + result = db.query(Auth).filter_by(id=id).update({"email": email}) + db.commit() + return True if result == 1 else False except: return False def delete_auth_by_id(self, id: str) -> bool: try: - # Delete User - result = Users.delete_user_by_id(id) + with get_db() as db: - if result: - # Delete Auth - query = Auth.delete().where(Auth.id == id) - query.execute() # Remove the rows, return number of rows removed. + # Delete User + result = Users.delete_user_by_id(id) - return True - else: - return False + if result: + db.query(Auth).filter_by(id=id).delete() + db.commit() + + return True + else: + return False except: return False -Auths = AuthsTable(DB) +Auths = AuthsTable() diff --git a/backend/apps/webui/models/chats.py b/backend/apps/webui/models/chats.py index a6f1ae923..c03abb233 100644 --- a/backend/apps/webui/models/chats.py +++ b/backend/apps/webui/models/chats.py @@ -1,36 +1,38 @@ -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict from typing import List, Union, Optional -from peewee import * -from playhouse.shortcuts import model_to_dict import json import uuid import time -from apps.webui.internal.db import DB +from sqlalchemy import Column, String, BigInteger, Boolean, Text + +from apps.webui.internal.db import Base, get_db + #################### # Chat DB Schema #################### -class Chat(Model): - id = CharField(unique=True) - user_id = CharField() - title = TextField() - chat = TextField() # Save Chat JSON as Text +class Chat(Base): + __tablename__ = "chat" - created_at = BigIntegerField() - updated_at = BigIntegerField() + id = Column(String, primary_key=True) + user_id = Column(String) + title = Column(Text) + chat = Column(Text) # Save Chat JSON as Text - share_id = CharField(null=True, unique=True) - archived = BooleanField(default=False) + created_at = Column(BigInteger) + updated_at = Column(BigInteger) - class Meta: - database = DB + share_id = Column(Text, unique=True, nullable=True) + archived = Column(Boolean, default=False) class ChatModel(BaseModel): + model_config = ConfigDict(from_attributes=True) + id: str user_id: str title: str @@ -75,91 +77,104 @@ class ChatTitleIdResponse(BaseModel): class ChatTable: - def __init__(self, db): - self.db = db - db.create_tables([Chat]) def insert_new_chat(self, user_id: str, form_data: ChatForm) -> Optional[ChatModel]: - id = str(uuid.uuid4()) - chat = ChatModel( - **{ - "id": id, - "user_id": user_id, - "title": ( - form_data.chat["title"] if "title" in form_data.chat else "New Chat" - ), - "chat": json.dumps(form_data.chat), - "created_at": int(time.time()), - "updated_at": int(time.time()), - } - ) + with get_db() as db: - result = Chat.create(**chat.model_dump()) - return chat if result else None + id = str(uuid.uuid4()) + chat = ChatModel( + **{ + "id": id, + "user_id": user_id, + "title": ( + form_data.chat["title"] + if "title" in form_data.chat + else "New Chat" + ), + "chat": json.dumps(form_data.chat), + "created_at": int(time.time()), + "updated_at": int(time.time()), + } + ) + + result = Chat(**chat.model_dump()) + db.add(result) + db.commit() + db.refresh(result) + return ChatModel.model_validate(result) if result else None def update_chat_by_id(self, id: str, chat: dict) -> Optional[ChatModel]: try: - query = Chat.update( - chat=json.dumps(chat), - title=chat["title"] if "title" in chat else "New Chat", - updated_at=int(time.time()), - ).where(Chat.id == id) - query.execute() + with get_db() as db: - chat = Chat.get(Chat.id == id) - return ChatModel(**model_to_dict(chat)) - except: + chat_obj = db.get(Chat, id) + chat_obj.chat = json.dumps(chat) + chat_obj.title = chat["title"] if "title" in chat else "New Chat" + chat_obj.updated_at = int(time.time()) + db.commit() + db.refresh(chat_obj) + + return ChatModel.model_validate(chat_obj) + except Exception as e: return None def insert_shared_chat_by_chat_id(self, chat_id: str) -> Optional[ChatModel]: - # Get the existing chat to share - chat = Chat.get(Chat.id == chat_id) - # Check if the chat is already shared - if chat.share_id: - return self.get_chat_by_id_and_user_id(chat.share_id, "shared") - # Create a new chat with the same data, but with a new ID - shared_chat = ChatModel( - **{ - "id": str(uuid.uuid4()), - "user_id": f"shared-{chat_id}", - "title": chat.title, - "chat": chat.chat, - "created_at": chat.created_at, - "updated_at": int(time.time()), - } - ) - shared_result = Chat.create(**shared_chat.model_dump()) - # Update the original chat with the share_id - result = ( - Chat.update(share_id=shared_chat.id).where(Chat.id == chat_id).execute() - ) + with get_db() as db: - return shared_chat if (shared_result and result) else None + # Get the existing chat to share + chat = db.get(Chat, chat_id) + # Check if the chat is already shared + if chat.share_id: + return self.get_chat_by_id_and_user_id(chat.share_id, "shared") + # Create a new chat with the same data, but with a new ID + shared_chat = ChatModel( + **{ + "id": str(uuid.uuid4()), + "user_id": f"shared-{chat_id}", + "title": chat.title, + "chat": chat.chat, + "created_at": chat.created_at, + "updated_at": int(time.time()), + } + ) + shared_result = Chat(**shared_chat.model_dump()) + db.add(shared_result) + db.commit() + db.refresh(shared_result) + + # Update the original chat with the share_id + result = ( + db.query(Chat) + .filter_by(id=chat_id) + .update({"share_id": shared_chat.id}) + ) + db.commit() + return shared_chat if (shared_result and result) else None def update_shared_chat_by_chat_id(self, chat_id: str) -> Optional[ChatModel]: try: - print("update_shared_chat_by_id") - chat = Chat.get(Chat.id == chat_id) - print(chat) + with get_db() as db: - query = Chat.update( - title=chat.title, - chat=chat.chat, - ).where(Chat.id == chat.share_id) + print("update_shared_chat_by_id") + chat = db.get(Chat, chat_id) + print(chat) + chat.title = chat.title + chat.chat = chat.chat + db.commit() + db.refresh(chat) - query.execute() - - chat = Chat.get(Chat.id == chat.share_id) - return ChatModel(**model_to_dict(chat)) + return self.get_chat_by_id(chat.share_id) except: return None def delete_shared_chat_by_chat_id(self, chat_id: str) -> bool: try: - query = Chat.delete().where(Chat.user_id == f"shared-{chat_id}") - query.execute() # Remove the rows, return number of rows removed. + with get_db() as db: - return True + db.query(Chat).filter_by(user_id=f"shared-{chat_id}").delete() + db.commit() + + return True except: return False @@ -167,56 +182,50 @@ class ChatTable: self, id: str, share_id: Optional[str] ) -> Optional[ChatModel]: try: - query = Chat.update( - share_id=share_id, - ).where(Chat.id == id) - query.execute() + with get_db() as db: - chat = Chat.get(Chat.id == id) - return ChatModel(**model_to_dict(chat)) + chat = db.get(Chat, id) + chat.share_id = share_id + db.commit() + db.refresh(chat) + return ChatModel.model_validate(chat) except: return None def toggle_chat_archive_by_id(self, id: str) -> Optional[ChatModel]: try: - chat = self.get_chat_by_id(id) - query = Chat.update( - archived=(not chat.archived), - ).where(Chat.id == id) + with get_db() as db: - query.execute() - - chat = Chat.get(Chat.id == id) - return ChatModel(**model_to_dict(chat)) + chat = db.get(Chat, id) + chat.archived = not chat.archived + db.commit() + db.refresh(chat) + return ChatModel.model_validate(chat) except: return None def archive_all_chats_by_user_id(self, user_id: str) -> bool: try: - chats = self.get_chats_by_user_id(user_id) - for chat in chats: - query = Chat.update( - archived=True, - ).where(Chat.id == chat.id) - - query.execute() - - return True + with get_db() as db: + db.query(Chat).filter_by(user_id=user_id).update({"archived": True}) + db.commit() + return True except: return False def get_archived_chat_list_by_user_id( self, user_id: str, skip: int = 0, limit: int = 50 ) -> List[ChatModel]: - return [ - ChatModel(**model_to_dict(chat)) - for chat in Chat.select() - .where(Chat.archived == True) - .where(Chat.user_id == user_id) - .order_by(Chat.updated_at.desc()) - # .limit(limit) - # .offset(skip) - ] + with get_db() as db: + + all_chats = ( + db.query(Chat) + .filter_by(user_id=user_id, archived=True) + .order_by(Chat.updated_at.desc()) + # .limit(limit).offset(skip) + .all() + ) + return [ChatModel.model_validate(chat) for chat in all_chats] def get_chat_list_by_user_id( self, @@ -225,131 +234,141 @@ class ChatTable: skip: int = 0, limit: int = 50, ) -> List[ChatModel]: - if include_archived: - return [ - ChatModel(**model_to_dict(chat)) - for chat in Chat.select() - .where(Chat.user_id == user_id) - .order_by(Chat.updated_at.desc()) - # .limit(limit) - # .offset(skip) - ] - else: - return [ - ChatModel(**model_to_dict(chat)) - for chat in Chat.select() - .where(Chat.archived == False) - .where(Chat.user_id == user_id) - .order_by(Chat.updated_at.desc()) - # .limit(limit) - # .offset(skip) - ] + with get_db() as db: + query = db.query(Chat).filter_by(user_id=user_id) + if not include_archived: + query = query.filter_by(archived=False) + all_chats = ( + query.order_by(Chat.updated_at.desc()) + # .limit(limit).offset(skip) + .all() + ) + return [ChatModel.model_validate(chat) for chat in all_chats] def get_chat_list_by_chat_ids( self, chat_ids: List[str], skip: int = 0, limit: int = 50 ) -> List[ChatModel]: - return [ - ChatModel(**model_to_dict(chat)) - for chat in Chat.select() - .where(Chat.archived == False) - .where(Chat.id.in_(chat_ids)) - .order_by(Chat.updated_at.desc()) - ] + with get_db() as db: + all_chats = ( + db.query(Chat) + .filter(Chat.id.in_(chat_ids)) + .filter_by(archived=False) + .order_by(Chat.updated_at.desc()) + .all() + ) + return [ChatModel.model_validate(chat) for chat in all_chats] def get_chat_by_id(self, id: str) -> Optional[ChatModel]: try: - chat = Chat.get(Chat.id == id) - return ChatModel(**model_to_dict(chat)) + with get_db() as db: + + chat = db.get(Chat, id) + return ChatModel.model_validate(chat) except: return None def get_chat_by_share_id(self, id: str) -> Optional[ChatModel]: try: - chat = Chat.get(Chat.share_id == id) + with get_db() as db: - if chat: - chat = Chat.get(Chat.id == id) - return ChatModel(**model_to_dict(chat)) - else: - return None - except: + chat = db.query(Chat).filter_by(share_id=id).first() + + if chat: + return self.get_chat_by_id(id) + else: + return None + except Exception as e: return None def get_chat_by_id_and_user_id(self, id: str, user_id: str) -> Optional[ChatModel]: try: - chat = Chat.get(Chat.id == id, Chat.user_id == user_id) - return ChatModel(**model_to_dict(chat)) + with get_db() as db: + + chat = db.query(Chat).filter_by(id=id, user_id=user_id).first() + return ChatModel.model_validate(chat) except: return None def get_chats(self, skip: int = 0, limit: int = 50) -> List[ChatModel]: - return [ - ChatModel(**model_to_dict(chat)) - for chat in Chat.select().order_by(Chat.updated_at.desc()) - # .limit(limit).offset(skip) - ] + with get_db() as db: + + all_chats = ( + db.query(Chat) + # .limit(limit).offset(skip) + .order_by(Chat.updated_at.desc()) + ) + return [ChatModel.model_validate(chat) for chat in all_chats] def get_chats_by_user_id(self, user_id: str) -> List[ChatModel]: - return [ - ChatModel(**model_to_dict(chat)) - for chat in Chat.select() - .where(Chat.user_id == user_id) - .order_by(Chat.updated_at.desc()) - # .limit(limit).offset(skip) - ] + with get_db() as db: + + all_chats = ( + db.query(Chat) + .filter_by(user_id=user_id) + .order_by(Chat.updated_at.desc()) + ) + return [ChatModel.model_validate(chat) for chat in all_chats] def get_archived_chats_by_user_id(self, user_id: str) -> List[ChatModel]: - return [ - ChatModel(**model_to_dict(chat)) - for chat in Chat.select() - .where(Chat.archived == True) - .where(Chat.user_id == user_id) - .order_by(Chat.updated_at.desc()) - ] + with get_db() as db: + + all_chats = ( + db.query(Chat) + .filter_by(user_id=user_id, archived=True) + .order_by(Chat.updated_at.desc()) + ) + return [ChatModel.model_validate(chat) for chat in all_chats] def delete_chat_by_id(self, id: str) -> bool: try: - query = Chat.delete().where((Chat.id == id)) - query.execute() # Remove the rows, return number of rows removed. + with get_db() as db: - return True and self.delete_shared_chat_by_chat_id(id) + db.query(Chat).filter_by(id=id).delete() + db.commit() + + return True and self.delete_shared_chat_by_chat_id(id) except: return False def delete_chat_by_id_and_user_id(self, id: str, user_id: str) -> bool: try: - query = Chat.delete().where((Chat.id == id) & (Chat.user_id == user_id)) - query.execute() # Remove the rows, return number of rows removed. + with get_db() as db: - return True and self.delete_shared_chat_by_chat_id(id) + db.query(Chat).filter_by(id=id, user_id=user_id).delete() + db.commit() + + return True and self.delete_shared_chat_by_chat_id(id) except: return False def delete_chats_by_user_id(self, user_id: str) -> bool: try: - self.delete_shared_chats_by_user_id(user_id) + with get_db() as db: - query = Chat.delete().where(Chat.user_id == user_id) - query.execute() # Remove the rows, return number of rows removed. + self.delete_shared_chats_by_user_id(user_id) - return True + db.query(Chat).filter_by(user_id=user_id).delete() + db.commit() + + return True except: return False def delete_shared_chats_by_user_id(self, user_id: str) -> bool: try: - shared_chat_ids = [ - f"shared-{chat.id}" - for chat in Chat.select().where(Chat.user_id == user_id) - ] - query = Chat.delete().where(Chat.user_id << shared_chat_ids) - query.execute() # Remove the rows, return number of rows removed. + with get_db() as db: - return True + chats_by_user = db.query(Chat).filter_by(user_id=user_id).all() + shared_chat_ids = [f"shared-{chat.id}" for chat in chats_by_user] + + db.query(Chat).filter(Chat.user_id.in_(shared_chat_ids)).delete() + db.commit() + + return True except: return False -Chats = ChatTable(DB) +Chats = ChatTable() diff --git a/backend/apps/webui/models/documents.py b/backend/apps/webui/models/documents.py index 3b730535f..ac8655da9 100644 --- a/backend/apps/webui/models/documents.py +++ b/backend/apps/webui/models/documents.py @@ -1,14 +1,11 @@ -from pydantic import BaseModel -from peewee import * -from playhouse.shortcuts import model_to_dict -from typing import List, Union, Optional +from pydantic import BaseModel, ConfigDict +from typing import List, Optional import time import logging -from utils.utils import decode_token -from utils.misc import get_gravatar_url +from sqlalchemy import String, Column, BigInteger, Text -from apps.webui.internal.db import DB +from apps.webui.internal.db import Base, get_db import json @@ -22,20 +19,21 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"]) #################### -class Document(Model): - collection_name = CharField(unique=True) - name = CharField(unique=True) - title = TextField() - filename = TextField() - content = TextField(null=True) - user_id = CharField() - timestamp = BigIntegerField() +class Document(Base): + __tablename__ = "document" - class Meta: - database = DB + collection_name = Column(String, primary_key=True) + name = Column(String, unique=True) + title = Column(Text) + filename = Column(Text) + content = Column(Text, nullable=True) + user_id = Column(String) + timestamp = Column(BigInteger) class DocumentModel(BaseModel): + model_config = ConfigDict(from_attributes=True) + collection_name: str name: str title: str @@ -72,57 +70,63 @@ class DocumentForm(DocumentUpdateForm): class DocumentsTable: - def __init__(self, db): - self.db = db - self.db.create_tables([Document]) def insert_new_doc( self, user_id: str, form_data: DocumentForm ) -> Optional[DocumentModel]: - document = DocumentModel( - **{ - **form_data.model_dump(), - "user_id": user_id, - "timestamp": int(time.time()), - } - ) + with get_db() as db: - try: - result = Document.create(**document.model_dump()) - if result: - return document - else: + document = DocumentModel( + **{ + **form_data.model_dump(), + "user_id": user_id, + "timestamp": int(time.time()), + } + ) + + try: + result = Document(**document.model_dump()) + db.add(result) + db.commit() + db.refresh(result) + if result: + return DocumentModel.model_validate(result) + else: + return None + except: return None - except: - return None def get_doc_by_name(self, name: str) -> Optional[DocumentModel]: try: - document = Document.get(Document.name == name) - return DocumentModel(**model_to_dict(document)) + with get_db() as db: + + document = db.query(Document).filter_by(name=name).first() + return DocumentModel.model_validate(document) if document else None except: return None def get_docs(self) -> List[DocumentModel]: - return [ - DocumentModel(**model_to_dict(doc)) - for doc in Document.select() - # .limit(limit).offset(skip) - ] + with get_db() as db: + + return [ + DocumentModel.model_validate(doc) for doc in db.query(Document).all() + ] def update_doc_by_name( self, name: str, form_data: DocumentUpdateForm ) -> Optional[DocumentModel]: try: - query = Document.update( - title=form_data.title, - name=form_data.name, - timestamp=int(time.time()), - ).where(Document.name == name) - query.execute() + with get_db() as db: - doc = Document.get(Document.name == form_data.name) - return DocumentModel(**model_to_dict(doc)) + db.query(Document).filter_by(name=name).update( + { + "title": form_data.title, + "name": form_data.name, + "timestamp": int(time.time()), + } + ) + db.commit() + return self.get_doc_by_name(form_data.name) except Exception as e: log.exception(e) return None @@ -135,26 +139,29 @@ class DocumentsTable: doc_content = json.loads(doc.content if doc.content else "{}") doc_content = {**doc_content, **updated} - query = Document.update( - content=json.dumps(doc_content), - timestamp=int(time.time()), - ).where(Document.name == name) - query.execute() + with get_db() as db: - doc = Document.get(Document.name == name) - return DocumentModel(**model_to_dict(doc)) + db.query(Document).filter_by(name=name).update( + { + "content": json.dumps(doc_content), + "timestamp": int(time.time()), + } + ) + db.commit() + return self.get_doc_by_name(name) except Exception as e: log.exception(e) return None def delete_doc_by_name(self, name: str) -> bool: try: - query = Document.delete().where((Document.name == name)) - query.execute() # Remove the rows, return number of rows removed. + with get_db() as db: - return True + db.query(Document).filter_by(name=name).delete() + db.commit() + return True except: return False -Documents = DocumentsTable(DB) +Documents = DocumentsTable() diff --git a/backend/apps/webui/models/files.py b/backend/apps/webui/models/files.py index 6459ad725..16272f24a 100644 --- a/backend/apps/webui/models/files.py +++ b/backend/apps/webui/models/files.py @@ -1,10 +1,11 @@ -from pydantic import BaseModel -from peewee import * -from playhouse.shortcuts import model_to_dict +from pydantic import BaseModel, ConfigDict from typing import List, Union, Optional import time import logging -from apps.webui.internal.db import DB, JSONField + +from sqlalchemy import Column, String, BigInteger, Text + +from apps.webui.internal.db import JSONField, Base, get_db import json @@ -18,15 +19,14 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"]) #################### -class File(Model): - id = CharField(unique=True) - user_id = CharField() - filename = TextField() - meta = JSONField() - created_at = BigIntegerField() +class File(Base): + __tablename__ = "file" - class Meta: - database = DB + id = Column(String, primary_key=True) + user_id = Column(String) + filename = Column(Text) + meta = Column(JSONField) + created_at = Column(BigInteger) class FileModel(BaseModel): @@ -36,6 +36,8 @@ class FileModel(BaseModel): meta: dict created_at: int # timestamp in epoch + model_config = ConfigDict(from_attributes=True) + #################### # Forms @@ -57,56 +59,68 @@ class FileForm(BaseModel): class FilesTable: - def __init__(self, db): - self.db = db - self.db.create_tables([File]) def insert_new_file(self, user_id: str, form_data: FileForm) -> Optional[FileModel]: - file = FileModel( - **{ - **form_data.model_dump(), - "user_id": user_id, - "created_at": int(time.time()), - } - ) + with get_db() as db: - try: - result = File.create(**file.model_dump()) - if result: - return file - else: + file = FileModel( + **{ + **form_data.model_dump(), + "user_id": user_id, + "created_at": int(time.time()), + } + ) + + try: + result = File(**file.model_dump()) + db.add(result) + db.commit() + db.refresh(result) + if result: + return FileModel.model_validate(result) + else: + return None + except Exception as e: + print(f"Error creating tool: {e}") return None - except Exception as e: - print(f"Error creating tool: {e}") - return None def get_file_by_id(self, id: str) -> Optional[FileModel]: - try: - file = File.get(File.id == id) - return FileModel(**model_to_dict(file)) - except: - return None + with get_db() as db: + + try: + file = db.get(File, id) + return FileModel.model_validate(file) + except: + return None def get_files(self) -> List[FileModel]: - return [FileModel(**model_to_dict(file)) for file in File.select()] + with get_db() as db: + + return [FileModel.model_validate(file) for file in db.query(File).all()] def delete_file_by_id(self, id: str) -> bool: - try: - query = File.delete().where((File.id == id)) - query.execute() # Remove the rows, return number of rows removed. - return True - except: - return False + with get_db() as db: + + try: + db.query(File).filter_by(id=id).delete() + db.commit() + + return True + except: + return False def delete_all_files(self) -> bool: - try: - query = File.delete() - query.execute() # Remove the rows, return number of rows removed. - return True - except: - return False + with get_db() as db: + + try: + db.query(File).delete() + db.commit() + + return True + except: + return False -Files = FilesTable(DB) +Files = FilesTable() diff --git a/backend/apps/webui/models/functions.py b/backend/apps/webui/models/functions.py index 2cace54c4..907576b80 100644 --- a/backend/apps/webui/models/functions.py +++ b/backend/apps/webui/models/functions.py @@ -1,10 +1,11 @@ -from pydantic import BaseModel -from peewee import * -from playhouse.shortcuts import model_to_dict +from pydantic import BaseModel, ConfigDict from typing import List, Union, Optional import time import logging -from apps.webui.internal.db import DB, JSONField + +from sqlalchemy import Column, String, Text, BigInteger, Boolean + +from apps.webui.internal.db import JSONField, Base, get_db from apps.webui.models.users import Users import json @@ -21,21 +22,20 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"]) #################### -class Function(Model): - id = CharField(unique=True) - user_id = CharField() - name = TextField() - type = TextField() - content = TextField() - meta = JSONField() - valves = JSONField() - is_active = BooleanField(default=False) - is_global = BooleanField(default=False) - updated_at = BigIntegerField() - created_at = BigIntegerField() +class Function(Base): + __tablename__ = "function" - class Meta: - database = DB + id = Column(String, primary_key=True) + user_id = Column(String) + name = Column(Text) + type = Column(Text) + content = Column(Text) + meta = Column(JSONField) + valves = Column(JSONField) + is_active = Column(Boolean) + is_global = Column(Boolean) + updated_at = Column(BigInteger) + created_at = Column(BigInteger) class FunctionMeta(BaseModel): @@ -55,6 +55,8 @@ class FunctionModel(BaseModel): updated_at: int # timestamp in epoch created_at: int # timestamp in epoch + model_config = ConfigDict(from_attributes=True) + #################### # Forms @@ -85,13 +87,11 @@ class FunctionValves(BaseModel): class FunctionsTable: - def __init__(self, db): - self.db = db - self.db.create_tables([Function]) def insert_new_function( self, user_id: str, type: str, form_data: FunctionForm ) -> Optional[FunctionModel]: + function = FunctionModel( **{ **form_data.model_dump(), @@ -103,89 +103,102 @@ class FunctionsTable: ) try: - result = Function.create(**function.model_dump()) - if result: - return function - else: - return None + with get_db() as db: + result = Function(**function.model_dump()) + db.add(result) + db.commit() + db.refresh(result) + if result: + return FunctionModel.model_validate(result) + else: + return None except Exception as e: print(f"Error creating tool: {e}") return None def get_function_by_id(self, id: str) -> Optional[FunctionModel]: try: - function = Function.get(Function.id == id) - return FunctionModel(**model_to_dict(function)) + with get_db() as db: + + function = db.get(Function, id) + return FunctionModel.model_validate(function) except: return None def get_functions(self, active_only=False) -> List[FunctionModel]: - if active_only: - return [ - FunctionModel(**model_to_dict(function)) - for function in Function.select().where(Function.is_active == True) - ] - else: - return [ - FunctionModel(**model_to_dict(function)) - for function in Function.select() - ] + with get_db() as db: + + if active_only: + return [ + FunctionModel.model_validate(function) + for function in db.query(Function).filter_by(is_active=True).all() + ] + else: + return [ + FunctionModel.model_validate(function) + for function in db.query(Function).all() + ] def get_functions_by_type( self, type: str, active_only=False ) -> List[FunctionModel]: - if active_only: - return [ - FunctionModel(**model_to_dict(function)) - for function in Function.select().where( - Function.type == type, Function.is_active == True - ) - ] - else: - return [ - FunctionModel(**model_to_dict(function)) - for function in Function.select().where(Function.type == type) - ] + with get_db() as db: + + if active_only: + return [ + FunctionModel.model_validate(function) + for function in db.query(Function) + .filter_by(type=type, is_active=True) + .all() + ] + else: + return [ + FunctionModel.model_validate(function) + for function in db.query(Function).filter_by(type=type).all() + ] def get_global_filter_functions(self) -> List[FunctionModel]: - return [ - FunctionModel(**model_to_dict(function)) - for function in Function.select().where( - Function.type == "filter", - Function.is_active == True, - Function.is_global == True, - ) - ] + with get_db() as db: + + return [ + FunctionModel.model_validate(function) + for function in db.query(Function) + .filter_by(type="filter", is_active=True, is_global=True) + .all() + ] def get_function_valves_by_id(self, id: str) -> Optional[dict]: - try: - function = Function.get(Function.id == id) - return function.valves if function.valves else {} - except Exception as e: - print(f"An error occurred: {e}") - return None + with get_db() as db: + + try: + function = db.get(Function, id) + return function.valves if function.valves else {} + except Exception as e: + print(f"An error occurred: {e}") + return None def update_function_valves_by_id( self, id: str, valves: dict ) -> Optional[FunctionValves]: - try: - query = Function.update( - **{"valves": valves}, - updated_at=int(time.time()), - ).where(Function.id == id) - query.execute() + with get_db() as db: - function = Function.get(Function.id == id) - return FunctionValves(**model_to_dict(function)) - except: - return None + try: + function = db.get(Function, id) + function.valves = valves + function.updated_at = int(time.time()) + db.commit() + db.refresh(function) + return self.get_function_by_id(id) + except: + return None def get_user_valves_by_id_and_user_id( self, id: str, user_id: str ) -> Optional[dict]: + try: user = Users.get_user_by_id(user_id) - user_settings = user.settings.model_dump() + user_settings = user.settings.model_dump() if user.settings else {} # Check if user has "functions" and "valves" settings if "functions" not in user_settings: @@ -201,9 +214,10 @@ class FunctionsTable: def update_user_valves_by_id_and_user_id( self, id: str, user_id: str, valves: dict ) -> Optional[dict]: + try: user = Users.get_user_by_id(user_id) - user_settings = user.settings.model_dump() + user_settings = user.settings.model_dump() if user.settings else {} # Check if user has "functions" and "valves" settings if "functions" not in user_settings: @@ -214,8 +228,7 @@ class FunctionsTable: user_settings["functions"]["valves"][id] = valves # Update the user settings in the database - query = Users.update_user_by_id(user_id, {"settings": user_settings}) - query.execute() + Users.update_user_by_id(user_id, {"settings": user_settings}) return user_settings["functions"]["valves"][id] except Exception as e: @@ -223,39 +236,44 @@ class FunctionsTable: return None def update_function_by_id(self, id: str, updated: dict) -> Optional[FunctionModel]: - try: - query = Function.update( - **updated, - updated_at=int(time.time()), - ).where(Function.id == id) - query.execute() + with get_db() as db: - function = Function.get(Function.id == id) - return FunctionModel(**model_to_dict(function)) - except: - return None + try: + db.query(Function).filter_by(id=id).update( + { + **updated, + "updated_at": int(time.time()), + } + ) + db.commit() + return self.get_function_by_id(id) + except: + return None def deactivate_all_functions(self) -> Optional[bool]: - try: - query = Function.update( - **{"is_active": False}, - updated_at=int(time.time()), - ) + with get_db() as db: - query.execute() - - return True - except: - return None + try: + db.query(Function).update( + { + "is_active": False, + "updated_at": int(time.time()), + } + ) + db.commit() + return True + except: + return None def delete_function_by_id(self, id: str) -> bool: - try: - query = Function.delete().where((Function.id == id)) - query.execute() # Remove the rows, return number of rows removed. + with get_db() as db: + try: + db.query(Function).filter_by(id=id).delete() + db.commit() - return True - except: - return False + return True + except: + return False -Functions = FunctionsTable(DB) +Functions = FunctionsTable() diff --git a/backend/apps/webui/models/memories.py b/backend/apps/webui/models/memories.py index ef63674ab..02d4b6924 100644 --- a/backend/apps/webui/models/memories.py +++ b/backend/apps/webui/models/memories.py @@ -1,10 +1,9 @@ -from pydantic import BaseModel -from peewee import * -from playhouse.shortcuts import model_to_dict +from pydantic import BaseModel, ConfigDict from typing import List, Union, Optional -from apps.webui.internal.db import DB -from apps.webui.models.chats import Chats +from sqlalchemy import Column, String, BigInteger, Text + +from apps.webui.internal.db import Base, get_db import time import uuid @@ -14,15 +13,14 @@ import uuid #################### -class Memory(Model): - id = CharField(unique=True) - user_id = CharField() - content = TextField() - updated_at = BigIntegerField() - created_at = BigIntegerField() +class Memory(Base): + __tablename__ = "memory" - class Meta: - database = DB + id = Column(String, primary_key=True) + user_id = Column(String) + content = Column(Text) + updated_at = Column(BigInteger) + created_at = Column(BigInteger) class MemoryModel(BaseModel): @@ -32,6 +30,8 @@ class MemoryModel(BaseModel): updated_at: int # timestamp in epoch created_at: int # timestamp in epoch + model_config = ConfigDict(from_attributes=True) + #################### # Forms @@ -39,94 +39,110 @@ class MemoryModel(BaseModel): class MemoriesTable: - def __init__(self, db): - self.db = db - self.db.create_tables([Memory]) def insert_new_memory( self, user_id: str, content: str, ) -> Optional[MemoryModel]: - id = str(uuid.uuid4()) - memory = MemoryModel( - **{ - "id": id, - "user_id": user_id, - "content": content, - "created_at": int(time.time()), - "updated_at": int(time.time()), - } - ) - result = Memory.create(**memory.model_dump()) - if result: - return memory - else: - return None + with get_db() as db: + id = str(uuid.uuid4()) + + memory = MemoryModel( + **{ + "id": id, + "user_id": user_id, + "content": content, + "created_at": int(time.time()), + "updated_at": int(time.time()), + } + ) + result = Memory(**memory.model_dump()) + db.add(result) + db.commit() + db.refresh(result) + if result: + return MemoryModel.model_validate(result) + else: + return None def update_memory_by_id( self, id: str, content: str, ) -> Optional[MemoryModel]: - try: - memory = Memory.get(Memory.id == id) - memory.content = content - memory.updated_at = int(time.time()) - memory.save() - return MemoryModel(**model_to_dict(memory)) - except: - return None + with get_db() as db: + + try: + db.query(Memory).filter_by(id=id).update( + {"content": content, "updated_at": int(time.time())} + ) + db.commit() + return self.get_memory_by_id(id) + except: + return None def get_memories(self) -> List[MemoryModel]: - try: - memories = Memory.select() - return [MemoryModel(**model_to_dict(memory)) for memory in memories] - except: - return None + with get_db() as db: + + try: + memories = db.query(Memory).all() + return [MemoryModel.model_validate(memory) for memory in memories] + except: + return None def get_memories_by_user_id(self, user_id: str) -> List[MemoryModel]: - try: - memories = Memory.select().where(Memory.user_id == user_id) - return [MemoryModel(**model_to_dict(memory)) for memory in memories] - except: - return None + with get_db() as db: - def get_memory_by_id(self, id) -> Optional[MemoryModel]: - try: - memory = Memory.get(Memory.id == id) - return MemoryModel(**model_to_dict(memory)) - except: - return None + try: + memories = db.query(Memory).filter_by(user_id=user_id).all() + return [MemoryModel.model_validate(memory) for memory in memories] + except: + return None + + def get_memory_by_id(self, id: str) -> Optional[MemoryModel]: + with get_db() as db: + + try: + memory = db.get(Memory, id) + return MemoryModel.model_validate(memory) + except: + return None def delete_memory_by_id(self, id: str) -> bool: - try: - query = Memory.delete().where(Memory.id == id) - query.execute() # Remove the rows, return number of rows removed. + with get_db() as db: - return True + try: + db.query(Memory).filter_by(id=id).delete() + db.commit() - except: - return False + return True + + except: + return False def delete_memories_by_user_id(self, user_id: str) -> bool: - try: - query = Memory.delete().where(Memory.user_id == user_id) - query.execute() + with get_db() as db: - return True - except: - return False + try: + db.query(Memory).filter_by(user_id=user_id).delete() + db.commit() + + return True + except: + return False def delete_memory_by_id_and_user_id(self, id: str, user_id: str) -> bool: - try: - query = Memory.delete().where(Memory.id == id, Memory.user_id == user_id) - query.execute() + with get_db() as db: - return True - except: - return False + try: + db.query(Memory).filter_by(id=id, user_id=user_id).delete() + db.commit() + + return True + except: + return False -Memories = MemoriesTable(DB) +Memories = MemoriesTable() diff --git a/backend/apps/webui/models/models.py b/backend/apps/webui/models/models.py index 851352398..3b128c7d6 100644 --- a/backend/apps/webui/models/models.py +++ b/backend/apps/webui/models/models.py @@ -2,13 +2,10 @@ import json import logging from typing import Optional -import peewee as pw -from peewee import * - -from playhouse.shortcuts import model_to_dict from pydantic import BaseModel, ConfigDict +from sqlalchemy import String, Column, BigInteger, Text -from apps.webui.internal.db import DB, JSONField +from apps.webui.internal.db import Base, JSONField, get_db from typing import List, Union, Optional from config import SRC_LOG_LEVELS @@ -32,7 +29,7 @@ class ModelParams(BaseModel): # ModelMeta is a model for the data stored in the meta field of the Model table class ModelMeta(BaseModel): - profile_image_url: Optional[str] = "/favicon.png" + profile_image_url: Optional[str] = "/static/favicon.png" description: Optional[str] = None """ @@ -46,38 +43,37 @@ class ModelMeta(BaseModel): pass -class Model(pw.Model): - id = pw.TextField(unique=True) +class Model(Base): + __tablename__ = "model" + + id = Column(Text, primary_key=True) """ The model's id as used in the API. If set to an existing model, it will override the model. """ - user_id = pw.TextField() + user_id = Column(Text) - base_model_id = pw.TextField(null=True) + base_model_id = Column(Text, nullable=True) """ An optional pointer to the actual model that should be used when proxying requests. """ - name = pw.TextField() + name = Column(Text) """ The human-readable display name of the model. """ - params = JSONField() + params = Column(JSONField) """ Holds a JSON encoded blob of parameters, see `ModelParams`. """ - meta = JSONField() + meta = Column(JSONField) """ Holds a JSON encoded blob of metadata, see `ModelMeta`. """ - updated_at = BigIntegerField() - created_at = BigIntegerField() - - class Meta: - database = DB + updated_at = Column(BigInteger) + created_at = Column(BigInteger) class ModelModel(BaseModel): @@ -92,6 +88,8 @@ class ModelModel(BaseModel): updated_at: int # timestamp in epoch created_at: int # timestamp in epoch + model_config = ConfigDict(from_attributes=True) + #################### # Forms @@ -115,12 +113,6 @@ class ModelForm(BaseModel): class ModelsTable: - def __init__( - self, - db: pw.SqliteDatabase | pw.PostgresqlDatabase, - ): - self.db = db - self.db.create_tables([Model]) def insert_new_model( self, form_data: ModelForm, user_id: str @@ -134,34 +126,50 @@ class ModelsTable: } ) try: - result = Model.create(**model.model_dump()) - if result: - return model - else: - return None + with get_db() as db: + + result = Model(**model.model_dump()) + db.add(result) + db.commit() + db.refresh(result) + + if result: + return ModelModel.model_validate(result) + else: + return None except Exception as e: print(e) return None def get_all_models(self) -> List[ModelModel]: - return [ModelModel(**model_to_dict(model)) for model in Model.select()] + with get_db() as db: + + return [ModelModel.model_validate(model) for model in db.query(Model).all()] def get_model_by_id(self, id: str) -> Optional[ModelModel]: try: - model = Model.get(Model.id == id) - return ModelModel(**model_to_dict(model)) + with get_db() as db: + + model = db.get(Model, id) + return ModelModel.model_validate(model) except: return None def update_model_by_id(self, id: str, model: ModelForm) -> Optional[ModelModel]: try: - # update only the fields that are present in the model - query = Model.update(**model.model_dump()).where(Model.id == id) - query.execute() + with get_db() as db: + # update only the fields that are present in the model + result = ( + db.query(Model) + .filter_by(id=id) + .update(model.model_dump(exclude={"id"}, exclude_none=True)) + ) + db.commit() - model = Model.get(Model.id == id) - return ModelModel(**model_to_dict(model)) + model = db.get(Model, id) + db.refresh(model) + return ModelModel.model_validate(model) except Exception as e: print(e) @@ -169,11 +177,14 @@ class ModelsTable: def delete_model_by_id(self, id: str) -> bool: try: - query = Model.delete().where(Model.id == id) - query.execute() - return True + with get_db() as db: + + db.query(Model).filter_by(id=id).delete() + db.commit() + + return True except: return False -Models = ModelsTable(DB) +Models = ModelsTable() diff --git a/backend/apps/webui/models/prompts.py b/backend/apps/webui/models/prompts.py index c4ac6be14..b8467b631 100644 --- a/backend/apps/webui/models/prompts.py +++ b/backend/apps/webui/models/prompts.py @@ -1,13 +1,10 @@ -from pydantic import BaseModel -from peewee import * -from playhouse.shortcuts import model_to_dict -from typing import List, Union, Optional +from pydantic import BaseModel, ConfigDict +from typing import List, Optional import time -from utils.utils import decode_token -from utils.misc import get_gravatar_url +from sqlalchemy import String, Column, BigInteger, Text -from apps.webui.internal.db import DB +from apps.webui.internal.db import Base, get_db import json @@ -16,15 +13,14 @@ import json #################### -class Prompt(Model): - command = CharField(unique=True) - user_id = CharField() - title = TextField() - content = TextField() - timestamp = BigIntegerField() +class Prompt(Base): + __tablename__ = "prompt" - class Meta: - database = DB + command = Column(String, primary_key=True) + user_id = Column(String) + title = Column(Text) + content = Column(Text) + timestamp = Column(BigInteger) class PromptModel(BaseModel): @@ -34,6 +30,8 @@ class PromptModel(BaseModel): content: str timestamp: int # timestamp in epoch + model_config = ConfigDict(from_attributes=True) + #################### # Forms @@ -48,10 +46,6 @@ class PromptForm(BaseModel): class PromptsTable: - def __init__(self, db): - self.db = db - self.db.create_tables([Prompt]) - def insert_new_prompt( self, user_id: str, form_data: PromptForm ) -> Optional[PromptModel]: @@ -66,53 +60,60 @@ class PromptsTable: ) try: - result = Prompt.create(**prompt.model_dump()) - if result: - return prompt - else: - return None - except: + with get_db() as db: + + result = Prompt(**prompt.dict()) + db.add(result) + db.commit() + db.refresh(result) + if result: + return PromptModel.model_validate(result) + else: + return None + except Exception as e: return None def get_prompt_by_command(self, command: str) -> Optional[PromptModel]: try: - prompt = Prompt.get(Prompt.command == command) - return PromptModel(**model_to_dict(prompt)) + with get_db() as db: + + prompt = db.query(Prompt).filter_by(command=command).first() + return PromptModel.model_validate(prompt) except: return None def get_prompts(self) -> List[PromptModel]: - return [ - PromptModel(**model_to_dict(prompt)) - for prompt in Prompt.select() - # .limit(limit).offset(skip) - ] + with get_db() as db: + + return [ + PromptModel.model_validate(prompt) for prompt in db.query(Prompt).all() + ] def update_prompt_by_command( self, command: str, form_data: PromptForm ) -> Optional[PromptModel]: try: - query = Prompt.update( - title=form_data.title, - content=form_data.content, - timestamp=int(time.time()), - ).where(Prompt.command == command) + with get_db() as db: - query.execute() - - prompt = Prompt.get(Prompt.command == command) - return PromptModel(**model_to_dict(prompt)) + prompt = db.query(Prompt).filter_by(command=command).first() + prompt.title = form_data.title + prompt.content = form_data.content + prompt.timestamp = int(time.time()) + db.commit() + return PromptModel.model_validate(prompt) except: return None def delete_prompt_by_command(self, command: str) -> bool: try: - query = Prompt.delete().where((Prompt.command == command)) - query.execute() # Remove the rows, return number of rows removed. + with get_db() as db: - return True + db.query(Prompt).filter_by(command=command).delete() + db.commit() + + return True except: return False -Prompts = PromptsTable(DB) +Prompts = PromptsTable() diff --git a/backend/apps/webui/models/tags.py b/backend/apps/webui/models/tags.py index 4c4fa82e6..bb98d3a38 100644 --- a/backend/apps/webui/models/tags.py +++ b/backend/apps/webui/models/tags.py @@ -1,14 +1,14 @@ -from pydantic import BaseModel -from typing import List, Union, Optional -from peewee import * -from playhouse.shortcuts import model_to_dict +from pydantic import BaseModel, ConfigDict +from typing import List, Optional import json import uuid import time import logging -from apps.webui.internal.db import DB +from sqlalchemy import String, Column, BigInteger, Text + +from apps.webui.internal.db import Base, get_db from config import SRC_LOG_LEVELS @@ -20,25 +20,23 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"]) #################### -class Tag(Model): - id = CharField(unique=True) - name = CharField() - user_id = CharField() - data = TextField(null=True) +class Tag(Base): + __tablename__ = "tag" - class Meta: - database = DB + id = Column(String, primary_key=True) + name = Column(String) + user_id = Column(String) + data = Column(Text, nullable=True) -class ChatIdTag(Model): - id = CharField(unique=True) - tag_name = CharField() - chat_id = CharField() - user_id = CharField() - timestamp = BigIntegerField() +class ChatIdTag(Base): + __tablename__ = "chatidtag" - class Meta: - database = DB + id = Column(String, primary_key=True) + tag_name = Column(String) + chat_id = Column(String) + user_id = Column(String) + timestamp = Column(BigInteger) class TagModel(BaseModel): @@ -47,6 +45,8 @@ class TagModel(BaseModel): user_id: str data: Optional[str] = None + model_config = ConfigDict(from_attributes=True) + class ChatIdTagModel(BaseModel): id: str @@ -55,6 +55,8 @@ class ChatIdTagModel(BaseModel): user_id: str timestamp: int + model_config = ConfigDict(from_attributes=True) + #################### # Forms @@ -75,28 +77,31 @@ class ChatTagsResponse(BaseModel): class TagTable: - def __init__(self, db): - self.db = db - db.create_tables([Tag, ChatIdTag]) def insert_new_tag(self, name: str, user_id: str) -> Optional[TagModel]: - id = str(uuid.uuid4()) - tag = TagModel(**{"id": id, "user_id": user_id, "name": name}) - try: - result = Tag.create(**tag.model_dump()) - if result: - return tag - else: + with get_db() as db: + + id = str(uuid.uuid4()) + tag = TagModel(**{"id": id, "user_id": user_id, "name": name}) + try: + result = Tag(**tag.model_dump()) + db.add(result) + db.commit() + db.refresh(result) + if result: + return TagModel.model_validate(result) + else: + return None + except Exception as e: return None - except Exception as e: - return None def get_tag_by_name_and_user_id( self, name: str, user_id: str ) -> Optional[TagModel]: try: - tag = Tag.get(Tag.name == name, Tag.user_id == user_id) - return TagModel(**model_to_dict(tag)) + with get_db() as db: + tag = db.query(Tag).filter(name=name, user_id=user_id).first() + return TagModel.model_validate(tag) except Exception as e: return None @@ -118,82 +123,110 @@ class TagTable: } ) try: - result = ChatIdTag.create(**chatIdTag.model_dump()) - if result: - return chatIdTag - else: - return None + with get_db() as db: + result = ChatIdTag(**chatIdTag.model_dump()) + db.add(result) + db.commit() + db.refresh(result) + if result: + return ChatIdTagModel.model_validate(result) + else: + return None except: return None def get_tags_by_user_id(self, user_id: str) -> List[TagModel]: - tag_names = [ - ChatIdTagModel(**model_to_dict(chat_id_tag)).tag_name - for chat_id_tag in ChatIdTag.select() - .where(ChatIdTag.user_id == user_id) - .order_by(ChatIdTag.timestamp.desc()) - ] + with get_db() as db: + tag_names = [ + chat_id_tag.tag_name + for chat_id_tag in ( + db.query(ChatIdTag) + .filter_by(user_id=user_id) + .order_by(ChatIdTag.timestamp.desc()) + .all() + ) + ] - return [ - TagModel(**model_to_dict(tag)) - for tag in Tag.select() - .where(Tag.user_id == user_id) - .where(Tag.name.in_(tag_names)) - ] + return [ + TagModel.model_validate(tag) + for tag in ( + db.query(Tag) + .filter_by(user_id=user_id) + .filter(Tag.name.in_(tag_names)) + .all() + ) + ] def get_tags_by_chat_id_and_user_id( self, chat_id: str, user_id: str ) -> List[TagModel]: - tag_names = [ - ChatIdTagModel(**model_to_dict(chat_id_tag)).tag_name - for chat_id_tag in ChatIdTag.select() - .where((ChatIdTag.user_id == user_id) & (ChatIdTag.chat_id == chat_id)) - .order_by(ChatIdTag.timestamp.desc()) - ] + with get_db() as db: - return [ - TagModel(**model_to_dict(tag)) - for tag in Tag.select() - .where(Tag.user_id == user_id) - .where(Tag.name.in_(tag_names)) - ] + tag_names = [ + chat_id_tag.tag_name + for chat_id_tag in ( + db.query(ChatIdTag) + .filter_by(user_id=user_id, chat_id=chat_id) + .order_by(ChatIdTag.timestamp.desc()) + .all() + ) + ] + + return [ + TagModel.model_validate(tag) + for tag in ( + db.query(Tag) + .filter_by(user_id=user_id) + .filter(Tag.name.in_(tag_names)) + .all() + ) + ] def get_chat_ids_by_tag_name_and_user_id( self, tag_name: str, user_id: str - ) -> Optional[ChatIdTagModel]: - return [ - ChatIdTagModel(**model_to_dict(chat_id_tag)) - for chat_id_tag in ChatIdTag.select() - .where((ChatIdTag.user_id == user_id) & (ChatIdTag.tag_name == tag_name)) - .order_by(ChatIdTag.timestamp.desc()) - ] + ) -> List[ChatIdTagModel]: + with get_db() as db: + + return [ + ChatIdTagModel.model_validate(chat_id_tag) + for chat_id_tag in ( + db.query(ChatIdTag) + .filter_by(user_id=user_id, tag_name=tag_name) + .order_by(ChatIdTag.timestamp.desc()) + .all() + ) + ] def count_chat_ids_by_tag_name_and_user_id( self, tag_name: str, user_id: str ) -> int: - return ( - ChatIdTag.select() - .where((ChatIdTag.tag_name == tag_name) & (ChatIdTag.user_id == user_id)) - .count() - ) + with get_db() as db: + + return ( + db.query(ChatIdTag) + .filter_by(tag_name=tag_name, user_id=user_id) + .count() + ) def delete_tag_by_tag_name_and_user_id(self, tag_name: str, user_id: str) -> bool: try: - query = ChatIdTag.delete().where( - (ChatIdTag.tag_name == tag_name) & (ChatIdTag.user_id == user_id) - ) - res = query.execute() # Remove the rows, return number of rows removed. - log.debug(f"res: {res}") - - tag_count = self.count_chat_ids_by_tag_name_and_user_id(tag_name, user_id) - if tag_count == 0: - # Remove tag item from Tag col as well - query = Tag.delete().where( - (Tag.name == tag_name) & (Tag.user_id == user_id) + with get_db() as db: + res = ( + db.query(ChatIdTag) + .filter_by(tag_name=tag_name, user_id=user_id) + .delete() ) - query.execute() # Remove the rows, return number of rows removed. + log.debug(f"res: {res}") + db.commit() - return True + tag_count = self.count_chat_ids_by_tag_name_and_user_id( + tag_name, user_id + ) + if tag_count == 0: + # Remove tag item from Tag col as well + db.query(Tag).filter_by(name=tag_name, user_id=user_id).delete() + db.commit() + return True except Exception as e: log.error(f"delete_tag: {e}") return False @@ -202,23 +235,25 @@ class TagTable: self, tag_name: str, chat_id: str, user_id: str ) -> bool: try: - query = ChatIdTag.delete().where( - (ChatIdTag.tag_name == tag_name) - & (ChatIdTag.chat_id == chat_id) - & (ChatIdTag.user_id == user_id) - ) - res = query.execute() # Remove the rows, return number of rows removed. - log.debug(f"res: {res}") + with get_db() as db: - tag_count = self.count_chat_ids_by_tag_name_and_user_id(tag_name, user_id) - if tag_count == 0: - # Remove tag item from Tag col as well - query = Tag.delete().where( - (Tag.name == tag_name) & (Tag.user_id == user_id) + res = ( + db.query(ChatIdTag) + .filter_by(tag_name=tag_name, chat_id=chat_id, user_id=user_id) + .delete() ) - query.execute() # Remove the rows, return number of rows removed. + log.debug(f"res: {res}") + db.commit() - return True + tag_count = self.count_chat_ids_by_tag_name_and_user_id( + tag_name, user_id + ) + if tag_count == 0: + # Remove tag item from Tag col as well + db.query(Tag).filter_by(name=tag_name, user_id=user_id).delete() + db.commit() + + return True except Exception as e: log.error(f"delete_tag: {e}") return False @@ -234,4 +269,4 @@ class TagTable: return True -Tags = TagTable(DB) +Tags = TagTable() diff --git a/backend/apps/webui/models/tools.py b/backend/apps/webui/models/tools.py index 694081df9..685ce6fcf 100644 --- a/backend/apps/webui/models/tools.py +++ b/backend/apps/webui/models/tools.py @@ -1,10 +1,10 @@ -from pydantic import BaseModel -from peewee import * -from playhouse.shortcuts import model_to_dict -from typing import List, Union, Optional +from pydantic import BaseModel, ConfigDict +from typing import List, Optional import time import logging -from apps.webui.internal.db import DB, JSONField +from sqlalchemy import String, Column, BigInteger, Text + +from apps.webui.internal.db import Base, JSONField, get_db from apps.webui.models.users import Users import json @@ -21,19 +21,18 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"]) #################### -class Tool(Model): - id = CharField(unique=True) - user_id = CharField() - name = TextField() - content = TextField() - specs = JSONField() - meta = JSONField() - valves = JSONField() - updated_at = BigIntegerField() - created_at = BigIntegerField() +class Tool(Base): + __tablename__ = "tool" - class Meta: - database = DB + id = Column(String, primary_key=True) + user_id = Column(String) + name = Column(Text) + content = Column(Text) + specs = Column(JSONField) + meta = Column(JSONField) + valves = Column(JSONField) + updated_at = Column(BigInteger) + created_at = Column(BigInteger) class ToolMeta(BaseModel): @@ -51,6 +50,8 @@ class ToolModel(BaseModel): updated_at: int # timestamp in epoch created_at: int # timestamp in epoch + model_config = ConfigDict(from_attributes=True) + #################### # Forms @@ -78,61 +79,68 @@ class ToolValves(BaseModel): class ToolsTable: - def __init__(self, db): - self.db = db - self.db.create_tables([Tool]) def insert_new_tool( self, user_id: str, form_data: ToolForm, specs: List[dict] ) -> Optional[ToolModel]: - tool = ToolModel( - **{ - **form_data.model_dump(), - "specs": specs, - "user_id": user_id, - "updated_at": int(time.time()), - "created_at": int(time.time()), - } - ) - try: - result = Tool.create(**tool.model_dump()) - if result: - return tool - else: + with get_db() as db: + + tool = ToolModel( + **{ + **form_data.model_dump(), + "specs": specs, + "user_id": user_id, + "updated_at": int(time.time()), + "created_at": int(time.time()), + } + ) + + try: + result = Tool(**tool.model_dump()) + db.add(result) + db.commit() + db.refresh(result) + if result: + return ToolModel.model_validate(result) + else: + return None + except Exception as e: + print(f"Error creating tool: {e}") return None - except Exception as e: - print(f"Error creating tool: {e}") - return None def get_tool_by_id(self, id: str) -> Optional[ToolModel]: try: - tool = Tool.get(Tool.id == id) - return ToolModel(**model_to_dict(tool)) + with get_db() as db: + + tool = db.get(Tool, id) + return ToolModel.model_validate(tool) except: return None def get_tools(self) -> List[ToolModel]: - return [ToolModel(**model_to_dict(tool)) for tool in Tool.select()] + with get_db() as db: + return [ToolModel.model_validate(tool) for tool in db.query(Tool).all()] def get_tool_valves_by_id(self, id: str) -> Optional[dict]: try: - tool = Tool.get(Tool.id == id) - return tool.valves if tool.valves else {} + with get_db() as db: + + tool = db.get(Tool, id) + return tool.valves if tool.valves else {} except Exception as e: print(f"An error occurred: {e}") return None def update_tool_valves_by_id(self, id: str, valves: dict) -> Optional[ToolValves]: try: - query = Tool.update( - **{"valves": valves}, - updated_at=int(time.time()), - ).where(Tool.id == id) - query.execute() + with get_db() as db: - tool = Tool.get(Tool.id == id) - return ToolValves(**model_to_dict(tool)) + db.query(Tool).filter_by(id=id).update( + {"valves": valves, "updated_at": int(time.time())} + ) + db.commit() + return self.get_tool_by_id(id) except: return None @@ -141,7 +149,7 @@ class ToolsTable: ) -> Optional[dict]: try: user = Users.get_user_by_id(user_id) - user_settings = user.settings.model_dump() + user_settings = user.settings.model_dump() if user.settings else {} # Check if user has "tools" and "valves" settings if "tools" not in user_settings: @@ -159,7 +167,7 @@ class ToolsTable: ) -> Optional[dict]: try: user = Users.get_user_by_id(user_id) - user_settings = user.settings.model_dump() + user_settings = user.settings.model_dump() if user.settings else {} # Check if user has "tools" and "valves" settings if "tools" not in user_settings: @@ -170,8 +178,7 @@ class ToolsTable: user_settings["tools"]["valves"][id] = valves # Update the user settings in the database - query = Users.update_user_by_id(user_id, {"settings": user_settings}) - query.execute() + Users.update_user_by_id(user_id, {"settings": user_settings}) return user_settings["tools"]["valves"][id] except Exception as e: @@ -180,25 +187,27 @@ class ToolsTable: def update_tool_by_id(self, id: str, updated: dict) -> Optional[ToolModel]: try: - query = Tool.update( - **updated, - updated_at=int(time.time()), - ).where(Tool.id == id) - query.execute() + with get_db() as db: + db.query(Tool).filter_by(id=id).update( + {**updated, "updated_at": int(time.time())} + ) + db.commit() - tool = Tool.get(Tool.id == id) - return ToolModel(**model_to_dict(tool)) + tool = db.query(Tool).get(id) + db.refresh(tool) + return ToolModel.model_validate(tool) except: return None def delete_tool_by_id(self, id: str) -> bool: try: - query = Tool.delete().where((Tool.id == id)) - query.execute() # Remove the rows, return number of rows removed. + with get_db() as db: + db.query(Tool).filter_by(id=id).delete() + db.commit() - return True + return True except: return False -Tools = ToolsTable(DB) +Tools = ToolsTable() diff --git a/backend/apps/webui/models/users.py b/backend/apps/webui/models/users.py index e3e1842b8..2f30cda02 100644 --- a/backend/apps/webui/models/users.py +++ b/backend/apps/webui/models/users.py @@ -1,11 +1,12 @@ -from pydantic import BaseModel, ConfigDict -from peewee import * -from playhouse.shortcuts import model_to_dict +from pydantic import BaseModel, ConfigDict, parse_obj_as from typing import List, Union, Optional import time + +from sqlalchemy import String, Column, BigInteger, Text + from utils.misc import get_gravatar_url -from apps.webui.internal.db import DB, JSONField +from apps.webui.internal.db import Base, JSONField, Session, get_db from apps.webui.models.chats import Chats #################### @@ -13,25 +14,24 @@ from apps.webui.models.chats import Chats #################### -class User(Model): - id = CharField(unique=True) - name = CharField() - email = CharField() - role = CharField() - profile_image_url = TextField() +class User(Base): + __tablename__ = "user" - last_active_at = BigIntegerField() - updated_at = BigIntegerField() - created_at = BigIntegerField() + id = Column(String, primary_key=True) + name = Column(String) + email = Column(String) + role = Column(String) + profile_image_url = Column(Text) - api_key = CharField(null=True, unique=True) - settings = JSONField(null=True) - info = JSONField(null=True) + last_active_at = Column(BigInteger) + updated_at = Column(BigInteger) + created_at = Column(BigInteger) - oauth_sub = TextField(null=True, unique=True) + api_key = Column(String, nullable=True, unique=True) + settings = Column(JSONField, nullable=True) + info = Column(JSONField, nullable=True) - class Meta: - database = DB + oauth_sub = Column(Text, unique=True) class UserSettings(BaseModel): @@ -57,6 +57,8 @@ class UserModel(BaseModel): oauth_sub: Optional[str] = None + model_config = ConfigDict(from_attributes=True) + #################### # Forms @@ -76,9 +78,6 @@ class UserUpdateForm(BaseModel): class UsersTable: - def __init__(self, db): - self.db = db - self.db.create_tables([User]) def insert_new_user( self, @@ -89,77 +88,92 @@ class UsersTable: role: str = "pending", oauth_sub: Optional[str] = None, ) -> Optional[UserModel]: - user = UserModel( - **{ - "id": id, - "name": name, - "email": email, - "role": role, - "profile_image_url": profile_image_url, - "last_active_at": int(time.time()), - "created_at": int(time.time()), - "updated_at": int(time.time()), - "oauth_sub": oauth_sub, - } - ) - result = User.create(**user.model_dump()) - if result: - return user - else: - return None + with get_db() as db: + user = UserModel( + **{ + "id": id, + "name": name, + "email": email, + "role": role, + "profile_image_url": profile_image_url, + "last_active_at": int(time.time()), + "created_at": int(time.time()), + "updated_at": int(time.time()), + "oauth_sub": oauth_sub, + } + ) + result = User(**user.model_dump()) + db.add(result) + db.commit() + db.refresh(result) + if result: + return user + else: + return None def get_user_by_id(self, id: str) -> Optional[UserModel]: try: - user = User.get(User.id == id) - return UserModel(**model_to_dict(user)) - except: + with get_db() as db: + user = db.query(User).filter_by(id=id).first() + return UserModel.model_validate(user) + except Exception as e: return None def get_user_by_api_key(self, api_key: str) -> Optional[UserModel]: try: - user = User.get(User.api_key == api_key) - return UserModel(**model_to_dict(user)) + with get_db() as db: + + user = db.query(User).filter_by(api_key=api_key).first() + return UserModel.model_validate(user) except: return None def get_user_by_email(self, email: str) -> Optional[UserModel]: try: - user = User.get(User.email == email) - return UserModel(**model_to_dict(user)) + with get_db() as db: + + user = db.query(User).filter_by(email=email).first() + return UserModel.model_validate(user) except: return None def get_user_by_oauth_sub(self, sub: str) -> Optional[UserModel]: try: - user = User.get(User.oauth_sub == sub) - return UserModel(**model_to_dict(user)) + with get_db() as db: + + user = db.query(User).filter_by(oauth_sub=sub).first() + return UserModel.model_validate(user) except: return None def get_users(self, skip: int = 0, limit: int = 50) -> List[UserModel]: - return [ - UserModel(**model_to_dict(user)) - for user in User.select() - # .limit(limit).offset(skip) - ] + with get_db() as db: + users = ( + db.query(User) + # .offset(skip).limit(limit) + .all() + ) + return [UserModel.model_validate(user) for user in users] def get_num_users(self) -> Optional[int]: - return User.select().count() + with get_db() as db: + return db.query(User).count() def get_first_user(self) -> UserModel: try: - user = User.select().order_by(User.created_at).first() - return UserModel(**model_to_dict(user)) + with get_db() as db: + user = db.query(User).order_by(User.created_at).first() + return UserModel.model_validate(user) except: return None def update_user_role_by_id(self, id: str, role: str) -> Optional[UserModel]: try: - query = User.update(role=role).where(User.id == id) - query.execute() - - user = User.get(User.id == id) - return UserModel(**model_to_dict(user)) + with get_db() as db: + db.query(User).filter_by(id=id).update({"role": role}) + db.commit() + user = db.query(User).filter_by(id=id).first() + return UserModel.model_validate(user) except: return None @@ -167,23 +181,28 @@ class UsersTable: self, id: str, profile_image_url: str ) -> Optional[UserModel]: try: - query = User.update(profile_image_url=profile_image_url).where( - User.id == id - ) - query.execute() + with get_db() as db: + db.query(User).filter_by(id=id).update( + {"profile_image_url": profile_image_url} + ) + db.commit() - user = User.get(User.id == id) - return UserModel(**model_to_dict(user)) + user = db.query(User).filter_by(id=id).first() + return UserModel.model_validate(user) except: return None def update_user_last_active_by_id(self, id: str) -> Optional[UserModel]: try: - query = User.update(last_active_at=int(time.time())).where(User.id == id) - query.execute() + with get_db() as db: - user = User.get(User.id == id) - return UserModel(**model_to_dict(user)) + db.query(User).filter_by(id=id).update( + {"last_active_at": int(time.time())} + ) + db.commit() + + user = db.query(User).filter_by(id=id).first() + return UserModel.model_validate(user) except: return None @@ -191,22 +210,25 @@ class UsersTable: self, id: str, oauth_sub: str ) -> Optional[UserModel]: try: - query = User.update(oauth_sub=oauth_sub).where(User.id == id) - query.execute() + with get_db() as db: + db.query(User).filter_by(id=id).update({"oauth_sub": oauth_sub}) + db.commit() - user = User.get(User.id == id) - return UserModel(**model_to_dict(user)) + user = db.query(User).filter_by(id=id).first() + return UserModel.model_validate(user) except: return None def update_user_by_id(self, id: str, updated: dict) -> Optional[UserModel]: try: - query = User.update(**updated).where(User.id == id) - query.execute() + with get_db() as db: + db.query(User).filter_by(id=id).update(updated) + db.commit() - user = User.get(User.id == id) - return UserModel(**model_to_dict(user)) - except: + user = db.query(User).filter_by(id=id).first() + return UserModel.model_validate(user) + # return UserModel(**user.dict()) + except Exception as e: return None def delete_user_by_id(self, id: str) -> bool: @@ -215,9 +237,10 @@ class UsersTable: result = Chats.delete_chats_by_user_id(id) if result: - # Delete User - query = User.delete().where(User.id == id) - query.execute() # Remove the rows, return number of rows removed. + with get_db() as db: + # Delete User + db.query(User).filter_by(id=id).delete() + db.commit() return True else: @@ -227,19 +250,20 @@ class UsersTable: def update_user_api_key_by_id(self, id: str, api_key: str) -> str: try: - query = User.update(api_key=api_key).where(User.id == id) - result = query.execute() - - return True if result == 1 else False + with get_db() as db: + result = db.query(User).filter_by(id=id).update({"api_key": api_key}) + db.commit() + return True if result == 1 else False except: return False def get_user_api_key_by_id(self, id: str) -> Optional[str]: try: - user = User.get(User.id == id) - return user.api_key - except: + with get_db() as db: + user = db.query(User).filter_by(id=id).first() + return user.api_key + except Exception as e: return None -Users = UsersTable(DB) +Users = UsersTable() diff --git a/backend/apps/webui/routers/chats.py b/backend/apps/webui/routers/chats.py index c4d6575c2..d3ccb9cce 100644 --- a/backend/apps/webui/routers/chats.py +++ b/backend/apps/webui/routers/chats.py @@ -76,7 +76,10 @@ async def delete_all_user_chats(request: Request, user=Depends(get_verified_user @router.get("/list/user/{user_id}", response_model=List[ChatTitleIdResponse]) async def get_user_chat_list_by_user_id( - user_id: str, user=Depends(get_admin_user), skip: int = 0, limit: int = 50 + user_id: str, + user=Depends(get_admin_user), + skip: int = 0, + limit: int = 50, ): return Chats.get_chat_list_by_user_id( user_id, include_archived=True, skip=skip, limit=limit @@ -119,7 +122,7 @@ async def get_user_chats(user=Depends(get_verified_user)): @router.get("/all/archived", response_model=List[ChatResponse]) -async def get_user_chats(user=Depends(get_verified_user)): +async def get_user_archived_chats(user=Depends(get_verified_user)): return [ ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)}) for chat in Chats.get_archived_chats_by_user_id(user.id) @@ -207,7 +210,6 @@ async def get_user_chat_list_by_tag_name( form_data: TagNameForm, user=Depends(get_verified_user) ): - print(form_data) chat_ids = [ chat_id_tag.chat_id for chat_id_tag in Tags.get_chat_ids_by_tag_name_and_user_id( diff --git a/backend/apps/webui/routers/documents.py b/backend/apps/webui/routers/documents.py index dc53b5246..2299b2fee 100644 --- a/backend/apps/webui/routers/documents.py +++ b/backend/apps/webui/routers/documents.py @@ -130,7 +130,9 @@ async def tag_doc_by_name(form_data: TagDocumentForm, user=Depends(get_verified_ @router.post("/doc/update", response_model=Optional[DocumentResponse]) async def update_doc_by_name( - name: str, form_data: DocumentUpdateForm, user=Depends(get_admin_user) + name: str, + form_data: DocumentUpdateForm, + user=Depends(get_admin_user), ): doc = Documents.update_doc_by_name(name, form_data) if doc: diff --git a/backend/apps/webui/routers/files.py b/backend/apps/webui/routers/files.py index 3b6d44aa5..fffe0743c 100644 --- a/backend/apps/webui/routers/files.py +++ b/backend/apps/webui/routers/files.py @@ -50,10 +50,7 @@ router = APIRouter() @router.post("/") -def upload_file( - file: UploadFile = File(...), - user=Depends(get_verified_user), -): +def upload_file(file: UploadFile = File(...), user=Depends(get_verified_user)): log.info(f"file.content_type: {file.content_type}") try: unsanitized_filename = file.filename diff --git a/backend/apps/webui/routers/functions.py b/backend/apps/webui/routers/functions.py index f01133a35..eb5216b20 100644 --- a/backend/apps/webui/routers/functions.py +++ b/backend/apps/webui/routers/functions.py @@ -233,7 +233,10 @@ async def delete_function_by_id( # delete the function file function_path = os.path.join(FUNCTIONS_DIR, f"{id}.py") - os.remove(function_path) + try: + os.remove(function_path) + except: + pass return result diff --git a/backend/apps/webui/routers/memories.py b/backend/apps/webui/routers/memories.py index e9ae96173..2c473ebe8 100644 --- a/backend/apps/webui/routers/memories.py +++ b/backend/apps/webui/routers/memories.py @@ -50,7 +50,9 @@ class MemoryUpdateModel(BaseModel): @router.post("/add", response_model=Optional[MemoryModel]) async def add_memory( - request: Request, form_data: AddMemoryForm, user=Depends(get_verified_user) + request: Request, + form_data: AddMemoryForm, + user=Depends(get_verified_user), ): memory = Memories.insert_new_memory(user.id, form_data.content) memory_embedding = request.app.state.EMBEDDING_FUNCTION(memory.content) diff --git a/backend/apps/webui/routers/models.py b/backend/apps/webui/routers/models.py index acc1c6b47..eeae9e1c4 100644 --- a/backend/apps/webui/routers/models.py +++ b/backend/apps/webui/routers/models.py @@ -5,6 +5,7 @@ from typing import List, Union, Optional from fastapi import APIRouter from pydantic import BaseModel import json + from apps.webui.models.models import Models, ModelModel, ModelForm, ModelResponse from utils.utils import get_verified_user, get_admin_user @@ -29,7 +30,9 @@ async def get_models(user=Depends(get_verified_user)): @router.post("/add", response_model=Optional[ModelModel]) async def add_new_model( - request: Request, form_data: ModelForm, user=Depends(get_admin_user) + request: Request, + form_data: ModelForm, + user=Depends(get_admin_user), ): if form_data.id in request.app.state.MODELS: raise HTTPException( @@ -73,7 +76,10 @@ async def get_model_by_id(id: str, user=Depends(get_verified_user)): @router.post("/update", response_model=Optional[ModelModel]) async def update_model_by_id( - request: Request, id: str, form_data: ModelForm, user=Depends(get_admin_user) + request: Request, + id: str, + form_data: ModelForm, + user=Depends(get_admin_user), ): model = Models.get_model_by_id(id) if model: diff --git a/backend/apps/webui/routers/prompts.py b/backend/apps/webui/routers/prompts.py index e609a0a1b..c674590e9 100644 --- a/backend/apps/webui/routers/prompts.py +++ b/backend/apps/webui/routers/prompts.py @@ -71,7 +71,9 @@ async def get_prompt_by_command(command: str, user=Depends(get_verified_user)): @router.post("/command/{command}/update", response_model=Optional[PromptModel]) async def update_prompt_by_command( - command: str, form_data: PromptForm, user=Depends(get_admin_user) + command: str, + form_data: PromptForm, + user=Depends(get_admin_user), ): prompt = Prompts.update_prompt_by_command(f"/{command}", form_data) if prompt: diff --git a/backend/apps/webui/routers/tools.py b/backend/apps/webui/routers/tools.py index d20584c22..ea9db8180 100644 --- a/backend/apps/webui/routers/tools.py +++ b/backend/apps/webui/routers/tools.py @@ -6,7 +6,6 @@ from fastapi import APIRouter from pydantic import BaseModel import json - from apps.webui.models.users import Users from apps.webui.models.tools import Tools, ToolForm, ToolModel, ToolResponse from apps.webui.utils import load_toolkit_module_by_id @@ -57,7 +56,9 @@ async def get_toolkits(user=Depends(get_admin_user)): @router.post("/create", response_model=Optional[ToolResponse]) async def create_new_toolkit( - request: Request, form_data: ToolForm, user=Depends(get_admin_user) + request: Request, + form_data: ToolForm, + user=Depends(get_admin_user), ): if not form_data.id.isidentifier(): raise HTTPException( @@ -131,7 +132,10 @@ async def get_toolkit_by_id(id: str, user=Depends(get_admin_user)): @router.post("/id/{id}/update", response_model=Optional[ToolModel]) async def update_toolkit_by_id( - request: Request, id: str, form_data: ToolForm, user=Depends(get_admin_user) + request: Request, + id: str, + form_data: ToolForm, + user=Depends(get_admin_user), ): toolkit_path = os.path.join(TOOLS_DIR, f"{id}.py") diff --git a/backend/apps/webui/routers/users.py b/backend/apps/webui/routers/users.py index 270d72a23..9627f0b06 100644 --- a/backend/apps/webui/routers/users.py +++ b/backend/apps/webui/routers/users.py @@ -138,7 +138,7 @@ async def get_user_info_by_session_user(user=Depends(get_verified_user)): @router.post("/user/info/update", response_model=Optional[dict]) -async def update_user_settings_by_session_user( +async def update_user_info_by_session_user( form_data: dict, user=Depends(get_verified_user) ): user = Users.get_user_by_id(user.id) @@ -205,7 +205,9 @@ async def get_user_by_id(user_id: str, user=Depends(get_verified_user)): @router.post("/{user_id}/update", response_model=Optional[UserModel]) async def update_user_by_id( - user_id: str, form_data: UserUpdateForm, session_user=Depends(get_admin_user) + user_id: str, + form_data: UserUpdateForm, + session_user=Depends(get_admin_user), ): user = Users.get_user_by_id(user_id) diff --git a/backend/apps/webui/routers/utils.py b/backend/apps/webui/routers/utils.py index 8f6d663b4..780ed6b43 100644 --- a/backend/apps/webui/routers/utils.py +++ b/backend/apps/webui/routers/utils.py @@ -1,6 +1,5 @@ from fastapi import APIRouter, UploadFile, File, Response from fastapi import Depends, HTTPException, status -from peewee import SqliteDatabase from starlette.responses import StreamingResponse, FileResponse from pydantic import BaseModel @@ -10,7 +9,6 @@ import markdown import black -from apps.webui.internal.db import DB from utils.utils import get_admin_user from utils.misc import calculate_sha256, get_gravatar_url @@ -114,13 +112,15 @@ async def download_db(user=Depends(get_admin_user)): status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.ACCESS_PROHIBITED, ) - if not isinstance(DB, SqliteDatabase): + from apps.webui.internal.db import engine + + if engine.name != "sqlite": raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.DB_NOT_SQLITE, ) return FileResponse( - DB.database, + engine.url.database, media_type="application/octet-stream", filename="webui.db", ) diff --git a/backend/config.py b/backend/config.py index 4b4b5deac..fe68eee34 100644 --- a/backend/config.py +++ b/backend/config.py @@ -5,9 +5,8 @@ import importlib.metadata import pkgutil import chromadb from chromadb import Settings -from base64 import b64encode from bs4 import BeautifulSoup -from typing import TypeVar, Generic, Union +from typing import TypeVar, Generic from pydantic import BaseModel from typing import Optional @@ -19,7 +18,6 @@ import markdown import requests import shutil -from secrets import token_bytes from constants import ERROR_MESSAGES #################################### @@ -395,6 +393,18 @@ OAUTH_PROVIDER_NAME = PersistentConfig( os.environ.get("OAUTH_PROVIDER_NAME", "SSO"), ) +OAUTH_USERNAME_CLAIM = PersistentConfig( + "OAUTH_USERNAME_CLAIM", + "oauth.oidc.username_claim", + os.environ.get("OAUTH_USERNAME_CLAIM", "name"), +) + +OAUTH_PICTURE_CLAIM = PersistentConfig( + "OAUTH_USERNAME_CLAIM", + "oauth.oidc.avatar_claim", + os.environ.get("OAUTH_PICTURE_CLAIM", "picture"), +) + def load_oauth_providers(): OAUTH_PROVIDERS.clear() @@ -440,16 +450,27 @@ load_oauth_providers() STATIC_DIR = Path(os.getenv("STATIC_DIR", BACKEND_DIR / "static")).resolve() -frontend_favicon = FRONTEND_BUILD_DIR / "favicon.png" +frontend_favicon = FRONTEND_BUILD_DIR / "static" / "favicon.png" + if frontend_favicon.exists(): try: shutil.copyfile(frontend_favicon, STATIC_DIR / "favicon.png") except Exception as e: logging.error(f"An error occurred: {e}") - else: logging.warning(f"Frontend favicon not found at {frontend_favicon}") +frontend_splash = FRONTEND_BUILD_DIR / "static" / "splash.png" + +if frontend_splash.exists(): + try: + shutil.copyfile(frontend_splash, STATIC_DIR / "splash.png") + except Exception as e: + logging.error(f"An error occurred: {e}") +else: + logging.warning(f"Frontend splash not found at {frontend_splash}") + + #################################### # CUSTOM_NAME #################################### @@ -474,6 +495,19 @@ if CUSTOM_NAME: r.raw.decode_content = True shutil.copyfileobj(r.raw, f) + if "splash" in data: + url = ( + f"https://api.openwebui.com{data['splash']}" + if data["splash"][0] == "/" + else data["splash"] + ) + + r = requests.get(url, stream=True) + if r.status_code == 200: + with open(f"{STATIC_DIR}/splash.png", "wb") as f: + r.raw.decode_content = True + shutil.copyfileobj(r.raw, f) + WEBUI_NAME = data["name"] except Exception as e: log.exception(e) @@ -769,11 +803,14 @@ class BannerModel(BaseModel): timestamp: int -WEBUI_BANNERS = PersistentConfig( - "WEBUI_BANNERS", - "ui.banners", - [BannerModel(**banner) for banner in json.loads("[]")], -) +try: + banners = json.loads(os.environ.get("WEBUI_BANNERS", "[]")) + banners = [BannerModel(**banner) for banner in banners] +except Exception as e: + print(f"Error loading WEBUI_BANNERS: {e}") + banners = [] + +WEBUI_BANNERS = PersistentConfig("WEBUI_BANNERS", "ui.banners", banners) SHOW_ADMIN_DETAILS = PersistentConfig( @@ -885,6 +922,22 @@ WEBUI_SESSION_COOKIE_SECURE = os.environ.get( if WEBUI_AUTH and WEBUI_SECRET_KEY == "": raise ValueError(ERROR_MESSAGES.ENV_VAR_NOT_FOUND) +#################################### +# RAG document content extraction +#################################### + +CONTENT_EXTRACTION_ENGINE = PersistentConfig( + "CONTENT_EXTRACTION_ENGINE", + "rag.CONTENT_EXTRACTION_ENGINE", + os.environ.get("CONTENT_EXTRACTION_ENGINE", "").lower(), +) + +TIKA_SERVER_URL = PersistentConfig( + "TIKA_SERVER_URL", + "rag.tika_server_url", + os.getenv("TIKA_SERVER_URL", "http://tika:9998"), # Default for sidecar deployment +) + #################################### # RAG #################################### @@ -1302,3 +1355,7 @@ AUDIO_TTS_VOICE = PersistentConfig( #################################### DATABASE_URL = os.environ.get("DATABASE_URL", f"sqlite:///{DATA_DIR}/webui.db") + +# Replace the postgres:// with postgresql:// +if "postgres://" in DATABASE_URL: + DATABASE_URL = DATABASE_URL.replace("postgres://", "postgresql://") diff --git a/backend/constants.py b/backend/constants.py index f1eed43d3..7c366c222 100644 --- a/backend/constants.py +++ b/backend/constants.py @@ -89,3 +89,14 @@ class ERROR_MESSAGES(str, Enum): OLLAMA_API_DISABLED = ( "The Ollama API is disabled. Please enable it to use this feature." ) + + +class TASKS(str, Enum): + def __str__(self) -> str: + return super().__str__() + + DEFAULT = lambda task="": f"{task if task else 'default'}" + TITLE_GENERATION = "Title Generation" + EMOJI_GENERATION = "Emoji Generation" + QUERY_GENERATION = "Query Generation" + FUNCTION_CALLING = "Function Calling" diff --git a/backend/main.py b/backend/main.py index df3f884b3..89252e164 100644 --- a/backend/main.py +++ b/backend/main.py @@ -4,9 +4,7 @@ from contextlib import asynccontextmanager from authlib.integrations.starlette_client import OAuth from authlib.oidc.core import UserInfo -from bs4 import BeautifulSoup import json -import markdown import time import os import sys @@ -18,25 +16,22 @@ import shutil import os import uuid import inspect -import asyncio -from fastapi.concurrency import run_in_threadpool from fastapi import FastAPI, Request, Depends, status, UploadFile, File, Form from fastapi.staticfiles import StaticFiles from fastapi.responses import JSONResponse from fastapi import HTTPException -from fastapi.middleware.wsgi import WSGIMiddleware from fastapi.middleware.cors import CORSMiddleware +from sqlalchemy import text from starlette.exceptions import HTTPException as StarletteHTTPException from starlette.middleware.base import BaseHTTPMiddleware from starlette.middleware.sessions import SessionMiddleware from starlette.responses import StreamingResponse, Response, RedirectResponse -from apps.socket.main import app as socket_app +from apps.socket.main import sio, app as socket_app from apps.ollama.main import ( app as ollama_app, - OpenAIChatCompletionForm, get_all_models as get_ollama_models, generate_openai_chat_completion as generate_ollama_chat_completion, ) @@ -54,13 +49,14 @@ from apps.webui.main import ( get_pipe_models, generate_function_chat_completion, ) +from apps.webui.internal.db import Session from pydantic import BaseModel -from typing import List, Optional, Iterator, Generator, Union +from typing import List, Optional from apps.webui.models.auths import Auths -from apps.webui.models.models import Models, ModelModel +from apps.webui.models.models import Models from apps.webui.models.tools import Tools from apps.webui.models.functions import Functions from apps.webui.models.users import Users @@ -83,14 +79,12 @@ from utils.task import ( from utils.misc import ( get_last_user_message, add_or_update_system_message, - stream_message_template, parse_duration, ) from apps.rag.utils import get_rag_context, rag_template from config import ( - CONFIG_DATA, WEBUI_NAME, WEBUI_URL, WEBUI_AUTH, @@ -98,7 +92,6 @@ from config import ( VERSION, CHANGELOG, FRONTEND_BUILD_DIR, - UPLOAD_DIR, CACHE_DIR, STATIC_DIR, DEFAULT_LOCALE, @@ -126,7 +119,8 @@ from config import ( WEBUI_SESSION_COOKIE_SECURE, AppConfig, ) -from constants import ERROR_MESSAGES, WEBHOOK_MESSAGES + +from constants import ERROR_MESSAGES, WEBHOOK_MESSAGES, TASKS from utils.webhook import post_webhook if SAFE_MODE: @@ -167,8 +161,20 @@ https://github.com/open-webui/open-webui ) +def run_migrations(): + try: + from alembic.config import Config + from alembic import command + + alembic_cfg = Config("alembic.ini") + command.upgrade(alembic_cfg, "head") + except Exception as e: + print(f"Error: {e}") + + @asynccontextmanager async def lifespan(app: FastAPI): + run_migrations() yield @@ -212,8 +218,79 @@ origins = ["*"] ################################## +async def get_body_and_model_and_user(request): + # Read the original request body + body = await request.body() + body_str = body.decode("utf-8") + body = json.loads(body_str) if body_str else {} + + model_id = body["model"] + if model_id not in app.state.MODELS: + raise Exception("Model not found") + model = app.state.MODELS[model_id] + + user = get_current_user( + request, + get_http_authorization_cred(request.headers.get("Authorization")), + ) + + return body, model, user + + +def get_task_model_id(default_model_id): + # Set the task model + task_model_id = default_model_id + # Check if the user has a custom task model and use that model + if app.state.MODELS[task_model_id]["owned_by"] == "ollama": + if ( + app.state.config.TASK_MODEL + and app.state.config.TASK_MODEL in app.state.MODELS + ): + task_model_id = app.state.config.TASK_MODEL + else: + if ( + app.state.config.TASK_MODEL_EXTERNAL + and app.state.config.TASK_MODEL_EXTERNAL in app.state.MODELS + ): + task_model_id = app.state.config.TASK_MODEL_EXTERNAL + + return task_model_id + + +def get_filter_function_ids(model): + def get_priority(function_id): + function = Functions.get_function_by_id(function_id) + if function is not None and hasattr(function, "valves"): + return (function.valves if function.valves else {}).get("priority", 0) + return 0 + + filter_ids = [function.id for function in Functions.get_global_filter_functions()] + if "info" in model and "meta" in model["info"]: + filter_ids.extend(model["info"]["meta"].get("filterIds", [])) + filter_ids = list(set(filter_ids)) + + enabled_filter_ids = [ + function.id + for function in Functions.get_functions_by_type("filter", active_only=True) + ] + + filter_ids = [ + filter_id for filter_id in filter_ids if filter_id in enabled_filter_ids + ] + + filter_ids.sort(key=get_priority) + return filter_ids + + async def get_function_call_response( - messages, files, tool_id, template, task_model_id, user + messages, + files, + tool_id, + template, + task_model_id, + user, + __event_emitter__=None, + __event_call__=None, ): tool = Tools.get_tool_by_id(tool_id) tools_specs = json.dumps(tool.specs, indent=2) @@ -240,6 +317,7 @@ async def get_function_call_response( {"role": "user", "content": f"Query: {prompt}"}, ], "stream": False, + "task": TASKS.FUNCTION_CALLING, } try: @@ -252,7 +330,6 @@ async def get_function_call_response( response = None try: response = await generate_chat_completions(form_data=payload, user=user) - content = None if hasattr(response, "body_iterator"): @@ -266,334 +343,367 @@ async def get_function_call_response( else: content = response["choices"][0]["message"]["content"] + if content is None: + return None, None, False + # Parse the function response - if content is not None: - print(f"content: {content}") - result = json.loads(content) - print(result) + print(f"content: {content}") + result = json.loads(content) + print(result) - citation = None - # Call the function - if "name" in result: - if tool_id in webui_app.state.TOOLS: - toolkit_module = webui_app.state.TOOLS[tool_id] - else: - toolkit_module, frontmatter = load_toolkit_module_by_id(tool_id) - webui_app.state.TOOLS[tool_id] = toolkit_module + citation = None - file_handler = False - # check if toolkit_module has file_handler self variable - if hasattr(toolkit_module, "file_handler"): - file_handler = True - print("file_handler: ", file_handler) + if "name" not in result: + return None, None, False - if hasattr(toolkit_module, "valves") and hasattr( - toolkit_module, "Valves" - ): - valves = Tools.get_tool_valves_by_id(tool_id) - toolkit_module.valves = toolkit_module.Valves( - **(valves if valves else {}) - ) + # Call the function + if tool_id in webui_app.state.TOOLS: + toolkit_module = webui_app.state.TOOLS[tool_id] + else: + toolkit_module, _ = load_toolkit_module_by_id(tool_id) + webui_app.state.TOOLS[tool_id] = toolkit_module + + file_handler = False + # check if toolkit_module has file_handler self variable + if hasattr(toolkit_module, "file_handler"): + file_handler = True + print("file_handler: ", file_handler) + + if hasattr(toolkit_module, "valves") and hasattr(toolkit_module, "Valves"): + valves = Tools.get_tool_valves_by_id(tool_id) + toolkit_module.valves = toolkit_module.Valves(**(valves if valves else {})) + + function = getattr(toolkit_module, result["name"]) + function_result = None + try: + # Get the signature of the function + sig = inspect.signature(function) + params = result["parameters"] + + # Extra parameters to be passed to the function + extra_params = { + "__model__": model, + "__id__": tool_id, + "__messages__": messages, + "__files__": files, + "__event_emitter__": __event_emitter__, + "__event_call__": __event_call__, + } + + # Add extra params in contained in function signature + for key, value in extra_params.items(): + if key in sig.parameters: + params[key] = value + + if "__user__" in sig.parameters: + # Call the function with the '__user__' parameter included + __user__ = { + "id": user.id, + "email": user.email, + "name": user.name, + "role": user.role, + } - function = getattr(toolkit_module, result["name"]) - function_result = None try: - # Get the signature of the function - sig = inspect.signature(function) - params = result["parameters"] - - if "__user__" in sig.parameters: - # Call the function with the '__user__' parameter included - __user__ = { - "id": user.id, - "email": user.email, - "name": user.name, - "role": user.role, - } - - try: - if hasattr(toolkit_module, "UserValves"): - __user__["valves"] = toolkit_module.UserValves( - **Tools.get_user_valves_by_id_and_user_id( - tool_id, user.id - ) - ) - except Exception as e: - print(e) - - params = {**params, "__user__": __user__} - if "__messages__" in sig.parameters: - # Call the function with the '__messages__' parameter included - params = { - **params, - "__messages__": messages, - } - - if "__files__" in sig.parameters: - # Call the function with the '__files__' parameter included - params = { - **params, - "__files__": files, - } - - if "__model__" in sig.parameters: - # Call the function with the '__model__' parameter included - params = { - **params, - "__model__": model, - } - - if "__id__" in sig.parameters: - # Call the function with the '__id__' parameter included - params = { - **params, - "__id__": tool_id, - } - - if inspect.iscoroutinefunction(function): - function_result = await function(**params) - else: - function_result = function(**params) - - if hasattr(toolkit_module, "citation") and toolkit_module.citation: - citation = { - "source": {"name": f"TOOL:{tool.name}/{result['name']}"}, - "document": [function_result], - "metadata": [{"source": result["name"]}], - } + if hasattr(toolkit_module, "UserValves"): + __user__["valves"] = toolkit_module.UserValves( + **Tools.get_user_valves_by_id_and_user_id(tool_id, user.id) + ) except Exception as e: print(e) - # Add the function result to the system prompt - if function_result is not None: - return function_result, citation, file_handler + params = {**params, "__user__": __user__} + + if inspect.iscoroutinefunction(function): + function_result = await function(**params) + else: + function_result = function(**params) + + if hasattr(toolkit_module, "citation") and toolkit_module.citation: + citation = { + "source": {"name": f"TOOL:{tool.name}/{result['name']}"}, + "document": [function_result], + "metadata": [{"source": result["name"]}], + } + except Exception as e: + print(e) + + # Add the function result to the system prompt + if function_result is not None: + return function_result, citation, file_handler except Exception as e: print(f"Error: {e}") return None, None, False +async def chat_completion_functions_handler( + body, model, user, __event_emitter__, __event_call__ +): + skip_files = None + + filter_ids = get_filter_function_ids(model) + for filter_id in filter_ids: + filter = Functions.get_function_by_id(filter_id) + if not filter: + continue + + if filter_id in webui_app.state.FUNCTIONS: + function_module = webui_app.state.FUNCTIONS[filter_id] + else: + function_module, _, _ = load_function_module_by_id(filter_id) + webui_app.state.FUNCTIONS[filter_id] = function_module + + # Check if the function has a file_handler variable + if hasattr(function_module, "file_handler"): + skip_files = function_module.file_handler + + if hasattr(function_module, "valves") and hasattr(function_module, "Valves"): + valves = Functions.get_function_valves_by_id(filter_id) + function_module.valves = function_module.Valves( + **(valves if valves else {}) + ) + + if not hasattr(function_module, "inlet"): + continue + + try: + inlet = function_module.inlet + + # Get the signature of the function + sig = inspect.signature(inlet) + params = {"body": body} + + # Extra parameters to be passed to the function + extra_params = { + "__model__": model, + "__id__": filter_id, + "__event_emitter__": __event_emitter__, + "__event_call__": __event_call__, + } + + # Add extra params in contained in function signature + for key, value in extra_params.items(): + if key in sig.parameters: + params[key] = value + + if "__user__" in sig.parameters: + __user__ = { + "id": user.id, + "email": user.email, + "name": user.name, + "role": user.role, + } + + try: + if hasattr(function_module, "UserValves"): + __user__["valves"] = function_module.UserValves( + **Functions.get_user_valves_by_id_and_user_id( + filter_id, user.id + ) + ) + except Exception as e: + print(e) + + params = {**params, "__user__": __user__} + + if inspect.iscoroutinefunction(inlet): + body = await inlet(**params) + else: + body = inlet(**params) + + except Exception as e: + print(f"Error: {e}") + raise e + + if skip_files: + if "files" in body: + del body["files"] + + return body, {} + + +async def chat_completion_tools_handler(body, user, __event_emitter__, __event_call__): + skip_files = None + + contexts = [] + citations = None + + task_model_id = get_task_model_id(body["model"]) + + # If tool_ids field is present, call the functions + if "tool_ids" in body: + print(body["tool_ids"]) + for tool_id in body["tool_ids"]: + print(tool_id) + try: + response, citation, file_handler = await get_function_call_response( + messages=body["messages"], + files=body.get("files", []), + tool_id=tool_id, + template=app.state.config.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE, + task_model_id=task_model_id, + user=user, + __event_emitter__=__event_emitter__, + __event_call__=__event_call__, + ) + + print(file_handler) + if isinstance(response, str): + contexts.append(response) + + if citation: + if citations is None: + citations = [citation] + else: + citations.append(citation) + + if file_handler: + skip_files = True + + except Exception as e: + print(f"Error: {e}") + del body["tool_ids"] + print(f"tool_contexts: {contexts}") + + if skip_files: + if "files" in body: + del body["files"] + + return body, { + **({"contexts": contexts} if contexts is not None else {}), + **({"citations": citations} if citations is not None else {}), + } + + +async def chat_completion_files_handler(body): + contexts = [] + citations = None + + if "files" in body: + files = body["files"] + del body["files"] + + contexts, citations = get_rag_context( + files=files, + messages=body["messages"], + embedding_function=rag_app.state.EMBEDDING_FUNCTION, + k=rag_app.state.config.TOP_K, + reranking_function=rag_app.state.sentence_transformer_rf, + r=rag_app.state.config.RELEVANCE_THRESHOLD, + hybrid_search=rag_app.state.config.ENABLE_RAG_HYBRID_SEARCH, + ) + + log.debug(f"rag_contexts: {contexts}, citations: {citations}") + + return body, { + **({"contexts": contexts} if contexts is not None else {}), + **({"citations": citations} if citations is not None else {}), + } + + class ChatCompletionMiddleware(BaseHTTPMiddleware): async def dispatch(self, request: Request, call_next): - data_items = [] - - show_citations = False - citations = [] - if request.method == "POST" and any( endpoint in request.url.path for endpoint in ["/ollama/api/chat", "/chat/completions"] ): log.debug(f"request.url.path: {request.url.path}") - # Read the original request body - body = await request.body() - body_str = body.decode("utf-8") - data = json.loads(body_str) if body_str else {} - - user = get_current_user( - request, - get_http_authorization_cred(request.headers.get("Authorization")), - ) - # Flag to skip RAG completions if file_handler is present in tools/functions - skip_files = False - if data.get("citations"): - show_citations = True - del data["citations"] - - model_id = data["model"] - if model_id not in app.state.MODELS: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail="Model not found", + try: + body, model, user = await get_body_and_model_and_user(request) + except Exception as e: + return JSONResponse( + status_code=status.HTTP_400_BAD_REQUEST, + content={"detail": str(e)}, ) - model = app.state.MODELS[model_id] - def get_priority(function_id): - function = Functions.get_function_by_id(function_id) - if function is not None and hasattr(function, "valves"): - return (function.valves if function.valves else {}).get( - "priority", 0 - ) - return 0 + # Extract session_id, chat_id and message_id from the request body + session_id = None + if "session_id" in body: + session_id = body["session_id"] + del body["session_id"] + chat_id = None + if "chat_id" in body: + chat_id = body["chat_id"] + del body["chat_id"] + message_id = None + if "id" in body: + message_id = body["id"] + del body["id"] - filter_ids = [ - function.id for function in Functions.get_global_filter_functions() - ] - if "info" in model and "meta" in model["info"]: - filter_ids.extend(model["info"]["meta"].get("filterIds", [])) - filter_ids = list(set(filter_ids)) - - enabled_filter_ids = [ - function.id - for function in Functions.get_functions_by_type( - "filter", active_only=True + async def __event_emitter__(data): + await sio.emit( + "chat-events", + { + "chat_id": chat_id, + "message_id": message_id, + "data": data, + }, + to=session_id, ) - ] - filter_ids = [ - filter_id for filter_id in filter_ids if filter_id in enabled_filter_ids - ] - filter_ids.sort(key=get_priority) - for filter_id in filter_ids: - filter = Functions.get_function_by_id(filter_id) - if filter: - if filter_id in webui_app.state.FUNCTIONS: - function_module = webui_app.state.FUNCTIONS[filter_id] - else: - function_module, function_type, frontmatter = ( - load_function_module_by_id(filter_id) - ) - webui_app.state.FUNCTIONS[filter_id] = function_module + async def __event_call__(data): + response = await sio.call( + "chat-events", + {"chat_id": chat_id, "message_id": message_id, "data": data}, + to=session_id, + ) + return response - # Check if the function has a file_handler variable - if hasattr(function_module, "file_handler"): - skip_files = function_module.file_handler + # Initialize data_items to store additional data to be sent to the client + data_items = [] - if hasattr(function_module, "valves") and hasattr( - function_module, "Valves" - ): - valves = Functions.get_function_valves_by_id(filter_id) - function_module.valves = function_module.Valves( - **(valves if valves else {}) - ) + # Initialize context, and citations + contexts = [] + citations = [] - try: - if hasattr(function_module, "inlet"): - inlet = function_module.inlet + try: + body, flags = await chat_completion_functions_handler( + body, model, user, __event_emitter__, __event_call__ + ) + except Exception as e: + return JSONResponse( + status_code=status.HTTP_400_BAD_REQUEST, + content={"detail": str(e)}, + ) - # Get the signature of the function - sig = inspect.signature(inlet) - params = {"body": data} + try: + body, flags = await chat_completion_tools_handler( + body, user, __event_emitter__, __event_call__ + ) - if "__user__" in sig.parameters: - __user__ = { - "id": user.id, - "email": user.email, - "name": user.name, - "role": user.role, - } + contexts.extend(flags.get("contexts", [])) + citations.extend(flags.get("citations", [])) + except Exception as e: + print(e) + pass - try: - if hasattr(function_module, "UserValves"): - __user__["valves"] = function_module.UserValves( - **Functions.get_user_valves_by_id_and_user_id( - filter_id, user.id - ) - ) - except Exception as e: - print(e) + try: + body, flags = await chat_completion_files_handler(body) - params = {**params, "__user__": __user__} + contexts.extend(flags.get("contexts", [])) + citations.extend(flags.get("citations", [])) + except Exception as e: + print(e) + pass - if "__id__" in sig.parameters: - params = { - **params, - "__id__": filter_id, - } + # If context is not empty, insert it into the messages + if len(contexts) > 0: + context_string = "/n".join(contexts).strip() + prompt = get_last_user_message(body["messages"]) + body["messages"] = add_or_update_system_message( + rag_template( + rag_app.state.config.RAG_TEMPLATE, context_string, prompt + ), + body["messages"], + ) - if inspect.iscoroutinefunction(inlet): - data = await inlet(**params) - else: - data = inlet(**params) - - except Exception as e: - print(f"Error: {e}") - return JSONResponse( - status_code=status.HTTP_400_BAD_REQUEST, - content={"detail": str(e)}, - ) - - # Set the task model - task_model_id = data["model"] - # Check if the user has a custom task model and use that model - if app.state.MODELS[task_model_id]["owned_by"] == "ollama": - if ( - app.state.config.TASK_MODEL - and app.state.config.TASK_MODEL in app.state.MODELS - ): - task_model_id = app.state.config.TASK_MODEL - else: - if ( - app.state.config.TASK_MODEL_EXTERNAL - and app.state.config.TASK_MODEL_EXTERNAL in app.state.MODELS - ): - task_model_id = app.state.config.TASK_MODEL_EXTERNAL - - prompt = get_last_user_message(data["messages"]) - context = "" - - # If tool_ids field is present, call the functions - if "tool_ids" in data: - print(data["tool_ids"]) - for tool_id in data["tool_ids"]: - print(tool_id) - try: - response, citation, file_handler = ( - await get_function_call_response( - messages=data["messages"], - files=data.get("files", []), - tool_id=tool_id, - template=app.state.config.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE, - task_model_id=task_model_id, - user=user, - ) - ) - - print(file_handler) - if isinstance(response, str): - context += ("\n" if context != "" else "") + response - - if citation: - citations.append(citation) - show_citations = True - - if file_handler: - skip_files = True - - except Exception as e: - print(f"Error: {e}") - del data["tool_ids"] - - print(f"tool_context: {context}") - - # If files field is present, generate RAG completions - # If skip_files is True, skip the RAG completions - if "files" in data: - if not skip_files: - data = {**data} - rag_context, rag_citations = get_rag_context( - files=data["files"], - messages=data["messages"], - embedding_function=rag_app.state.EMBEDDING_FUNCTION, - k=rag_app.state.config.TOP_K, - reranking_function=rag_app.state.sentence_transformer_rf, - r=rag_app.state.config.RELEVANCE_THRESHOLD, - hybrid_search=rag_app.state.config.ENABLE_RAG_HYBRID_SEARCH, - ) - if rag_context: - context += ("\n" if context != "" else "") + rag_context - - log.debug(f"rag_context: {rag_context}, citations: {citations}") - - if rag_citations: - citations.extend(rag_citations) - - del data["files"] - - if show_citations and len(citations) > 0: + # If there are citations, add them to the data_items + if len(citations) > 0: data_items.append({"citations": citations}) - if context != "": - system_prompt = rag_template( - rag_app.state.config.RAG_TEMPLATE, context, prompt - ) - print(system_prompt) - data["messages"] = add_or_update_system_message( - system_prompt, data["messages"] - ) - - modified_body_bytes = json.dumps(data).encode("utf-8") + modified_body_bytes = json.dumps(body).encode("utf-8") # Replace the request body with the modified one request._body = modified_body_bytes # Set custom header to ensure content-length matches new body length @@ -654,9 +764,7 @@ app.add_middleware(ChatCompletionMiddleware) ################################## -def filter_pipeline(payload, user): - user = {"id": user.id, "email": user.email, "name": user.name, "role": user.role} - model_id = payload["model"] +def get_sorted_filters(model_id): filters = [ model for model in app.state.MODELS.values() @@ -672,6 +780,13 @@ def filter_pipeline(payload, user): ) ] sorted_filters = sorted(filters, key=lambda x: x["pipeline"]["priority"]) + return sorted_filters + + +def filter_pipeline(payload, user): + user = {"id": user.id, "email": user.email, "name": user.name, "role": user.role} + model_id = payload["model"] + sorted_filters = get_sorted_filters(model_id) model = app.state.MODELS[model_id] @@ -704,25 +819,12 @@ def filter_pipeline(payload, user): print(f"Connection error: {e}") if r is not None: - try: - res = r.json() - except: - pass + res = r.json() if "detail" in res: raise Exception(r.status_code, res["detail"]) - else: - pass - - if "pipeline" not in app.state.MODELS[model_id]: - if "chat_id" in payload: - del payload["chat_id"] - - if "title" in payload: - del payload["title"] - - if "task" in payload: - del payload["task"] + if "pipeline" not in app.state.MODELS[model_id] and "task" in payload: + del payload["task"] return payload @@ -787,6 +889,14 @@ app.add_middleware( ) +@app.middleware("http") +async def commit_session_after_request(request: Request, call_next): + response = await call_next(request) + log.debug("Commit session after request") + Session.commit() + return response + + @app.middleware("http") async def check_url(request: Request, call_next): if len(app.state.MODELS) == 0: @@ -863,12 +973,16 @@ async def get_all_models(): model["info"] = custom_model.model_dump() else: owned_by = "openai" + pipe = None + for model in models: if ( custom_model.base_model_id == model["id"] or custom_model.base_model_id == model["id"].split(":")[0] ): owned_by = model["owned_by"] + if "pipe" in model: + pipe = model["pipe"] break models.append( @@ -880,11 +994,11 @@ async def get_all_models(): "owned_by": owned_by, "info": custom_model.model_dump(), "preset": True, + **({"pipe": pipe} if pipe is not None else {}), } ) app.state.MODELS = {model["id"]: model for model in models} - webui_app.state.MODELS = app.state.MODELS return models @@ -945,22 +1059,7 @@ async def chat_completed(form_data: dict, user=Depends(get_verified_user)): ) model = app.state.MODELS[model_id] - filters = [ - model - for model in app.state.MODELS.values() - if "pipeline" in model - and "type" in model["pipeline"] - and model["pipeline"]["type"] == "filter" - and ( - model["pipeline"]["pipelines"] == ["*"] - or any( - model_id == target_model_id - for target_model_id in model["pipeline"]["pipelines"] - ) - ) - ] - - sorted_filters = sorted(filters, key=lambda x: x["pipeline"]["priority"]) + sorted_filters = get_sorted_filters(model_id) if "pipeline" in model: sorted_filters = [model] + sorted_filters @@ -1008,6 +1107,25 @@ async def chat_completed(form_data: dict, user=Depends(get_verified_user)): else: pass + async def __event_emitter__(event_data): + await sio.emit( + "chat-events", + { + "chat_id": data["chat_id"], + "message_id": data["id"], + "data": event_data, + }, + to=data["session_id"], + ) + + async def __event_call__(event_data): + response = await sio.call( + "chat-events", + {"chat_id": data["chat_id"], "message_id": data["id"], "data": event_data}, + to=data["session_id"], + ) + return response + def get_priority(function_id): function = Functions.get_function_by_id(function_id) if function is not None and hasattr(function, "valves"): @@ -1032,68 +1150,74 @@ async def chat_completed(form_data: dict, user=Depends(get_verified_user)): for filter_id in filter_ids: filter = Functions.get_function_by_id(filter_id) - if filter: - if filter_id in webui_app.state.FUNCTIONS: - function_module = webui_app.state.FUNCTIONS[filter_id] + if not filter: + continue + + if filter_id in webui_app.state.FUNCTIONS: + function_module = webui_app.state.FUNCTIONS[filter_id] + else: + function_module, _, _ = load_function_module_by_id(filter_id) + webui_app.state.FUNCTIONS[filter_id] = function_module + + if hasattr(function_module, "valves") and hasattr(function_module, "Valves"): + valves = Functions.get_function_valves_by_id(filter_id) + function_module.valves = function_module.Valves( + **(valves if valves else {}) + ) + + if not hasattr(function_module, "outlet"): + continue + try: + outlet = function_module.outlet + + # Get the signature of the function + sig = inspect.signature(outlet) + params = {"body": data} + + # Extra parameters to be passed to the function + extra_params = { + "__model__": model, + "__id__": filter_id, + "__event_emitter__": __event_emitter__, + "__event_call__": __event_call__, + } + + # Add extra params in contained in function signature + for key, value in extra_params.items(): + if key in sig.parameters: + params[key] = value + + if "__user__" in sig.parameters: + __user__ = { + "id": user.id, + "email": user.email, + "name": user.name, + "role": user.role, + } + + try: + if hasattr(function_module, "UserValves"): + __user__["valves"] = function_module.UserValves( + **Functions.get_user_valves_by_id_and_user_id( + filter_id, user.id + ) + ) + except Exception as e: + print(e) + + params = {**params, "__user__": __user__} + + if inspect.iscoroutinefunction(outlet): + data = await outlet(**params) else: - function_module, function_type, frontmatter = ( - load_function_module_by_id(filter_id) - ) - webui_app.state.FUNCTIONS[filter_id] = function_module + data = outlet(**params) - if hasattr(function_module, "valves") and hasattr( - function_module, "Valves" - ): - valves = Functions.get_function_valves_by_id(filter_id) - function_module.valves = function_module.Valves( - **(valves if valves else {}) - ) - - try: - if hasattr(function_module, "outlet"): - outlet = function_module.outlet - - # Get the signature of the function - sig = inspect.signature(outlet) - params = {"body": data} - - if "__user__" in sig.parameters: - __user__ = { - "id": user.id, - "email": user.email, - "name": user.name, - "role": user.role, - } - - try: - if hasattr(function_module, "UserValves"): - __user__["valves"] = function_module.UserValves( - **Functions.get_user_valves_by_id_and_user_id( - filter_id, user.id - ) - ) - except Exception as e: - print(e) - - params = {**params, "__user__": __user__} - - if "__id__" in sig.parameters: - params = { - **params, - "__id__": filter_id, - } - - if inspect.iscoroutinefunction(outlet): - data = await outlet(**params) - else: - data = outlet(**params) - - except Exception as e: - print(f"Error: {e}") - return JSONResponse( - status_code=status.HTTP_400_BAD_REQUEST, - content={"detail": str(e)}, - ) + except Exception as e: + print(f"Error: {e}") + return JSONResponse( + status_code=status.HTTP_400_BAD_REQUEST, + content={"detail": str(e)}, + ) return data @@ -1169,19 +1293,9 @@ async def generate_title(form_data: dict, user=Depends(get_verified_user)): # Check if the user has a custom task model # If the user has a custom task model, use that model - if app.state.MODELS[model_id]["owned_by"] == "ollama": - if app.state.config.TASK_MODEL: - task_model_id = app.state.config.TASK_MODEL - if task_model_id in app.state.MODELS: - model_id = task_model_id - else: - if app.state.config.TASK_MODEL_EXTERNAL: - task_model_id = app.state.config.TASK_MODEL_EXTERNAL - if task_model_id in app.state.MODELS: - model_id = task_model_id + model_id = get_task_model_id(model_id) print(model_id) - model = app.state.MODELS[model_id] template = app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE @@ -1200,7 +1314,7 @@ async def generate_title(form_data: dict, user=Depends(get_verified_user)): "stream": False, "max_tokens": 50, "chat_id": form_data.get("chat_id", None), - "title": True, + "task": TASKS.TITLE_GENERATION, } log.debug(payload) @@ -1213,6 +1327,9 @@ async def generate_title(form_data: dict, user=Depends(get_verified_user)): content={"detail": e.args[1]}, ) + if "chat_id" in payload: + del payload["chat_id"] + return await generate_chat_completions(form_data=payload, user=user) @@ -1235,19 +1352,9 @@ async def generate_search_query(form_data: dict, user=Depends(get_verified_user) # Check if the user has a custom task model # If the user has a custom task model, use that model - if app.state.MODELS[model_id]["owned_by"] == "ollama": - if app.state.config.TASK_MODEL: - task_model_id = app.state.config.TASK_MODEL - if task_model_id in app.state.MODELS: - model_id = task_model_id - else: - if app.state.config.TASK_MODEL_EXTERNAL: - task_model_id = app.state.config.TASK_MODEL_EXTERNAL - if task_model_id in app.state.MODELS: - model_id = task_model_id + model_id = get_task_model_id(model_id) print(model_id) - model = app.state.MODELS[model_id] template = app.state.config.SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE @@ -1260,7 +1367,7 @@ async def generate_search_query(form_data: dict, user=Depends(get_verified_user) "messages": [{"role": "user", "content": content}], "stream": False, "max_tokens": 30, - "task": True, + "task": TASKS.QUERY_GENERATION, } print(payload) @@ -1273,6 +1380,9 @@ async def generate_search_query(form_data: dict, user=Depends(get_verified_user) content={"detail": e.args[1]}, ) + if "chat_id" in payload: + del payload["chat_id"] + return await generate_chat_completions(form_data=payload, user=user) @@ -1289,19 +1399,9 @@ async def generate_emoji(form_data: dict, user=Depends(get_verified_user)): # Check if the user has a custom task model # If the user has a custom task model, use that model - if app.state.MODELS[model_id]["owned_by"] == "ollama": - if app.state.config.TASK_MODEL: - task_model_id = app.state.config.TASK_MODEL - if task_model_id in app.state.MODELS: - model_id = task_model_id - else: - if app.state.config.TASK_MODEL_EXTERNAL: - task_model_id = app.state.config.TASK_MODEL_EXTERNAL - if task_model_id in app.state.MODELS: - model_id = task_model_id + model_id = get_task_model_id(model_id) print(model_id) - model = app.state.MODELS[model_id] template = ''' Your task is to reflect the speaker's likely facial expression through a fitting emoji. Interpret emotions from the message and reflect their facial expression using fitting, diverse emojis (e.g., 😊, 😢, 😡, 😱). @@ -1324,7 +1424,7 @@ Message: """{{prompt}}""" "stream": False, "max_tokens": 4, "chat_id": form_data.get("chat_id", None), - "task": True, + "task": TASKS.EMOJI_GENERATION, } log.debug(payload) @@ -1337,6 +1437,9 @@ Message: """{{prompt}}""" content={"detail": e.args[1]}, ) + if "chat_id" in payload: + del payload["chat_id"] + return await generate_chat_completions(form_data=payload, user=user) @@ -1353,22 +1456,13 @@ async def get_tools_function_calling(form_data: dict, user=Depends(get_verified_ # Check if the user has a custom task model # If the user has a custom task model, use that model - if app.state.MODELS[model_id]["owned_by"] == "ollama": - if app.state.config.TASK_MODEL: - task_model_id = app.state.config.TASK_MODEL - if task_model_id in app.state.MODELS: - model_id = task_model_id - else: - if app.state.config.TASK_MODEL_EXTERNAL: - task_model_id = app.state.config.TASK_MODEL_EXTERNAL - if task_model_id in app.state.MODELS: - model_id = task_model_id + model_id = get_task_model_id(model_id) print(model_id) template = app.state.config.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE try: - context, citation, file_handler = await get_function_call_response( + context, _, _ = await get_function_call_response( form_data["messages"], form_data.get("files", []), form_data["tool_id"], @@ -1432,6 +1526,7 @@ async def upload_pipeline( os.makedirs(upload_folder, exist_ok=True) file_path = os.path.join(upload_folder, file.filename) + r = None try: # Save the uploaded file with open(file_path, "wb") as buffer: @@ -1455,7 +1550,9 @@ async def upload_pipeline( print(f"Connection error: {e}") detail = "Pipeline not found" + status_code = status.HTTP_404_NOT_FOUND if r is not None: + status_code = r.status_code try: res = r.json() if "detail" in res: @@ -1464,7 +1561,7 @@ async def upload_pipeline( pass raise HTTPException( - status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND), + status_code=status_code, detail=detail, ) finally: @@ -1563,8 +1660,6 @@ async def delete_pipeline(form_data: DeletePipelineForm, user=Depends(get_admin_ async def get_pipelines(urlIdx: Optional[int] = None, user=Depends(get_admin_user)): r = None try: - urlIdx - url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx] key = openai_app.state.config.OPENAI_API_KEYS[urlIdx] @@ -1596,7 +1691,9 @@ async def get_pipelines(urlIdx: Optional[int] = None, user=Depends(get_admin_use @app.get("/api/pipelines/{pipeline_id}/valves") async def get_pipeline_valves( - urlIdx: Optional[int], pipeline_id: str, user=Depends(get_admin_user) + urlIdx: Optional[int], + pipeline_id: str, + user=Depends(get_admin_user), ): models = await get_all_models() r = None @@ -1634,7 +1731,9 @@ async def get_pipeline_valves( @app.get("/api/pipelines/{pipeline_id}/valves/spec") async def get_pipeline_valves_spec( - urlIdx: Optional[int], pipeline_id: str, user=Depends(get_admin_user) + urlIdx: Optional[int], + pipeline_id: str, + user=Depends(get_admin_user), ): models = await get_all_models() @@ -1920,7 +2019,8 @@ async def oauth_callback(provider: str, request: Request, response: Response): if existing_user: raise HTTPException(400, detail=ERROR_MESSAGES.EMAIL_TAKEN) - picture_url = user_data.get("picture", "") + picture_claim = webui_app.state.config.OAUTH_PICTURE_CLAIM + picture_url = user_data.get(picture_claim, "") if picture_url: # Download the profile image into a base64 string try: @@ -1940,6 +2040,7 @@ async def oauth_callback(provider: str, request: Request, response: Response): picture_url = "" if not picture_url: picture_url = "/user.png" + username_claim = webui_app.state.config.OAUTH_USERNAME_CLAIM role = ( "admin" if Users.get_num_users() == 0 @@ -1950,7 +2051,7 @@ async def oauth_callback(provider: str, request: Request, response: Response): password=get_password_hash( str(uuid.uuid4()) ), # Random password, not used - name=user_data.get("name", "User"), + name=user_data.get(username_claim, "User"), profile_image_url=picture_url, role=role, oauth_sub=provider_sub, @@ -2008,7 +2109,7 @@ async def get_opensearch_xml(): {WEBUI_NAME} Search {WEBUI_NAME} UTF-8 - {WEBUI_URL}/favicon.png + {WEBUI_URL}/static/favicon.png {WEBUI_URL} @@ -2021,6 +2122,12 @@ async def healthcheck(): return {"status": True} +@app.get("/health/db") +async def healthcheck_with_db(): + Session.execute(text("SELECT 1;")).all() + return {"status": True} + + app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static") app.mount("/cache", StaticFiles(directory=CACHE_DIR), name="cache") diff --git a/backend/migrations/README b/backend/migrations/README new file mode 100644 index 000000000..f1d93dff9 --- /dev/null +++ b/backend/migrations/README @@ -0,0 +1,4 @@ +Generic single-database configuration. + +Create new migrations with +DATABASE_URL= alembic revision --autogenerate -m "a description" diff --git a/backend/migrations/env.py b/backend/migrations/env.py new file mode 100644 index 000000000..7035cf917 --- /dev/null +++ b/backend/migrations/env.py @@ -0,0 +1,96 @@ +import os +from logging.config import fileConfig + +from sqlalchemy import engine_from_config +from sqlalchemy import pool + +from alembic import context + +from apps.webui.models.auths import Auth +from apps.webui.models.chats import Chat +from apps.webui.models.documents import Document +from apps.webui.models.memories import Memory +from apps.webui.models.models import Model +from apps.webui.models.prompts import Prompt +from apps.webui.models.tags import Tag, ChatIdTag +from apps.webui.models.tools import Tool +from apps.webui.models.users import User +from apps.webui.models.files import File +from apps.webui.models.functions import Function + +from config import DATABASE_URL + +# this is the Alembic Config object, which provides +# access to the values within the .ini file in use. +config = context.config + +# Interpret the config file for Python logging. +# This line sets up loggers basically. +if config.config_file_name is not None: + fileConfig(config.config_file_name) + +# add your model's MetaData object here +# for 'autogenerate' support +# from myapp import mymodel +# target_metadata = mymodel.Base.metadata +target_metadata = Auth.metadata + +# other values from the config, defined by the needs of env.py, +# can be acquired: +# my_important_option = config.get_main_option("my_important_option") +# ... etc. + +DB_URL = DATABASE_URL + +if DB_URL: + config.set_main_option("sqlalchemy.url", DB_URL.replace("%", "%%")) + + +def run_migrations_offline() -> None: + """Run migrations in 'offline' mode. + + This configures the context with just a URL + and not an Engine, though an Engine is acceptable + here as well. By skipping the Engine creation + we don't even need a DBAPI to be available. + + Calls to context.execute() here emit the given string to the + script output. + + """ + url = config.get_main_option("sqlalchemy.url") + context.configure( + url=url, + target_metadata=target_metadata, + literal_binds=True, + dialect_opts={"paramstyle": "named"}, + ) + + with context.begin_transaction(): + context.run_migrations() + + +def run_migrations_online() -> None: + """Run migrations in 'online' mode. + + In this scenario we need to create an Engine + and associate a connection with the context. + + """ + connectable = engine_from_config( + config.get_section(config.config_ini_section, {}), + prefix="sqlalchemy.", + poolclass=pool.NullPool, + ) + + with connectable.connect() as connection: + context.configure(connection=connection, target_metadata=target_metadata) + + with context.begin_transaction(): + context.run_migrations() + + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online() diff --git a/backend/migrations/script.py.mako b/backend/migrations/script.py.mako new file mode 100644 index 000000000..5f667ccfe --- /dev/null +++ b/backend/migrations/script.py.mako @@ -0,0 +1,27 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +import apps.webui.internal.db +${imports if imports else ""} + +# revision identifiers, used by Alembic. +revision: str = ${repr(up_revision)} +down_revision: Union[str, None] = ${repr(down_revision)} +branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)} +depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)} + + +def upgrade() -> None: + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + ${downgrades if downgrades else "pass"} diff --git a/backend/migrations/util.py b/backend/migrations/util.py new file mode 100644 index 000000000..401bb94d0 --- /dev/null +++ b/backend/migrations/util.py @@ -0,0 +1,9 @@ +from alembic import op +from sqlalchemy import Inspector + + +def get_existing_tables(): + con = op.get_bind() + inspector = Inspector.from_engine(con) + tables = set(inspector.get_table_names()) + return tables diff --git a/backend/migrations/versions/7e5b5dc7342b_init.py b/backend/migrations/versions/7e5b5dc7342b_init.py new file mode 100644 index 000000000..b82627f5b --- /dev/null +++ b/backend/migrations/versions/7e5b5dc7342b_init.py @@ -0,0 +1,202 @@ +"""init + +Revision ID: 7e5b5dc7342b +Revises: +Create Date: 2024-06-24 13:15:33.808998 + +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +import apps.webui.internal.db +from migrations.util import get_existing_tables + +# revision identifiers, used by Alembic. +revision: str = "7e5b5dc7342b" +down_revision: Union[str, None] = None +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + existing_tables = set(get_existing_tables()) + + # ### commands auto generated by Alembic - please adjust! ### + if "auth" not in existing_tables: + op.create_table( + "auth", + sa.Column("id", sa.String(), nullable=False), + sa.Column("email", sa.String(), nullable=True), + sa.Column("password", sa.Text(), nullable=True), + sa.Column("active", sa.Boolean(), nullable=True), + sa.PrimaryKeyConstraint("id"), + ) + + if "chat" not in existing_tables: + op.create_table( + "chat", + sa.Column("id", sa.String(), nullable=False), + sa.Column("user_id", sa.String(), nullable=True), + sa.Column("title", sa.Text(), nullable=True), + sa.Column("chat", sa.Text(), nullable=True), + sa.Column("created_at", sa.BigInteger(), nullable=True), + sa.Column("updated_at", sa.BigInteger(), nullable=True), + sa.Column("share_id", sa.Text(), nullable=True), + sa.Column("archived", sa.Boolean(), nullable=True), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint("share_id"), + ) + + if "chatidtag" not in existing_tables: + op.create_table( + "chatidtag", + sa.Column("id", sa.String(), nullable=False), + sa.Column("tag_name", sa.String(), nullable=True), + sa.Column("chat_id", sa.String(), nullable=True), + sa.Column("user_id", sa.String(), nullable=True), + sa.Column("timestamp", sa.BigInteger(), nullable=True), + sa.PrimaryKeyConstraint("id"), + ) + + if "document" not in existing_tables: + op.create_table( + "document", + sa.Column("collection_name", sa.String(), nullable=False), + sa.Column("name", sa.String(), nullable=True), + sa.Column("title", sa.Text(), nullable=True), + sa.Column("filename", sa.Text(), nullable=True), + sa.Column("content", sa.Text(), nullable=True), + sa.Column("user_id", sa.String(), nullable=True), + sa.Column("timestamp", sa.BigInteger(), nullable=True), + sa.PrimaryKeyConstraint("collection_name"), + sa.UniqueConstraint("name"), + ) + + if "file" not in existing_tables: + op.create_table( + "file", + sa.Column("id", sa.String(), nullable=False), + sa.Column("user_id", sa.String(), nullable=True), + sa.Column("filename", sa.Text(), nullable=True), + sa.Column("meta", apps.webui.internal.db.JSONField(), nullable=True), + sa.Column("created_at", sa.BigInteger(), nullable=True), + sa.PrimaryKeyConstraint("id"), + ) + + if "function" not in existing_tables: + op.create_table( + "function", + sa.Column("id", sa.String(), nullable=False), + sa.Column("user_id", sa.String(), nullable=True), + sa.Column("name", sa.Text(), nullable=True), + sa.Column("type", sa.Text(), nullable=True), + sa.Column("content", sa.Text(), nullable=True), + sa.Column("meta", apps.webui.internal.db.JSONField(), nullable=True), + sa.Column("valves", apps.webui.internal.db.JSONField(), nullable=True), + sa.Column("is_active", sa.Boolean(), nullable=True), + sa.Column("is_global", sa.Boolean(), nullable=True), + sa.Column("updated_at", sa.BigInteger(), nullable=True), + sa.Column("created_at", sa.BigInteger(), nullable=True), + sa.PrimaryKeyConstraint("id"), + ) + + if "memory" not in existing_tables: + op.create_table( + "memory", + sa.Column("id", sa.String(), nullable=False), + sa.Column("user_id", sa.String(), nullable=True), + sa.Column("content", sa.Text(), nullable=True), + sa.Column("updated_at", sa.BigInteger(), nullable=True), + sa.Column("created_at", sa.BigInteger(), nullable=True), + sa.PrimaryKeyConstraint("id"), + ) + + if "model" not in existing_tables: + op.create_table( + "model", + sa.Column("id", sa.Text(), nullable=False), + sa.Column("user_id", sa.Text(), nullable=True), + sa.Column("base_model_id", sa.Text(), nullable=True), + sa.Column("name", sa.Text(), nullable=True), + sa.Column("params", apps.webui.internal.db.JSONField(), nullable=True), + sa.Column("meta", apps.webui.internal.db.JSONField(), nullable=True), + sa.Column("updated_at", sa.BigInteger(), nullable=True), + sa.Column("created_at", sa.BigInteger(), nullable=True), + sa.PrimaryKeyConstraint("id"), + ) + + if "prompt" not in existing_tables: + op.create_table( + "prompt", + sa.Column("command", sa.String(), nullable=False), + sa.Column("user_id", sa.String(), nullable=True), + sa.Column("title", sa.Text(), nullable=True), + sa.Column("content", sa.Text(), nullable=True), + sa.Column("timestamp", sa.BigInteger(), nullable=True), + sa.PrimaryKeyConstraint("command"), + ) + + if "tag" not in existing_tables: + op.create_table( + "tag", + sa.Column("id", sa.String(), nullable=False), + sa.Column("name", sa.String(), nullable=True), + sa.Column("user_id", sa.String(), nullable=True), + sa.Column("data", sa.Text(), nullable=True), + sa.PrimaryKeyConstraint("id"), + ) + + if "tool" not in existing_tables: + op.create_table( + "tool", + sa.Column("id", sa.String(), nullable=False), + sa.Column("user_id", sa.String(), nullable=True), + sa.Column("name", sa.Text(), nullable=True), + sa.Column("content", sa.Text(), nullable=True), + sa.Column("specs", apps.webui.internal.db.JSONField(), nullable=True), + sa.Column("meta", apps.webui.internal.db.JSONField(), nullable=True), + sa.Column("valves", apps.webui.internal.db.JSONField(), nullable=True), + sa.Column("updated_at", sa.BigInteger(), nullable=True), + sa.Column("created_at", sa.BigInteger(), nullable=True), + sa.PrimaryKeyConstraint("id"), + ) + + if "user" not in existing_tables: + op.create_table( + "user", + sa.Column("id", sa.String(), nullable=False), + sa.Column("name", sa.String(), nullable=True), + sa.Column("email", sa.String(), nullable=True), + sa.Column("role", sa.String(), nullable=True), + sa.Column("profile_image_url", sa.Text(), nullable=True), + sa.Column("last_active_at", sa.BigInteger(), nullable=True), + sa.Column("updated_at", sa.BigInteger(), nullable=True), + sa.Column("created_at", sa.BigInteger(), nullable=True), + sa.Column("api_key", sa.String(), nullable=True), + sa.Column("settings", apps.webui.internal.db.JSONField(), nullable=True), + sa.Column("info", apps.webui.internal.db.JSONField(), nullable=True), + sa.Column("oauth_sub", sa.Text(), nullable=True), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint("api_key"), + sa.UniqueConstraint("oauth_sub"), + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table("user") + op.drop_table("tool") + op.drop_table("tag") + op.drop_table("prompt") + op.drop_table("model") + op.drop_table("memory") + op.drop_table("function") + op.drop_table("file") + op.drop_table("document") + op.drop_table("chatidtag") + op.drop_table("chat") + op.drop_table("auth") + # ### end Alembic commands ### diff --git a/backend/requirements.txt b/backend/requirements.txt index 329637ef4..1a04cfcb2 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -10,9 +10,11 @@ python-socketio==5.11.3 python-jose==3.3.0 passlib[bcrypt]==1.7.4 -requests==2.32.2 +requests==2.32.3 aiohttp==3.9.5 -peewee==3.17.5 +sqlalchemy==2.0.30 +alembic==1.13.2 +peewee==3.17.6 peewee-migrate==1.12.2 psycopg2-binary==2.9.9 PyMySQL==1.1.1 @@ -30,26 +32,26 @@ openai anthropic google-generativeai==0.5.4 -langchain==0.2.0 -langchain-community==0.2.0 +langchain==0.2.6 +langchain-community==0.2.6 langchain-chroma==0.1.2 fake-useragent==1.5.1 chromadb==0.5.3 -sentence-transformers==2.7.0 +sentence-transformers==3.0.1 pypdf==4.2.0 docx2txt==0.8 python-pptx==0.6.23 -unstructured==0.14.0 +unstructured==0.14.9 Markdown==3.6 pypandoc==1.13 pandas==2.2.2 -openpyxl==3.1.2 +openpyxl==3.1.5 pyxlsb==1.0.10 xlrd==2.0.1 validators==0.28.1 -opencv-python-headless==4.9.0.80 +opencv-python-headless==4.10.0.84 rapidocr-onnxruntime==1.3.22 fpdf2==2.7.9 @@ -61,10 +63,15 @@ PyJWT[crypto]==2.8.0 authlib==1.3.1 black==24.4.2 -langfuse==2.33.0 +langfuse==2.38.0 youtube-transcript-api==0.6.2 pytube==15.0.0 extract_msg pydub -duckduckgo-search~=6.1.7 \ No newline at end of file +duckduckgo-search~=6.1.7 + +## Tests +docker~=7.1.0 +pytest~=8.2.2 +pytest-docker~=3.1.1 diff --git a/backend/static/splash.png b/backend/static/splash.png new file mode 100644 index 000000000..389196ca6 Binary files /dev/null and b/backend/static/splash.png differ diff --git a/backend/test/__init__.py b/backend/test/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/backend/test/apps/webui/routers/test_auths.py b/backend/test/apps/webui/routers/test_auths.py new file mode 100644 index 000000000..3a8695a69 --- /dev/null +++ b/backend/test/apps/webui/routers/test_auths.py @@ -0,0 +1,202 @@ +import pytest + +from test.util.abstract_integration_test import AbstractPostgresTest +from test.util.mock_user import mock_webui_user + + +class TestAuths(AbstractPostgresTest): + BASE_PATH = "/api/v1/auths" + + def setup_class(cls): + super().setup_class() + from apps.webui.models.users import Users + from apps.webui.models.auths import Auths + + cls.users = Users + cls.auths = Auths + + def test_get_session_user(self): + with mock_webui_user(): + response = self.fast_api_client.get(self.create_url("")) + assert response.status_code == 200 + assert response.json() == { + "id": "1", + "name": "John Doe", + "email": "john.doe@openwebui.com", + "role": "user", + "profile_image_url": "/user.png", + } + + def test_update_profile(self): + from utils.utils import get_password_hash + + user = self.auths.insert_new_auth( + email="john.doe@openwebui.com", + password=get_password_hash("old_password"), + name="John Doe", + profile_image_url="/user.png", + role="user", + ) + + with mock_webui_user(id=user.id): + response = self.fast_api_client.post( + self.create_url("/update/profile"), + json={"name": "John Doe 2", "profile_image_url": "/user2.png"}, + ) + assert response.status_code == 200 + db_user = self.users.get_user_by_id(user.id) + assert db_user.name == "John Doe 2" + assert db_user.profile_image_url == "/user2.png" + + def test_update_password(self): + from utils.utils import get_password_hash + + user = self.auths.insert_new_auth( + email="john.doe@openwebui.com", + password=get_password_hash("old_password"), + name="John Doe", + profile_image_url="/user.png", + role="user", + ) + + with mock_webui_user(id=user.id): + response = self.fast_api_client.post( + self.create_url("/update/password"), + json={"password": "old_password", "new_password": "new_password"}, + ) + assert response.status_code == 200 + + old_auth = self.auths.authenticate_user( + "john.doe@openwebui.com", "old_password" + ) + assert old_auth is None + new_auth = self.auths.authenticate_user( + "john.doe@openwebui.com", "new_password" + ) + assert new_auth is not None + + def test_signin(self): + from utils.utils import get_password_hash + + user = self.auths.insert_new_auth( + email="john.doe@openwebui.com", + password=get_password_hash("password"), + name="John Doe", + profile_image_url="/user.png", + role="user", + ) + response = self.fast_api_client.post( + self.create_url("/signin"), + json={"email": "john.doe@openwebui.com", "password": "password"}, + ) + assert response.status_code == 200 + data = response.json() + assert data["id"] == user.id + assert data["name"] == "John Doe" + assert data["email"] == "john.doe@openwebui.com" + assert data["role"] == "user" + assert data["profile_image_url"] == "/user.png" + assert data["token"] is not None and len(data["token"]) > 0 + assert data["token_type"] == "Bearer" + + def test_signup(self): + response = self.fast_api_client.post( + self.create_url("/signup"), + json={ + "name": "John Doe", + "email": "john.doe@openwebui.com", + "password": "password", + }, + ) + assert response.status_code == 200 + data = response.json() + assert data["id"] is not None and len(data["id"]) > 0 + assert data["name"] == "John Doe" + assert data["email"] == "john.doe@openwebui.com" + assert data["role"] in ["admin", "user", "pending"] + assert data["profile_image_url"] == "/user.png" + assert data["token"] is not None and len(data["token"]) > 0 + assert data["token_type"] == "Bearer" + + def test_add_user(self): + with mock_webui_user(): + response = self.fast_api_client.post( + self.create_url("/add"), + json={ + "name": "John Doe 2", + "email": "john.doe2@openwebui.com", + "password": "password2", + "role": "admin", + }, + ) + assert response.status_code == 200 + data = response.json() + assert data["id"] is not None and len(data["id"]) > 0 + assert data["name"] == "John Doe 2" + assert data["email"] == "john.doe2@openwebui.com" + assert data["role"] == "admin" + assert data["profile_image_url"] == "/user.png" + assert data["token"] is not None and len(data["token"]) > 0 + assert data["token_type"] == "Bearer" + + def test_get_admin_details(self): + self.auths.insert_new_auth( + email="john.doe@openwebui.com", + password="password", + name="John Doe", + profile_image_url="/user.png", + role="admin", + ) + with mock_webui_user(): + response = self.fast_api_client.get(self.create_url("/admin/details")) + + assert response.status_code == 200 + assert response.json() == { + "name": "John Doe", + "email": "john.doe@openwebui.com", + } + + def test_create_api_key_(self): + user = self.auths.insert_new_auth( + email="john.doe@openwebui.com", + password="password", + name="John Doe", + profile_image_url="/user.png", + role="admin", + ) + with mock_webui_user(id=user.id): + response = self.fast_api_client.post(self.create_url("/api_key")) + assert response.status_code == 200 + data = response.json() + assert data["api_key"] is not None + assert len(data["api_key"]) > 0 + + def test_delete_api_key(self): + user = self.auths.insert_new_auth( + email="john.doe@openwebui.com", + password="password", + name="John Doe", + profile_image_url="/user.png", + role="admin", + ) + self.users.update_user_api_key_by_id(user.id, "abc") + with mock_webui_user(id=user.id): + response = self.fast_api_client.delete(self.create_url("/api_key")) + assert response.status_code == 200 + assert response.json() == True + db_user = self.users.get_user_by_id(user.id) + assert db_user.api_key is None + + def test_get_api_key(self): + user = self.auths.insert_new_auth( + email="john.doe@openwebui.com", + password="password", + name="John Doe", + profile_image_url="/user.png", + role="admin", + ) + self.users.update_user_api_key_by_id(user.id, "abc") + with mock_webui_user(id=user.id): + response = self.fast_api_client.get(self.create_url("/api_key")) + assert response.status_code == 200 + assert response.json() == {"api_key": "abc"} diff --git a/backend/test/apps/webui/routers/test_chats.py b/backend/test/apps/webui/routers/test_chats.py new file mode 100644 index 000000000..f4661b625 --- /dev/null +++ b/backend/test/apps/webui/routers/test_chats.py @@ -0,0 +1,238 @@ +import uuid + +from test.util.abstract_integration_test import AbstractPostgresTest +from test.util.mock_user import mock_webui_user + + +class TestChats(AbstractPostgresTest): + + BASE_PATH = "/api/v1/chats" + + def setup_class(cls): + super().setup_class() + + def setup_method(self): + super().setup_method() + from apps.webui.models.chats import ChatForm + from apps.webui.models.chats import Chats + + self.chats = Chats + self.chats.insert_new_chat( + "2", + ChatForm( + **{ + "chat": { + "name": "chat1", + "description": "chat1 description", + "tags": ["tag1", "tag2"], + "history": {"currentId": "1", "messages": []}, + } + } + ), + ) + + def test_get_session_user_chat_list(self): + with mock_webui_user(id="2"): + response = self.fast_api_client.get(self.create_url("/")) + assert response.status_code == 200 + first_chat = response.json()[0] + assert first_chat["id"] is not None + assert first_chat["title"] == "New Chat" + assert first_chat["created_at"] is not None + assert first_chat["updated_at"] is not None + + def test_delete_all_user_chats(self): + with mock_webui_user(id="2"): + response = self.fast_api_client.delete(self.create_url("/")) + assert response.status_code == 200 + assert len(self.chats.get_chats()) == 0 + + def test_get_user_chat_list_by_user_id(self): + with mock_webui_user(id="3"): + response = self.fast_api_client.get(self.create_url("/list/user/2")) + assert response.status_code == 200 + first_chat = response.json()[0] + assert first_chat["id"] is not None + assert first_chat["title"] == "New Chat" + assert first_chat["created_at"] is not None + assert first_chat["updated_at"] is not None + + def test_create_new_chat(self): + with mock_webui_user(id="2"): + response = self.fast_api_client.post( + self.create_url("/new"), + json={ + "chat": { + "name": "chat2", + "description": "chat2 description", + "tags": ["tag1", "tag2"], + } + }, + ) + assert response.status_code == 200 + data = response.json() + assert data["archived"] is False + assert data["chat"] == { + "name": "chat2", + "description": "chat2 description", + "tags": ["tag1", "tag2"], + } + assert data["user_id"] == "2" + assert data["id"] is not None + assert data["share_id"] is None + assert data["title"] == "New Chat" + assert data["updated_at"] is not None + assert data["created_at"] is not None + assert len(self.chats.get_chats()) == 2 + + def test_get_user_chats(self): + self.test_get_session_user_chat_list() + + def test_get_user_archived_chats(self): + self.chats.archive_all_chats_by_user_id("2") + from apps.webui.internal.db import Session + + Session.commit() + with mock_webui_user(id="2"): + response = self.fast_api_client.get(self.create_url("/all/archived")) + assert response.status_code == 200 + first_chat = response.json()[0] + assert first_chat["id"] is not None + assert first_chat["title"] == "New Chat" + assert first_chat["created_at"] is not None + assert first_chat["updated_at"] is not None + + def test_get_all_user_chats_in_db(self): + with mock_webui_user(id="4"): + response = self.fast_api_client.get(self.create_url("/all/db")) + assert response.status_code == 200 + assert len(response.json()) == 1 + + def test_get_archived_session_user_chat_list(self): + self.test_get_user_archived_chats() + + def test_archive_all_chats(self): + with mock_webui_user(id="2"): + response = self.fast_api_client.post(self.create_url("/archive/all")) + assert response.status_code == 200 + assert len(self.chats.get_archived_chats_by_user_id("2")) == 1 + + def test_get_shared_chat_by_id(self): + chat_id = self.chats.get_chats()[0].id + self.chats.update_chat_share_id_by_id(chat_id, chat_id) + with mock_webui_user(id="2"): + response = self.fast_api_client.get(self.create_url(f"/share/{chat_id}")) + assert response.status_code == 200 + data = response.json() + assert data["id"] == chat_id + assert data["chat"] == { + "name": "chat1", + "description": "chat1 description", + "tags": ["tag1", "tag2"], + "history": {"currentId": "1", "messages": []}, + } + assert data["id"] == chat_id + assert data["share_id"] == chat_id + assert data["title"] == "New Chat" + + def test_get_chat_by_id(self): + chat_id = self.chats.get_chats()[0].id + with mock_webui_user(id="2"): + response = self.fast_api_client.get(self.create_url(f"/{chat_id}")) + assert response.status_code == 200 + data = response.json() + assert data["id"] == chat_id + assert data["chat"] == { + "name": "chat1", + "description": "chat1 description", + "tags": ["tag1", "tag2"], + "history": {"currentId": "1", "messages": []}, + } + assert data["share_id"] is None + assert data["title"] == "New Chat" + assert data["user_id"] == "2" + + def test_update_chat_by_id(self): + chat_id = self.chats.get_chats()[0].id + with mock_webui_user(id="2"): + response = self.fast_api_client.post( + self.create_url(f"/{chat_id}"), + json={ + "chat": { + "name": "chat2", + "description": "chat2 description", + "tags": ["tag2", "tag4"], + "title": "Just another title", + } + }, + ) + assert response.status_code == 200 + data = response.json() + assert data["id"] == chat_id + assert data["chat"] == { + "name": "chat2", + "title": "Just another title", + "description": "chat2 description", + "tags": ["tag2", "tag4"], + "history": {"currentId": "1", "messages": []}, + } + assert data["share_id"] is None + assert data["title"] == "Just another title" + assert data["user_id"] == "2" + + def test_delete_chat_by_id(self): + chat_id = self.chats.get_chats()[0].id + with mock_webui_user(id="2"): + response = self.fast_api_client.delete(self.create_url(f"/{chat_id}")) + assert response.status_code == 200 + assert response.json() is True + + def test_clone_chat_by_id(self): + chat_id = self.chats.get_chats()[0].id + with mock_webui_user(id="2"): + response = self.fast_api_client.get(self.create_url(f"/{chat_id}/clone")) + + assert response.status_code == 200 + data = response.json() + assert data["id"] != chat_id + assert data["chat"] == { + "branchPointMessageId": "1", + "description": "chat1 description", + "history": {"currentId": "1", "messages": []}, + "name": "chat1", + "originalChatId": chat_id, + "tags": ["tag1", "tag2"], + "title": "Clone of New Chat", + } + assert data["share_id"] is None + assert data["title"] == "Clone of New Chat" + assert data["user_id"] == "2" + + def test_archive_chat_by_id(self): + chat_id = self.chats.get_chats()[0].id + with mock_webui_user(id="2"): + response = self.fast_api_client.get(self.create_url(f"/{chat_id}/archive")) + assert response.status_code == 200 + + chat = self.chats.get_chat_by_id(chat_id) + assert chat.archived is True + + def test_share_chat_by_id(self): + chat_id = self.chats.get_chats()[0].id + with mock_webui_user(id="2"): + response = self.fast_api_client.post(self.create_url(f"/{chat_id}/share")) + assert response.status_code == 200 + + chat = self.chats.get_chat_by_id(chat_id) + assert chat.share_id is not None + + def test_delete_shared_chat_by_id(self): + chat_id = self.chats.get_chats()[0].id + share_id = str(uuid.uuid4()) + self.chats.update_chat_share_id_by_id(chat_id, share_id) + with mock_webui_user(id="2"): + response = self.fast_api_client.delete(self.create_url(f"/{chat_id}/share")) + assert response.status_code + + chat = self.chats.get_chat_by_id(chat_id) + assert chat.share_id is None diff --git a/backend/test/apps/webui/routers/test_documents.py b/backend/test/apps/webui/routers/test_documents.py new file mode 100644 index 000000000..14ca339fd --- /dev/null +++ b/backend/test/apps/webui/routers/test_documents.py @@ -0,0 +1,106 @@ +from test.util.abstract_integration_test import AbstractPostgresTest +from test.util.mock_user import mock_webui_user + + +class TestDocuments(AbstractPostgresTest): + + BASE_PATH = "/api/v1/documents" + + def setup_class(cls): + super().setup_class() + from apps.webui.models.documents import Documents + + cls.documents = Documents + + def test_documents(self): + # Empty database + assert len(self.documents.get_docs()) == 0 + with mock_webui_user(id="2"): + response = self.fast_api_client.get(self.create_url("/")) + assert response.status_code == 200 + assert len(response.json()) == 0 + + # Create a new document + with mock_webui_user(id="2"): + response = self.fast_api_client.post( + self.create_url("/create"), + json={ + "name": "doc_name", + "title": "doc title", + "collection_name": "custom collection", + "filename": "doc_name.pdf", + "content": "", + }, + ) + assert response.status_code == 200 + assert response.json()["name"] == "doc_name" + assert len(self.documents.get_docs()) == 1 + + # Get the document + with mock_webui_user(id="2"): + response = self.fast_api_client.get(self.create_url("/doc?name=doc_name")) + assert response.status_code == 200 + data = response.json() + assert data["collection_name"] == "custom collection" + assert data["name"] == "doc_name" + assert data["title"] == "doc title" + assert data["filename"] == "doc_name.pdf" + assert data["content"] == {} + + # Create another document + with mock_webui_user(id="2"): + response = self.fast_api_client.post( + self.create_url("/create"), + json={ + "name": "doc_name 2", + "title": "doc title 2", + "collection_name": "custom collection 2", + "filename": "doc_name2.pdf", + "content": "", + }, + ) + assert response.status_code == 200 + assert response.json()["name"] == "doc_name 2" + assert len(self.documents.get_docs()) == 2 + + # Get all documents + with mock_webui_user(id="2"): + response = self.fast_api_client.get(self.create_url("/")) + assert response.status_code == 200 + assert len(response.json()) == 2 + + # Update the first document + with mock_webui_user(id="2"): + response = self.fast_api_client.post( + self.create_url("/doc/update?name=doc_name"), + json={"name": "doc_name rework", "title": "updated title"}, + ) + assert response.status_code == 200 + data = response.json() + assert data["name"] == "doc_name rework" + assert data["title"] == "updated title" + + # Tag the first document + with mock_webui_user(id="2"): + response = self.fast_api_client.post( + self.create_url("/doc/tags"), + json={ + "name": "doc_name rework", + "tags": [{"name": "testing-tag"}, {"name": "another-tag"}], + }, + ) + assert response.status_code == 200 + data = response.json() + assert data["name"] == "doc_name rework" + assert data["content"] == { + "tags": [{"name": "testing-tag"}, {"name": "another-tag"}] + } + assert len(self.documents.get_docs()) == 2 + + # Delete the first document + with mock_webui_user(id="2"): + response = self.fast_api_client.delete( + self.create_url("/doc/delete?name=doc_name rework") + ) + assert response.status_code == 200 + assert len(self.documents.get_docs()) == 1 diff --git a/backend/test/apps/webui/routers/test_models.py b/backend/test/apps/webui/routers/test_models.py new file mode 100644 index 000000000..410c4516a --- /dev/null +++ b/backend/test/apps/webui/routers/test_models.py @@ -0,0 +1,62 @@ +from test.util.abstract_integration_test import AbstractPostgresTest +from test.util.mock_user import mock_webui_user + + +class TestModels(AbstractPostgresTest): + + BASE_PATH = "/api/v1/models" + + def setup_class(cls): + super().setup_class() + from apps.webui.models.models import Model + + cls.models = Model + + def test_models(self): + with mock_webui_user(id="2"): + response = self.fast_api_client.get(self.create_url("/")) + assert response.status_code == 200 + assert len(response.json()) == 0 + + with mock_webui_user(id="2"): + response = self.fast_api_client.post( + self.create_url("/add"), + json={ + "id": "my-model", + "base_model_id": "base-model-id", + "name": "Hello World", + "meta": { + "profile_image_url": "/static/favicon.png", + "description": "description", + "capabilities": None, + "model_config": {}, + }, + "params": {}, + }, + ) + assert response.status_code == 200 + + with mock_webui_user(id="2"): + response = self.fast_api_client.get(self.create_url("/")) + assert response.status_code == 200 + assert len(response.json()) == 1 + + with mock_webui_user(id="2"): + response = self.fast_api_client.get( + self.create_url(query_params={"id": "my-model"}) + ) + assert response.status_code == 200 + data = response.json()[0] + assert data["id"] == "my-model" + assert data["name"] == "Hello World" + + with mock_webui_user(id="2"): + response = self.fast_api_client.delete( + self.create_url("/delete?id=my-model") + ) + assert response.status_code == 200 + + with mock_webui_user(id="2"): + response = self.fast_api_client.get(self.create_url("/")) + assert response.status_code == 200 + assert len(response.json()) == 0 diff --git a/backend/test/apps/webui/routers/test_prompts.py b/backend/test/apps/webui/routers/test_prompts.py new file mode 100644 index 000000000..9f47be992 --- /dev/null +++ b/backend/test/apps/webui/routers/test_prompts.py @@ -0,0 +1,92 @@ +from test.util.abstract_integration_test import AbstractPostgresTest +from test.util.mock_user import mock_webui_user + + +class TestPrompts(AbstractPostgresTest): + + BASE_PATH = "/api/v1/prompts" + + def test_prompts(self): + # Get all prompts + with mock_webui_user(id="2"): + response = self.fast_api_client.get(self.create_url("/")) + assert response.status_code == 200 + assert len(response.json()) == 0 + + # Create a two new prompts + with mock_webui_user(id="2"): + response = self.fast_api_client.post( + self.create_url("/create"), + json={ + "command": "/my-command", + "title": "Hello World", + "content": "description", + }, + ) + assert response.status_code == 200 + with mock_webui_user(id="3"): + response = self.fast_api_client.post( + self.create_url("/create"), + json={ + "command": "/my-command2", + "title": "Hello World 2", + "content": "description 2", + }, + ) + assert response.status_code == 200 + + # Get all prompts + with mock_webui_user(id="2"): + response = self.fast_api_client.get(self.create_url("/")) + assert response.status_code == 200 + assert len(response.json()) == 2 + + # Get prompt by command + with mock_webui_user(id="2"): + response = self.fast_api_client.get(self.create_url("/command/my-command")) + assert response.status_code == 200 + data = response.json() + assert data["command"] == "/my-command" + assert data["title"] == "Hello World" + assert data["content"] == "description" + assert data["user_id"] == "2" + + # Update prompt + with mock_webui_user(id="2"): + response = self.fast_api_client.post( + self.create_url("/command/my-command2/update"), + json={ + "command": "irrelevant for request", + "title": "Hello World Updated", + "content": "description Updated", + }, + ) + assert response.status_code == 200 + data = response.json() + assert data["command"] == "/my-command2" + assert data["title"] == "Hello World Updated" + assert data["content"] == "description Updated" + assert data["user_id"] == "3" + + # Get prompt by command + with mock_webui_user(id="2"): + response = self.fast_api_client.get(self.create_url("/command/my-command2")) + assert response.status_code == 200 + data = response.json() + assert data["command"] == "/my-command2" + assert data["title"] == "Hello World Updated" + assert data["content"] == "description Updated" + assert data["user_id"] == "3" + + # Delete prompt + with mock_webui_user(id="2"): + response = self.fast_api_client.delete( + self.create_url("/command/my-command/delete") + ) + assert response.status_code == 200 + + # Get all prompts + with mock_webui_user(id="2"): + response = self.fast_api_client.get(self.create_url("/")) + assert response.status_code == 200 + assert len(response.json()) == 1 diff --git a/backend/test/apps/webui/routers/test_users.py b/backend/test/apps/webui/routers/test_users.py new file mode 100644 index 000000000..9736b4d32 --- /dev/null +++ b/backend/test/apps/webui/routers/test_users.py @@ -0,0 +1,168 @@ +from test.util.abstract_integration_test import AbstractPostgresTest +from test.util.mock_user import mock_webui_user + + +def _get_user_by_id(data, param): + return next((item for item in data if item["id"] == param), None) + + +def _assert_user(data, id, **kwargs): + user = _get_user_by_id(data, id) + assert user is not None + comparison_data = { + "name": f"user {id}", + "email": f"user{id}@openwebui.com", + "profile_image_url": f"/user{id}.png", + "role": "user", + **kwargs, + } + for key, value in comparison_data.items(): + assert user[key] == value + + +class TestUsers(AbstractPostgresTest): + + BASE_PATH = "/api/v1/users" + + def setup_class(cls): + super().setup_class() + from apps.webui.models.users import Users + + cls.users = Users + + def setup_method(self): + super().setup_method() + self.users.insert_new_user( + id="1", + name="user 1", + email="user1@openwebui.com", + profile_image_url="/user1.png", + role="user", + ) + self.users.insert_new_user( + id="2", + name="user 2", + email="user2@openwebui.com", + profile_image_url="/user2.png", + role="user", + ) + + def test_users(self): + # Get all users + with mock_webui_user(id="3"): + response = self.fast_api_client.get(self.create_url("")) + assert response.status_code == 200 + assert len(response.json()) == 2 + data = response.json() + _assert_user(data, "1") + _assert_user(data, "2") + + # update role + with mock_webui_user(id="3"): + response = self.fast_api_client.post( + self.create_url("/update/role"), json={"id": "2", "role": "admin"} + ) + assert response.status_code == 200 + _assert_user([response.json()], "2", role="admin") + + # Get all users + with mock_webui_user(id="3"): + response = self.fast_api_client.get(self.create_url("")) + assert response.status_code == 200 + assert len(response.json()) == 2 + data = response.json() + _assert_user(data, "1") + _assert_user(data, "2", role="admin") + + # Get (empty) user settings + with mock_webui_user(id="2"): + response = self.fast_api_client.get(self.create_url("/user/settings")) + assert response.status_code == 200 + assert response.json() is None + + # Update user settings + with mock_webui_user(id="2"): + response = self.fast_api_client.post( + self.create_url("/user/settings/update"), + json={ + "ui": {"attr1": "value1", "attr2": "value2"}, + "model_config": {"attr3": "value3", "attr4": "value4"}, + }, + ) + assert response.status_code == 200 + + # Get user settings + with mock_webui_user(id="2"): + response = self.fast_api_client.get(self.create_url("/user/settings")) + assert response.status_code == 200 + assert response.json() == { + "ui": {"attr1": "value1", "attr2": "value2"}, + "model_config": {"attr3": "value3", "attr4": "value4"}, + } + + # Get (empty) user info + with mock_webui_user(id="1"): + response = self.fast_api_client.get(self.create_url("/user/info")) + assert response.status_code == 200 + assert response.json() is None + + # Update user info + with mock_webui_user(id="1"): + response = self.fast_api_client.post( + self.create_url("/user/info/update"), + json={"attr1": "value1", "attr2": "value2"}, + ) + assert response.status_code == 200 + + # Get user info + with mock_webui_user(id="1"): + response = self.fast_api_client.get(self.create_url("/user/info")) + assert response.status_code == 200 + assert response.json() == {"attr1": "value1", "attr2": "value2"} + + # Get user by id + with mock_webui_user(id="1"): + response = self.fast_api_client.get(self.create_url("/2")) + assert response.status_code == 200 + assert response.json() == {"name": "user 2", "profile_image_url": "/user2.png"} + + # Update user by id + with mock_webui_user(id="1"): + response = self.fast_api_client.post( + self.create_url("/2/update"), + json={ + "name": "user 2 updated", + "email": "user2-updated@openwebui.com", + "profile_image_url": "/user2-updated.png", + }, + ) + assert response.status_code == 200 + + # Get all users + with mock_webui_user(id="3"): + response = self.fast_api_client.get(self.create_url("")) + assert response.status_code == 200 + assert len(response.json()) == 2 + data = response.json() + _assert_user(data, "1") + _assert_user( + data, + "2", + role="admin", + name="user 2 updated", + email="user2-updated@openwebui.com", + profile_image_url="/user2-updated.png", + ) + + # Delete user by id + with mock_webui_user(id="1"): + response = self.fast_api_client.delete(self.create_url("/2")) + assert response.status_code == 200 + + # Get all users + with mock_webui_user(id="3"): + response = self.fast_api_client.get(self.create_url("")) + assert response.status_code == 200 + assert len(response.json()) == 1 + data = response.json() + _assert_user(data, "1") diff --git a/backend/test/util/abstract_integration_test.py b/backend/test/util/abstract_integration_test.py new file mode 100644 index 000000000..8535221a8 --- /dev/null +++ b/backend/test/util/abstract_integration_test.py @@ -0,0 +1,161 @@ +import logging +import os +import time + +import docker +import pytest +from docker import DockerClient +from pytest_docker.plugin import get_docker_ip +from fastapi.testclient import TestClient +from sqlalchemy import text, create_engine + + +log = logging.getLogger(__name__) + + +def get_fast_api_client(): + from main import app + + with TestClient(app) as c: + return c + + +class AbstractIntegrationTest: + BASE_PATH = None + + def create_url(self, path="", query_params=None): + if self.BASE_PATH is None: + raise Exception("BASE_PATH is not set") + parts = self.BASE_PATH.split("/") + parts = [part.strip() for part in parts if part.strip() != ""] + path_parts = path.split("/") + path_parts = [part.strip() for part in path_parts if part.strip() != ""] + query_parts = "" + if query_params: + query_parts = "&".join( + [f"{key}={value}" for key, value in query_params.items()] + ) + query_parts = f"?{query_parts}" + return "/".join(parts + path_parts) + query_parts + + @classmethod + def setup_class(cls): + pass + + def setup_method(self): + pass + + @classmethod + def teardown_class(cls): + pass + + def teardown_method(self): + pass + + +class AbstractPostgresTest(AbstractIntegrationTest): + DOCKER_CONTAINER_NAME = "postgres-test-container-will-get-deleted" + docker_client: DockerClient + + @classmethod + def _create_db_url(cls, env_vars_postgres: dict) -> str: + host = get_docker_ip() + user = env_vars_postgres["POSTGRES_USER"] + pw = env_vars_postgres["POSTGRES_PASSWORD"] + port = 8081 + db = env_vars_postgres["POSTGRES_DB"] + return f"postgresql://{user}:{pw}@{host}:{port}/{db}" + + @classmethod + def setup_class(cls): + super().setup_class() + try: + env_vars_postgres = { + "POSTGRES_USER": "user", + "POSTGRES_PASSWORD": "example", + "POSTGRES_DB": "openwebui", + } + cls.docker_client = docker.from_env() + cls.docker_client.containers.run( + "postgres:16.2", + detach=True, + environment=env_vars_postgres, + name=cls.DOCKER_CONTAINER_NAME, + ports={5432: ("0.0.0.0", 8081)}, + command="postgres -c log_statement=all", + ) + time.sleep(0.5) + + database_url = cls._create_db_url(env_vars_postgres) + os.environ["DATABASE_URL"] = database_url + retries = 10 + db = None + while retries > 0: + try: + from config import BACKEND_DIR + + db = create_engine(database_url, pool_pre_ping=True) + db = db.connect() + log.info("postgres is ready!") + break + except Exception as e: + log.warning(e) + time.sleep(3) + retries -= 1 + + if db: + # import must be after setting env! + cls.fast_api_client = get_fast_api_client() + db.close() + else: + raise Exception("Could not connect to Postgres") + except Exception as ex: + log.error(ex) + cls.teardown_class() + pytest.fail(f"Could not setup test environment: {ex}") + + def _check_db_connection(self): + from apps.webui.internal.db import Session + + retries = 10 + while retries > 0: + try: + Session.execute(text("SELECT 1")) + Session.commit() + break + except Exception as e: + Session.rollback() + log.warning(e) + time.sleep(3) + retries -= 1 + + def setup_method(self): + super().setup_method() + self._check_db_connection() + + @classmethod + def teardown_class(cls) -> None: + super().teardown_class() + cls.docker_client.containers.get(cls.DOCKER_CONTAINER_NAME).remove(force=True) + + def teardown_method(self): + from apps.webui.internal.db import Session + + # rollback everything not yet committed + Session.commit() + + # truncate all tables + tables = [ + "auth", + "chat", + "chatidtag", + "document", + "memory", + "model", + "prompt", + "tag", + '"user"', + ] + for table in tables: + Session.execute(text(f"TRUNCATE TABLE {table}")) + Session.commit() diff --git a/backend/test/util/mock_user.py b/backend/test/util/mock_user.py new file mode 100644 index 000000000..8d0300d3f --- /dev/null +++ b/backend/test/util/mock_user.py @@ -0,0 +1,45 @@ +from contextlib import contextmanager + +from fastapi import FastAPI + + +@contextmanager +def mock_webui_user(**kwargs): + from apps.webui.main import app + + with mock_user(app, **kwargs): + yield + + +@contextmanager +def mock_user(app: FastAPI, **kwargs): + from utils.utils import ( + get_current_user, + get_verified_user, + get_admin_user, + get_current_user_by_api_key, + ) + from apps.webui.models.users import User + + def create_user(): + user_parameters = { + "id": "1", + "name": "John Doe", + "email": "john.doe@openwebui.com", + "role": "user", + "profile_image_url": "/user.png", + "last_active_at": 1627351200, + "updated_at": 1627351200, + "created_at": 162735120, + **kwargs, + } + return User(**user_parameters) + + app.dependency_overrides = { + get_current_user: create_user, + get_verified_user: create_user, + get_admin_user: create_user, + get_current_user_by_api_key: create_user, + } + yield + app.dependency_overrides = {} diff --git a/backend/utils/misc.py b/backend/utils/misc.py index b4e499df8..5a05f167d 100644 --- a/backend/utils/misc.py +++ b/backend/utils/misc.py @@ -8,14 +8,22 @@ import uuid import time -def get_last_user_message(messages: List[dict]) -> str: +def get_last_user_message_item(messages: List[dict]) -> str: for message in reversed(messages): if message["role"] == "user": - if isinstance(message["content"], list): - for item in message["content"]: - if item["type"] == "text": - return item["text"] - return message["content"] + return message + return None + + +def get_last_user_message(messages: List[dict]) -> str: + message = get_last_user_message_item(messages) + + if message is not None: + if isinstance(message["content"], list): + for item in message["content"]: + if item["type"] == "text": + return item["text"] + return message["content"] return None diff --git a/backend/utils/tools.py b/backend/utils/tools.py index c1c41ed37..3e5d82fd6 100644 --- a/backend/utils/tools.py +++ b/backend/utils/tools.py @@ -59,7 +59,10 @@ def get_tools_specs(tools) -> List[dict]: for param_name, param_annotation in get_type_hints( function ).items() - if param_name != "return" and param_name != "__user__" + if param_name != "return" + and not ( + param_name.startswith("__") and param_name.endswith("__") + ) }, "required": [ name diff --git a/backend/utils/utils.py b/backend/utils/utils.py index 8c3c899bd..fbc539af5 100644 --- a/backend/utils/utils.py +++ b/backend/utils/utils.py @@ -1,5 +1,6 @@ from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials from fastapi import HTTPException, status, Depends, Request +from sqlalchemy.orm import Session from apps.webui.models.users import Users diff --git a/package-lock.json b/package-lock.json index 9c009e356..9eb09d421 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "open-webui", - "version": "0.3.7", + "version": "0.3.8", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "open-webui", - "version": "0.3.7", + "version": "0.3.8", "dependencies": { "@codemirror/lang-javascript": "^6.2.2", "@codemirror/lang-python": "^6.1.6", diff --git a/package.json b/package.json index 0ad2445df..080f8ed5b 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "open-webui", - "version": "0.3.7", + "version": "0.3.8", "private": true, "scripts": { "dev": "npm run pyodide:fetch && vite dev --host", diff --git a/requirements-dev.lock b/requirements-dev.lock index f7660eae3..e56ad08f0 100644 --- a/requirements-dev.lock +++ b/requirements-dev.lock @@ -7,6 +7,7 @@ # all-features: false # with-sources: false # generate-hashes: false +# universal: false -e file:. aiohttp==3.9.5 @@ -31,7 +32,7 @@ asgiref==3.8.1 # via opentelemetry-instrumentation-asgi attrs==23.2.0 # via aiohttp -authlib==1.3.0 +authlib==1.3.1 # via open-webui av==11.0.0 # via faster-whisper @@ -398,7 +399,6 @@ pandas==2.2.2 # via open-webui passlib==1.7.4 # via open-webui - # via passlib pathspec==0.12.1 # via black pcodedmp==1.2.6 @@ -457,7 +457,6 @@ pygments==2.18.0 # via rich pyjwt==2.8.0 # via open-webui - # via pyjwt pymysql==1.1.0 # via open-webui pypandoc==1.13 @@ -559,6 +558,9 @@ scipy==1.13.0 # via sentence-transformers sentence-transformers==2.7.0 # via open-webui +setuptools==69.5.1 + # via ctranslate2 + # via opentelemetry-instrumentation shapely==2.0.4 # via rapidocr-onnxruntime shellingham==1.5.4 @@ -653,7 +655,6 @@ uvicorn==0.22.0 # via chromadb # via fastapi # via open-webui - # via uvicorn uvloop==0.19.0 # via uvicorn validators==0.28.1 @@ -681,6 +682,3 @@ youtube-transcript-api==0.6.2 # via open-webui zipp==3.18.1 # via importlib-metadata -setuptools==69.5.1 - # via ctranslate2 - # via opentelemetry-instrumentation diff --git a/requirements.lock b/requirements.lock index f7660eae3..e56ad08f0 100644 --- a/requirements.lock +++ b/requirements.lock @@ -7,6 +7,7 @@ # all-features: false # with-sources: false # generate-hashes: false +# universal: false -e file:. aiohttp==3.9.5 @@ -31,7 +32,7 @@ asgiref==3.8.1 # via opentelemetry-instrumentation-asgi attrs==23.2.0 # via aiohttp -authlib==1.3.0 +authlib==1.3.1 # via open-webui av==11.0.0 # via faster-whisper @@ -398,7 +399,6 @@ pandas==2.2.2 # via open-webui passlib==1.7.4 # via open-webui - # via passlib pathspec==0.12.1 # via black pcodedmp==1.2.6 @@ -457,7 +457,6 @@ pygments==2.18.0 # via rich pyjwt==2.8.0 # via open-webui - # via pyjwt pymysql==1.1.0 # via open-webui pypandoc==1.13 @@ -559,6 +558,9 @@ scipy==1.13.0 # via sentence-transformers sentence-transformers==2.7.0 # via open-webui +setuptools==69.5.1 + # via ctranslate2 + # via opentelemetry-instrumentation shapely==2.0.4 # via rapidocr-onnxruntime shellingham==1.5.4 @@ -653,7 +655,6 @@ uvicorn==0.22.0 # via chromadb # via fastapi # via open-webui - # via uvicorn uvloop==0.19.0 # via uvicorn validators==0.28.1 @@ -681,6 +682,3 @@ youtube-transcript-api==0.6.2 # via open-webui zipp==3.18.1 # via importlib-metadata -setuptools==69.5.1 - # via ctranslate2 - # via opentelemetry-instrumentation diff --git a/src/app.css b/src/app.css index 2dacf5d72..c3388f1d3 100644 --- a/src/app.css +++ b/src/app.css @@ -1,6 +1,12 @@ @font-face { - font-family: 'Arimo'; - src: url('/assets/fonts/Arimo-Variable.ttf'); + font-family: 'Inter'; + src: url('/assets/fonts/Inter-Variable.ttf'); + font-display: swap; +} + +@font-face { + font-family: 'Archivo'; + src: url('/assets/fonts/Archivo-Variable.ttf'); font-display: swap; } @@ -32,6 +38,10 @@ math { @apply underline; } +.font-primary { + font-family: 'Archivo', sans-serif; +} + iframe { @apply rounded-lg; } @@ -140,3 +150,7 @@ input[type='number'] { .cm-editor.cm-focused { outline: none; } + +.tippy-box[data-theme~='dark'] { + @apply rounded-lg bg-gray-950 text-xs border border-gray-900 shadow-xl; +} diff --git a/src/app.html b/src/app.html index da6af2cc4..6f637bf61 100644 --- a/src/app.html +++ b/src/app.html @@ -23,6 +23,8 @@ // On page load or when changing themes, best to add inline in `head` to avoid FOUC (() => { if (localStorage?.theme && localStorage?.theme.includes('oled')) { + document.documentElement.style.setProperty('--color-gray-800', '#101010'); + document.documentElement.style.setProperty('--color-gray-850', '#050505'); document.documentElement.style.setProperty('--color-gray-900', '#000000'); document.documentElement.style.setProperty('--color-gray-950', '#000000'); document.documentElement.classList.add('dark'); @@ -80,13 +82,13 @@ id="logo" style=" position: absolute; - width: 6rem; + width: auto; height: 6rem; - top: 41%; + top: 44%; left: 50%; margin-left: -3rem; " - src="/logo.svg" + src="/static/splash.png" />
diff --git a/src/lib/apis/rag/index.ts b/src/lib/apis/rag/index.ts index 50f236e06..b32e544ee 100644 --- a/src/lib/apis/rag/index.ts +++ b/src/lib/apis/rag/index.ts @@ -32,6 +32,11 @@ type ChunkConfigForm = { chunk_overlap: number; }; +type ContentExtractConfigForm = { + engine: string; + tika_server_url: string | null; +}; + type YoutubeConfigForm = { language: string[]; translation?: string | null; @@ -40,6 +45,7 @@ type YoutubeConfigForm = { type RAGConfigForm = { pdf_extract_images?: boolean; chunk?: ChunkConfigForm; + content_extraction?: ContentExtractConfigForm; web_loader_ssl_verification?: boolean; youtube?: YoutubeConfigForm; }; diff --git a/src/lib/components/ChangelogModal.svelte b/src/lib/components/ChangelogModal.svelte index d263c88e9..48156f924 100644 --- a/src/lib/components/ChangelogModal.svelte +++ b/src/lib/components/ChangelogModal.svelte @@ -24,7 +24,7 @@
-
+
{$i18n.t('What’s New in')} {$WEBUI_NAME} @@ -63,7 +63,7 @@ {#if changelog} {#each Object.keys(changelog) as version}
-
+
v{version} - {changelog[version].date}
@@ -72,7 +72,7 @@ {#each Object.keys(changelog[version]).filter((section) => section !== 'date') as section}
- import { getContext, tick } from 'svelte'; + import { getContext, tick, onMount } from 'svelte'; import { toast } from 'svelte-sonner'; import Database from './Settings/Database.svelte'; @@ -21,17 +21,31 @@ const i18n = getContext('i18n'); let selectedTab = 'general'; + + onMount(() => { + const containerElement = document.getElementById('admin-settings-tabs-container'); + + if (containerElement) { + containerElement.addEventListener('wheel', function (event) { + if (event.deltaY !== 0) { + // Adjust horizontal scroll position based on vertical scroll + containerElement.scrollLeft += event.deltaY; + } + }); + } + });