Merge branch 'dev' of https://github.com/open-webui/open-webui into dev
This commit is contained in:
Gabriel Ecegi
@@ -0,0 +1,703 @@
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import uuid
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
from pydub import AudioSegment
|
||||
from pydub.silence import split_on_silence
|
||||
|
||||
import aiohttp
|
||||
import aiofiles
|
||||
import requests
|
||||
|
||||
from fastapi import (
|
||||
Depends,
|
||||
FastAPI,
|
||||
File,
|
||||
HTTPException,
|
||||
Request,
|
||||
UploadFile,
|
||||
status,
|
||||
APIRouter,
|
||||
)
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import FileResponse
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
from open_webui.utils.auth import get_admin_user, get_verified_user
|
||||
from open_webui.config import (
|
||||
WHISPER_MODEL_AUTO_UPDATE,
|
||||
WHISPER_MODEL_DIR,
|
||||
CACHE_DIR,
|
||||
)
|
||||
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
from open_webui.env import (
|
||||
ENV,
|
||||
SRC_LOG_LEVELS,
|
||||
DEVICE_TYPE,
|
||||
ENABLE_FORWARD_USER_INFO_HEADERS,
|
||||
)
|
||||
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
# Constants
|
||||
MAX_FILE_SIZE_MB = 25
|
||||
MAX_FILE_SIZE = MAX_FILE_SIZE_MB * 1024 * 1024 # Convert MB to bytes
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["AUDIO"])
|
||||
|
||||
SPEECH_CACHE_DIR = Path(CACHE_DIR).joinpath("./audio/speech/")
|
||||
SPEECH_CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
|
||||
##########################################
|
||||
#
|
||||
# Utility functions
|
||||
#
|
||||
##########################################
|
||||
|
||||
from pydub import AudioSegment
|
||||
from pydub.utils import mediainfo
|
||||
|
||||
|
||||
def is_mp4_audio(file_path):
|
||||
"""Check if the given file is an MP4 audio file."""
|
||||
if not os.path.isfile(file_path):
|
||||
print(f"File not found: {file_path}")
|
||||
return False
|
||||
|
||||
info = mediainfo(file_path)
|
||||
if (
|
||||
info.get("codec_name") == "aac"
|
||||
and info.get("codec_type") == "audio"
|
||||
and info.get("codec_tag_string") == "mp4a"
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def convert_mp4_to_wav(file_path, output_path):
|
||||
"""Convert MP4 audio file to WAV format."""
|
||||
audio = AudioSegment.from_file(file_path, format="mp4")
|
||||
audio.export(output_path, format="wav")
|
||||
print(f"Converted {file_path} to {output_path}")
|
||||
|
||||
|
||||
def set_faster_whisper_model(model: str, auto_update: bool = False):
|
||||
whisper_model = None
|
||||
if model:
|
||||
from faster_whisper import WhisperModel
|
||||
|
||||
faster_whisper_kwargs = {
|
||||
"model_size_or_path": model,
|
||||
"device": DEVICE_TYPE if DEVICE_TYPE and DEVICE_TYPE == "cuda" else "cpu",
|
||||
"compute_type": "int8",
|
||||
"download_root": WHISPER_MODEL_DIR,
|
||||
"local_files_only": not auto_update,
|
||||
}
|
||||
|
||||
try:
|
||||
whisper_model = WhisperModel(**faster_whisper_kwargs)
|
||||
except Exception:
|
||||
log.warning(
|
||||
"WhisperModel initialization failed, attempting download with local_files_only=False"
|
||||
)
|
||||
faster_whisper_kwargs["local_files_only"] = False
|
||||
whisper_model = WhisperModel(**faster_whisper_kwargs)
|
||||
return whisper_model
|
||||
|
||||
|
||||
##########################################
|
||||
#
|
||||
# Audio API
|
||||
#
|
||||
##########################################
|
||||
|
||||
|
||||
class TTSConfigForm(BaseModel):
|
||||
OPENAI_API_BASE_URL: str
|
||||
OPENAI_API_KEY: str
|
||||
API_KEY: str
|
||||
ENGINE: str
|
||||
MODEL: str
|
||||
VOICE: str
|
||||
SPLIT_ON: str
|
||||
AZURE_SPEECH_REGION: str
|
||||
AZURE_SPEECH_OUTPUT_FORMAT: str
|
||||
|
||||
|
||||
class STTConfigForm(BaseModel):
|
||||
OPENAI_API_BASE_URL: str
|
||||
OPENAI_API_KEY: str
|
||||
ENGINE: str
|
||||
MODEL: str
|
||||
WHISPER_MODEL: str
|
||||
|
||||
|
||||
class AudioConfigUpdateForm(BaseModel):
|
||||
tts: TTSConfigForm
|
||||
stt: STTConfigForm
|
||||
|
||||
|
||||
@router.get("/config")
|
||||
async def get_audio_config(request: Request, user=Depends(get_admin_user)):
|
||||
return {
|
||||
"tts": {
|
||||
"OPENAI_API_BASE_URL": request.app.state.config.TTS_OPENAI_API_BASE_URL,
|
||||
"OPENAI_API_KEY": request.app.state.config.TTS_OPENAI_API_KEY,
|
||||
"API_KEY": request.app.state.config.TTS_API_KEY,
|
||||
"ENGINE": request.app.state.config.TTS_ENGINE,
|
||||
"MODEL": request.app.state.config.TTS_MODEL,
|
||||
"VOICE": request.app.state.config.TTS_VOICE,
|
||||
"SPLIT_ON": request.app.state.config.TTS_SPLIT_ON,
|
||||
"AZURE_SPEECH_REGION": request.app.state.config.TTS_AZURE_SPEECH_REGION,
|
||||
"AZURE_SPEECH_OUTPUT_FORMAT": request.app.state.config.TTS_AZURE_SPEECH_OUTPUT_FORMAT,
|
||||
},
|
||||
"stt": {
|
||||
"OPENAI_API_BASE_URL": request.app.state.config.STT_OPENAI_API_BASE_URL,
|
||||
"OPENAI_API_KEY": request.app.state.config.STT_OPENAI_API_KEY,
|
||||
"ENGINE": request.app.state.config.STT_ENGINE,
|
||||
"MODEL": request.app.state.config.STT_MODEL,
|
||||
"WHISPER_MODEL": request.app.state.config.WHISPER_MODEL,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@router.post("/config/update")
|
||||
async def update_audio_config(
|
||||
request: Request, form_data: AudioConfigUpdateForm, user=Depends(get_admin_user)
|
||||
):
|
||||
request.app.state.config.TTS_OPENAI_API_BASE_URL = form_data.tts.OPENAI_API_BASE_URL
|
||||
request.app.state.config.TTS_OPENAI_API_KEY = form_data.tts.OPENAI_API_KEY
|
||||
request.app.state.config.TTS_API_KEY = form_data.tts.API_KEY
|
||||
request.app.state.config.TTS_ENGINE = form_data.tts.ENGINE
|
||||
request.app.state.config.TTS_MODEL = form_data.tts.MODEL
|
||||
request.app.state.config.TTS_VOICE = form_data.tts.VOICE
|
||||
request.app.state.config.TTS_SPLIT_ON = form_data.tts.SPLIT_ON
|
||||
request.app.state.config.TTS_AZURE_SPEECH_REGION = form_data.tts.AZURE_SPEECH_REGION
|
||||
request.app.state.config.TTS_AZURE_SPEECH_OUTPUT_FORMAT = (
|
||||
form_data.tts.AZURE_SPEECH_OUTPUT_FORMAT
|
||||
)
|
||||
|
||||
request.app.state.config.STT_OPENAI_API_BASE_URL = form_data.stt.OPENAI_API_BASE_URL
|
||||
request.app.state.config.STT_OPENAI_API_KEY = form_data.stt.OPENAI_API_KEY
|
||||
request.app.state.config.STT_ENGINE = form_data.stt.ENGINE
|
||||
request.app.state.config.STT_MODEL = form_data.stt.MODEL
|
||||
request.app.state.config.WHISPER_MODEL = form_data.stt.WHISPER_MODEL
|
||||
|
||||
if request.app.state.config.STT_ENGINE == "":
|
||||
request.app.state.faster_whisper_model = set_faster_whisper_model(
|
||||
form_data.stt.WHISPER_MODEL, WHISPER_MODEL_AUTO_UPDATE
|
||||
)
|
||||
|
||||
return {
|
||||
"tts": {
|
||||
"OPENAI_API_BASE_URL": request.app.state.config.TTS_OPENAI_API_BASE_URL,
|
||||
"OPENAI_API_KEY": request.app.state.config.TTS_OPENAI_API_KEY,
|
||||
"API_KEY": request.app.state.config.TTS_API_KEY,
|
||||
"ENGINE": request.app.state.config.TTS_ENGINE,
|
||||
"MODEL": request.app.state.config.TTS_MODEL,
|
||||
"VOICE": request.app.state.config.TTS_VOICE,
|
||||
"SPLIT_ON": request.app.state.config.TTS_SPLIT_ON,
|
||||
"AZURE_SPEECH_REGION": request.app.state.config.TTS_AZURE_SPEECH_REGION,
|
||||
"AZURE_SPEECH_OUTPUT_FORMAT": request.app.state.config.TTS_AZURE_SPEECH_OUTPUT_FORMAT,
|
||||
},
|
||||
"stt": {
|
||||
"OPENAI_API_BASE_URL": request.app.state.config.STT_OPENAI_API_BASE_URL,
|
||||
"OPENAI_API_KEY": request.app.state.config.STT_OPENAI_API_KEY,
|
||||
"ENGINE": request.app.state.config.STT_ENGINE,
|
||||
"MODEL": request.app.state.config.STT_MODEL,
|
||||
"WHISPER_MODEL": request.app.state.config.WHISPER_MODEL,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def load_speech_pipeline():
|
||||
from transformers import pipeline
|
||||
from datasets import load_dataset
|
||||
|
||||
if request.app.state.speech_synthesiser is None:
|
||||
request.app.state.speech_synthesiser = pipeline(
|
||||
"text-to-speech", "microsoft/speecht5_tts"
|
||||
)
|
||||
|
||||
if request.app.state.speech_speaker_embeddings_dataset is None:
|
||||
request.app.state.speech_speaker_embeddings_dataset = load_dataset(
|
||||
"Matthijs/cmu-arctic-xvectors", split="validation"
|
||||
)
|
||||
|
||||
|
||||
@router.post("/speech")
|
||||
async def speech(request: Request, user=Depends(get_verified_user)):
|
||||
body = await request.body()
|
||||
name = hashlib.sha256(body).hexdigest()
|
||||
|
||||
file_path = SPEECH_CACHE_DIR.joinpath(f"{name}.mp3")
|
||||
file_body_path = SPEECH_CACHE_DIR.joinpath(f"{name}.json")
|
||||
|
||||
# Check if the file already exists in the cache
|
||||
if file_path.is_file():
|
||||
return FileResponse(file_path)
|
||||
|
||||
payload = None
|
||||
try:
|
||||
payload = json.loads(body.decode("utf-8"))
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
raise HTTPException(status_code=400, detail="Invalid JSON payload")
|
||||
|
||||
if request.app.state.config.TTS_ENGINE == "openai":
|
||||
payload["model"] = request.app.state.config.TTS_MODEL
|
||||
|
||||
try:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(
|
||||
url=f"{request.app.state.config.TTS_OPENAI_API_BASE_URL}/audio/speech",
|
||||
data=payload,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {request.app.state.config.TTS_OPENAI_API_KEY}",
|
||||
**(
|
||||
{
|
||||
"X-OpenWebUI-User-Name": user.name,
|
||||
"X-OpenWebUI-User-Id": user.id,
|
||||
"X-OpenWebUI-User-Email": user.email,
|
||||
"X-OpenWebUI-User-Role": user.role,
|
||||
}
|
||||
if ENABLE_FORWARD_USER_INFO_HEADERS
|
||||
else {}
|
||||
),
|
||||
},
|
||||
) as r:
|
||||
r.raise_for_status()
|
||||
|
||||
async with aiofiles.open(file_path, "wb") as f:
|
||||
await f.write(await r.read())
|
||||
|
||||
async with aiofiles.open(file_body_path, "w") as f:
|
||||
await f.write(json.dumps(json.loads(body.decode("utf-8"))))
|
||||
|
||||
return FileResponse(file_path)
|
||||
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
detail = None
|
||||
|
||||
try:
|
||||
if r.status != 200:
|
||||
res = await r.json()
|
||||
if "error" in res:
|
||||
detail = f"External: {res['error'].get('message', '')}"
|
||||
except Exception:
|
||||
detail = f"External: {e}"
|
||||
|
||||
raise HTTPException(
|
||||
status_code=getattr(r, "status", 500),
|
||||
detail=detail if detail else "Open WebUI: Server Connection Error",
|
||||
)
|
||||
|
||||
elif request.app.state.config.TTS_ENGINE == "elevenlabs":
|
||||
voice_id = payload.get("voice", "")
|
||||
|
||||
if voice_id not in get_available_voices():
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Invalid voice id",
|
||||
)
|
||||
|
||||
try:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(
|
||||
f"https://api.elevenlabs.io/v1/text-to-speech/{voice_id}",
|
||||
json={
|
||||
"text": payload["input"],
|
||||
"model_id": request.app.state.config.TTS_MODEL,
|
||||
"voice_settings": {"stability": 0.5, "similarity_boost": 0.5},
|
||||
},
|
||||
headers={
|
||||
"Accept": "audio/mpeg",
|
||||
"Content-Type": "application/json",
|
||||
"xi-api-key": request.app.state.config.TTS_API_KEY,
|
||||
},
|
||||
) as r:
|
||||
r.raise_for_status()
|
||||
|
||||
async with aiofiles.open(file_path, "wb") as f:
|
||||
await f.write(await r.read())
|
||||
|
||||
async with aiofiles.open(file_body_path, "w") as f:
|
||||
await f.write(json.dumps(json.loads(body.decode("utf-8"))))
|
||||
|
||||
return FileResponse(file_path)
|
||||
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
detail = None
|
||||
|
||||
try:
|
||||
if r.status != 200:
|
||||
res = await r.json()
|
||||
if "error" in res:
|
||||
detail = f"External: {res['error'].get('message', '')}"
|
||||
except Exception:
|
||||
detail = f"External: {e}"
|
||||
|
||||
raise HTTPException(
|
||||
status_code=getattr(r, "status", 500),
|
||||
detail=detail if detail else "Open WebUI: Server Connection Error",
|
||||
)
|
||||
|
||||
elif request.app.state.config.TTS_ENGINE == "azure":
|
||||
try:
|
||||
payload = json.loads(body.decode("utf-8"))
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
raise HTTPException(status_code=400, detail="Invalid JSON payload")
|
||||
|
||||
region = request.app.state.config.TTS_AZURE_SPEECH_REGION
|
||||
language = request.app.state.config.TTS_VOICE
|
||||
locale = "-".join(request.app.state.config.TTS_VOICE.split("-")[:1])
|
||||
output_format = request.app.state.config.TTS_AZURE_SPEECH_OUTPUT_FORMAT
|
||||
|
||||
try:
|
||||
data = f"""<speak version="1.0" xmlns="http://www.w3.org/2001/10/synthesis" xml:lang="{locale}">
|
||||
<voice name="{language}">{payload["input"]}</voice>
|
||||
</speak>"""
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(
|
||||
f"https://{region}.tts.speech.microsoft.com/cognitiveservices/v1",
|
||||
headers={
|
||||
"Ocp-Apim-Subscription-Key": request.app.state.config.TTS_API_KEY,
|
||||
"Content-Type": "application/ssml+xml",
|
||||
"X-Microsoft-OutputFormat": output_format,
|
||||
},
|
||||
data=data,
|
||||
) as r:
|
||||
r.raise_for_status()
|
||||
|
||||
async with aiofiles.open(file_path, "wb") as f:
|
||||
await f.write(await r.read())
|
||||
|
||||
return FileResponse(file_path)
|
||||
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
detail = None
|
||||
|
||||
try:
|
||||
if r.status != 200:
|
||||
res = await r.json()
|
||||
if "error" in res:
|
||||
detail = f"External: {res['error'].get('message', '')}"
|
||||
except Exception:
|
||||
detail = f"External: {e}"
|
||||
|
||||
raise HTTPException(
|
||||
status_code=getattr(r, "status", 500),
|
||||
detail=detail if detail else "Open WebUI: Server Connection Error",
|
||||
)
|
||||
|
||||
elif request.app.state.config.TTS_ENGINE == "transformers":
|
||||
payload = None
|
||||
try:
|
||||
payload = json.loads(body.decode("utf-8"))
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
raise HTTPException(status_code=400, detail="Invalid JSON payload")
|
||||
|
||||
import torch
|
||||
import soundfile as sf
|
||||
|
||||
load_speech_pipeline()
|
||||
|
||||
embeddings_dataset = request.app.state.speech_speaker_embeddings_dataset
|
||||
|
||||
speaker_index = 6799
|
||||
try:
|
||||
speaker_index = embeddings_dataset["filename"].index(
|
||||
request.app.state.config.TTS_MODEL
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
speaker_embedding = torch.tensor(
|
||||
embeddings_dataset[speaker_index]["xvector"]
|
||||
).unsqueeze(0)
|
||||
|
||||
speech = request.app.state.speech_synthesiser(
|
||||
payload["input"],
|
||||
forward_params={"speaker_embeddings": speaker_embedding},
|
||||
)
|
||||
|
||||
sf.write(file_path, speech["audio"], samplerate=speech["sampling_rate"])
|
||||
with open(file_body_path, "w") as f:
|
||||
json.dump(json.loads(body.decode("utf-8")), f)
|
||||
|
||||
return FileResponse(file_path)
|
||||
|
||||
|
||||
def transcribe(request: Request, file_path):
|
||||
print("transcribe", file_path)
|
||||
filename = os.path.basename(file_path)
|
||||
file_dir = os.path.dirname(file_path)
|
||||
id = filename.split(".")[0]
|
||||
|
||||
if request.app.state.config.STT_ENGINE == "":
|
||||
if request.app.state.faster_whisper_model is None:
|
||||
request.app.state.faster_whisper_model = set_faster_whisper_model(
|
||||
request.app.state.config.WHISPER_MODEL
|
||||
)
|
||||
|
||||
model = request.app.state.faster_whisper_model
|
||||
segments, info = model.transcribe(file_path, beam_size=5)
|
||||
log.info(
|
||||
"Detected language '%s' with probability %f"
|
||||
% (info.language, info.language_probability)
|
||||
)
|
||||
|
||||
transcript = "".join([segment.text for segment in list(segments)])
|
||||
data = {"text": transcript.strip()}
|
||||
|
||||
# save the transcript to a json file
|
||||
transcript_file = f"{file_dir}/{id}.json"
|
||||
with open(transcript_file, "w") as f:
|
||||
json.dump(data, f)
|
||||
|
||||
log.debug(data)
|
||||
return data
|
||||
elif request.app.state.config.STT_ENGINE == "openai":
|
||||
if is_mp4_audio(file_path):
|
||||
os.rename(file_path, file_path.replace(".wav", ".mp4"))
|
||||
# Convert MP4 audio file to WAV format
|
||||
convert_mp4_to_wav(file_path.replace(".wav", ".mp4"), file_path)
|
||||
|
||||
r = None
|
||||
try:
|
||||
r = requests.post(
|
||||
url=f"{request.app.state.config.STT_OPENAI_API_BASE_URL}/audio/transcriptions",
|
||||
headers={
|
||||
"Authorization": f"Bearer {request.app.state.config.STT_OPENAI_API_KEY}"
|
||||
},
|
||||
files={"file": (filename, open(file_path, "rb"))},
|
||||
data={"model": request.app.state.config.STT_MODEL},
|
||||
)
|
||||
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
|
||||
# save the transcript to a json file
|
||||
transcript_file = f"{file_dir}/{id}.json"
|
||||
with open(transcript_file, "w") as f:
|
||||
json.dump(data, f)
|
||||
|
||||
return data
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
|
||||
detail = None
|
||||
if r is not None:
|
||||
try:
|
||||
res = r.json()
|
||||
if "error" in res:
|
||||
detail = f"External: {res['error'].get('message', '')}"
|
||||
except Exception:
|
||||
detail = f"External: {e}"
|
||||
|
||||
raise Exception(detail if detail else "Open WebUI: Server Connection Error")
|
||||
|
||||
|
||||
def compress_audio(file_path):
|
||||
if os.path.getsize(file_path) > MAX_FILE_SIZE:
|
||||
file_dir = os.path.dirname(file_path)
|
||||
audio = AudioSegment.from_file(file_path)
|
||||
audio = audio.set_frame_rate(16000).set_channels(1) # Compress audio
|
||||
compressed_path = f"{file_dir}/{id}_compressed.opus"
|
||||
audio.export(compressed_path, format="opus", bitrate="32k")
|
||||
log.debug(f"Compressed audio to {compressed_path}")
|
||||
|
||||
if (
|
||||
os.path.getsize(compressed_path) > MAX_FILE_SIZE
|
||||
): # Still larger than MAX_FILE_SIZE after compression
|
||||
raise Exception(ERROR_MESSAGES.FILE_TOO_LARGE(size=f"{MAX_FILE_SIZE_MB}MB"))
|
||||
return compressed_path
|
||||
else:
|
||||
return file_path
|
||||
|
||||
|
||||
@router.post("/transcriptions")
|
||||
def transcription(
|
||||
request: Request,
|
||||
file: UploadFile = File(...),
|
||||
user=Depends(get_verified_user),
|
||||
):
|
||||
log.info(f"file.content_type: {file.content_type}")
|
||||
|
||||
if file.content_type not in ["audio/mpeg", "audio/wav", "audio/ogg", "audio/x-m4a"]:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.FILE_NOT_SUPPORTED,
|
||||
)
|
||||
|
||||
try:
|
||||
ext = file.filename.split(".")[-1]
|
||||
id = uuid.uuid4()
|
||||
|
||||
filename = f"{id}.{ext}"
|
||||
contents = file.file.read()
|
||||
|
||||
file_dir = f"{CACHE_DIR}/audio/transcriptions"
|
||||
os.makedirs(file_dir, exist_ok=True)
|
||||
file_path = f"{file_dir}/{filename}"
|
||||
|
||||
with open(file_path, "wb") as f:
|
||||
f.write(contents)
|
||||
|
||||
try:
|
||||
try:
|
||||
file_path = compress_audio(file_path)
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(e),
|
||||
)
|
||||
|
||||
data = transcribe(request, file_path)
|
||||
file_path = file_path.split("/")[-1]
|
||||
return {**data, "filename": file_path}
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(e),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(e),
|
||||
)
|
||||
|
||||
|
||||
def get_available_models(request: Request) -> list[dict]:
|
||||
available_models = []
|
||||
if request.app.state.config.TTS_ENGINE == "openai":
|
||||
available_models = [{"id": "tts-1"}, {"id": "tts-1-hd"}]
|
||||
elif request.app.state.config.TTS_ENGINE == "elevenlabs":
|
||||
try:
|
||||
response = requests.get(
|
||||
"https://api.elevenlabs.io/v1/models",
|
||||
headers={
|
||||
"xi-api-key": request.app.state.config.TTS_API_KEY,
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
timeout=5,
|
||||
)
|
||||
response.raise_for_status()
|
||||
models = response.json()
|
||||
|
||||
available_models = [
|
||||
{"name": model["name"], "id": model["model_id"]} for model in models
|
||||
]
|
||||
except requests.RequestException as e:
|
||||
log.error(f"Error fetching voices: {str(e)}")
|
||||
return available_models
|
||||
|
||||
|
||||
@router.get("/models")
|
||||
async def get_models(request: Request, user=Depends(get_verified_user)):
|
||||
return {"models": get_available_models(request)}
|
||||
|
||||
|
||||
def get_available_voices(request) -> dict:
|
||||
"""Returns {voice_id: voice_name} dict"""
|
||||
available_voices = {}
|
||||
if request.app.state.config.TTS_ENGINE == "openai":
|
||||
available_voices = {
|
||||
"alloy": "alloy",
|
||||
"echo": "echo",
|
||||
"fable": "fable",
|
||||
"onyx": "onyx",
|
||||
"nova": "nova",
|
||||
"shimmer": "shimmer",
|
||||
}
|
||||
elif request.app.state.config.TTS_ENGINE == "elevenlabs":
|
||||
try:
|
||||
available_voices = get_elevenlabs_voices(
|
||||
api_key=request.app.state.config.TTS_API_KEY
|
||||
)
|
||||
except Exception:
|
||||
# Avoided @lru_cache with exception
|
||||
pass
|
||||
elif request.app.state.config.TTS_ENGINE == "azure":
|
||||
try:
|
||||
region = request.app.state.config.TTS_AZURE_SPEECH_REGION
|
||||
url = f"https://{region}.tts.speech.microsoft.com/cognitiveservices/voices/list"
|
||||
headers = {
|
||||
"Ocp-Apim-Subscription-Key": request.app.state.config.TTS_API_KEY
|
||||
}
|
||||
|
||||
response = requests.get(url, headers=headers)
|
||||
response.raise_for_status()
|
||||
voices = response.json()
|
||||
|
||||
for voice in voices:
|
||||
available_voices[voice["ShortName"]] = (
|
||||
f"{voice['DisplayName']} ({voice['ShortName']})"
|
||||
)
|
||||
except requests.RequestException as e:
|
||||
log.error(f"Error fetching voices: {str(e)}")
|
||||
|
||||
return available_voices
|
||||
|
||||
|
||||
@lru_cache
|
||||
def get_elevenlabs_voices(api_key: str) -> dict:
|
||||
"""
|
||||
Note, set the following in your .env file to use Elevenlabs:
|
||||
AUDIO_TTS_ENGINE=elevenlabs
|
||||
AUDIO_TTS_API_KEY=sk_... # Your Elevenlabs API key
|
||||
AUDIO_TTS_VOICE=EXAVITQu4vr4xnSDxMaL # From https://api.elevenlabs.io/v1/voices
|
||||
AUDIO_TTS_MODEL=eleven_multilingual_v2
|
||||
"""
|
||||
|
||||
try:
|
||||
# TODO: Add retries
|
||||
response = requests.get(
|
||||
"https://api.elevenlabs.io/v1/voices",
|
||||
headers={
|
||||
"xi-api-key": api_key,
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
voices_data = response.json()
|
||||
|
||||
voices = {}
|
||||
for voice in voices_data.get("voices", []):
|
||||
voices[voice["voice_id"]] = voice["name"]
|
||||
except requests.RequestException as e:
|
||||
# Avoid @lru_cache with exception
|
||||
log.error(f"Error fetching voices: {str(e)}")
|
||||
raise RuntimeError(f"Error fetching voices: {str(e)}")
|
||||
|
||||
return voices
|
||||
|
||||
|
||||
@router.get("/voices")
|
||||
async def get_voices(request: Request, user=Depends(get_verified_user)):
|
||||
return {
|
||||
"voices": [
|
||||
{"id": k, "name": v} for k, v in get_available_voices(request).items()
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,809 @@
|
||||
import re
|
||||
import uuid
|
||||
import time
|
||||
import datetime
|
||||
import logging
|
||||
from aiohttp import ClientSession
|
||||
|
||||
from open_webui.models.auths import (
|
||||
AddUserForm,
|
||||
ApiKey,
|
||||
Auths,
|
||||
Token,
|
||||
LdapForm,
|
||||
SigninForm,
|
||||
SigninResponse,
|
||||
SignupForm,
|
||||
UpdatePasswordForm,
|
||||
UpdateProfileForm,
|
||||
UserResponse,
|
||||
)
|
||||
from open_webui.models.users import Users
|
||||
|
||||
from open_webui.constants import ERROR_MESSAGES, WEBHOOK_MESSAGES
|
||||
from open_webui.env import (
|
||||
WEBUI_AUTH,
|
||||
WEBUI_AUTH_TRUSTED_EMAIL_HEADER,
|
||||
WEBUI_AUTH_TRUSTED_NAME_HEADER,
|
||||
WEBUI_SESSION_COOKIE_SAME_SITE,
|
||||
WEBUI_SESSION_COOKIE_SECURE,
|
||||
SRC_LOG_LEVELS,
|
||||
)
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, status
|
||||
from fastapi.responses import RedirectResponse, Response
|
||||
from open_webui.config import (
|
||||
OPENID_PROVIDER_URL,
|
||||
ENABLE_OAUTH_SIGNUP,
|
||||
)
|
||||
from pydantic import BaseModel
|
||||
from open_webui.utils.misc import parse_duration, validate_email_format
|
||||
from open_webui.utils.auth import (
|
||||
create_api_key,
|
||||
create_token,
|
||||
get_admin_user,
|
||||
get_verified_user,
|
||||
get_current_user,
|
||||
get_password_hash,
|
||||
)
|
||||
from open_webui.utils.webhook import post_webhook
|
||||
from open_webui.utils.access_control import get_permissions
|
||||
|
||||
from typing import Optional, List
|
||||
|
||||
from ssl import CERT_REQUIRED, PROTOCOL_TLS
|
||||
from ldap3 import Server, Connection, ALL, Tls
|
||||
from ldap3.utils.conv import escape_filter_chars
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["MAIN"])
|
||||
|
||||
############################
|
||||
# GetSessionUser
|
||||
############################
|
||||
|
||||
|
||||
class SessionUserResponse(Token, UserResponse):
|
||||
expires_at: Optional[int] = None
|
||||
permissions: Optional[dict] = None
|
||||
|
||||
|
||||
@router.get("/", response_model=SessionUserResponse)
|
||||
async def get_session_user(
|
||||
request: Request, response: Response, user=Depends(get_current_user)
|
||||
):
|
||||
expires_delta = parse_duration(request.app.state.config.JWT_EXPIRES_IN)
|
||||
expires_at = None
|
||||
if expires_delta:
|
||||
expires_at = int(time.time()) + int(expires_delta.total_seconds())
|
||||
|
||||
token = create_token(
|
||||
data={"id": user.id},
|
||||
expires_delta=expires_delta,
|
||||
)
|
||||
|
||||
datetime_expires_at = (
|
||||
datetime.datetime.fromtimestamp(expires_at, datetime.timezone.utc)
|
||||
if expires_at
|
||||
else None
|
||||
)
|
||||
|
||||
# Set the cookie token
|
||||
response.set_cookie(
|
||||
key="token",
|
||||
value=token,
|
||||
expires=datetime_expires_at,
|
||||
httponly=True, # Ensures the cookie is not accessible via JavaScript
|
||||
samesite=WEBUI_SESSION_COOKIE_SAME_SITE,
|
||||
secure=WEBUI_SESSION_COOKIE_SECURE,
|
||||
)
|
||||
|
||||
user_permissions = get_permissions(
|
||||
user.id, request.app.state.config.USER_PERMISSIONS
|
||||
)
|
||||
|
||||
return {
|
||||
"token": token,
|
||||
"token_type": "Bearer",
|
||||
"expires_at": expires_at,
|
||||
"id": user.id,
|
||||
"email": user.email,
|
||||
"name": user.name,
|
||||
"role": user.role,
|
||||
"profile_image_url": user.profile_image_url,
|
||||
"permissions": user_permissions,
|
||||
}
|
||||
|
||||
|
||||
############################
|
||||
# Update Profile
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/update/profile", response_model=UserResponse)
|
||||
async def update_profile(
|
||||
form_data: UpdateProfileForm, session_user=Depends(get_verified_user)
|
||||
):
|
||||
if session_user:
|
||||
user = Users.update_user_by_id(
|
||||
session_user.id,
|
||||
{"profile_image_url": form_data.profile_image_url, "name": form_data.name},
|
||||
)
|
||||
if user:
|
||||
return user
|
||||
else:
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.DEFAULT())
|
||||
else:
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)
|
||||
|
||||
|
||||
############################
|
||||
# Update Password
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/update/password", response_model=bool)
|
||||
async def update_password(
|
||||
form_data: UpdatePasswordForm, session_user=Depends(get_current_user)
|
||||
):
|
||||
if WEBUI_AUTH_TRUSTED_EMAIL_HEADER:
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.ACTION_PROHIBITED)
|
||||
if session_user:
|
||||
user = Auths.authenticate_user(session_user.email, form_data.password)
|
||||
|
||||
if user:
|
||||
hashed = get_password_hash(form_data.new_password)
|
||||
return Auths.update_user_password_by_id(user.id, hashed)
|
||||
else:
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_PASSWORD)
|
||||
else:
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)
|
||||
|
||||
|
||||
############################
|
||||
# LDAP Authentication
|
||||
############################
|
||||
@router.post("/ldap", response_model=SigninResponse)
|
||||
async def ldap_auth(request: Request, response: Response, form_data: LdapForm):
|
||||
ENABLE_LDAP = request.app.state.config.ENABLE_LDAP
|
||||
LDAP_SERVER_LABEL = request.app.state.config.LDAP_SERVER_LABEL
|
||||
LDAP_SERVER_HOST = request.app.state.config.LDAP_SERVER_HOST
|
||||
LDAP_SERVER_PORT = request.app.state.config.LDAP_SERVER_PORT
|
||||
LDAP_ATTRIBUTE_FOR_USERNAME = request.app.state.config.LDAP_ATTRIBUTE_FOR_USERNAME
|
||||
LDAP_SEARCH_BASE = request.app.state.config.LDAP_SEARCH_BASE
|
||||
LDAP_SEARCH_FILTERS = request.app.state.config.LDAP_SEARCH_FILTERS
|
||||
LDAP_APP_DN = request.app.state.config.LDAP_APP_DN
|
||||
LDAP_APP_PASSWORD = request.app.state.config.LDAP_APP_PASSWORD
|
||||
LDAP_USE_TLS = request.app.state.config.LDAP_USE_TLS
|
||||
LDAP_CA_CERT_FILE = request.app.state.config.LDAP_CA_CERT_FILE
|
||||
LDAP_CIPHERS = (
|
||||
request.app.state.config.LDAP_CIPHERS
|
||||
if request.app.state.config.LDAP_CIPHERS
|
||||
else "ALL"
|
||||
)
|
||||
|
||||
if not ENABLE_LDAP:
|
||||
raise HTTPException(400, detail="LDAP authentication is not enabled")
|
||||
|
||||
try:
|
||||
tls = Tls(
|
||||
validate=CERT_REQUIRED,
|
||||
version=PROTOCOL_TLS,
|
||||
ca_certs_file=LDAP_CA_CERT_FILE,
|
||||
ciphers=LDAP_CIPHERS,
|
||||
)
|
||||
except Exception as e:
|
||||
log.error(f"An error occurred on TLS: {str(e)}")
|
||||
raise HTTPException(400, detail=str(e))
|
||||
|
||||
try:
|
||||
server = Server(
|
||||
host=LDAP_SERVER_HOST,
|
||||
port=LDAP_SERVER_PORT,
|
||||
get_info=ALL,
|
||||
use_ssl=LDAP_USE_TLS,
|
||||
tls=tls,
|
||||
)
|
||||
connection_app = Connection(
|
||||
server,
|
||||
LDAP_APP_DN,
|
||||
LDAP_APP_PASSWORD,
|
||||
auto_bind="NONE",
|
||||
authentication="SIMPLE",
|
||||
)
|
||||
if not connection_app.bind():
|
||||
raise HTTPException(400, detail="Application account bind failed")
|
||||
|
||||
search_success = connection_app.search(
|
||||
search_base=LDAP_SEARCH_BASE,
|
||||
search_filter=f"(&({LDAP_ATTRIBUTE_FOR_USERNAME}={escape_filter_chars(form_data.user.lower())}){LDAP_SEARCH_FILTERS})",
|
||||
attributes=[f"{LDAP_ATTRIBUTE_FOR_USERNAME}", "mail", "cn"],
|
||||
)
|
||||
|
||||
if not search_success:
|
||||
raise HTTPException(400, detail="User not found in the LDAP server")
|
||||
|
||||
entry = connection_app.entries[0]
|
||||
username = str(entry[f"{LDAP_ATTRIBUTE_FOR_USERNAME}"]).lower()
|
||||
mail = str(entry["mail"])
|
||||
cn = str(entry["cn"])
|
||||
user_dn = entry.entry_dn
|
||||
|
||||
if username == form_data.user.lower():
|
||||
connection_user = Connection(
|
||||
server,
|
||||
user_dn,
|
||||
form_data.password,
|
||||
auto_bind="NONE",
|
||||
authentication="SIMPLE",
|
||||
)
|
||||
if not connection_user.bind():
|
||||
raise HTTPException(400, f"Authentication failed for {form_data.user}")
|
||||
|
||||
user = Users.get_user_by_email(mail)
|
||||
if not user:
|
||||
try:
|
||||
role = (
|
||||
"admin"
|
||||
if Users.get_num_users() == 0
|
||||
else request.app.state.config.DEFAULT_USER_ROLE
|
||||
)
|
||||
|
||||
user = Auths.insert_new_auth(
|
||||
email=mail, password=str(uuid.uuid4()), name=cn, role=role
|
||||
)
|
||||
|
||||
if not user:
|
||||
raise HTTPException(
|
||||
500, detail=ERROR_MESSAGES.CREATE_USER_ERROR
|
||||
)
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as err:
|
||||
raise HTTPException(500, detail=ERROR_MESSAGES.DEFAULT(err))
|
||||
|
||||
user = Auths.authenticate_user_by_trusted_header(mail)
|
||||
|
||||
if user:
|
||||
token = create_token(
|
||||
data={"id": user.id},
|
||||
expires_delta=parse_duration(
|
||||
request.app.state.config.JWT_EXPIRES_IN
|
||||
),
|
||||
)
|
||||
|
||||
# Set the cookie token
|
||||
response.set_cookie(
|
||||
key="token",
|
||||
value=token,
|
||||
httponly=True, # Ensures the cookie is not accessible via JavaScript
|
||||
)
|
||||
|
||||
return {
|
||||
"token": token,
|
||||
"token_type": "Bearer",
|
||||
"id": user.id,
|
||||
"email": user.email,
|
||||
"name": user.name,
|
||||
"role": user.role,
|
||||
"profile_image_url": user.profile_image_url,
|
||||
}
|
||||
else:
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)
|
||||
else:
|
||||
raise HTTPException(
|
||||
400,
|
||||
f"User {form_data.user} does not match the record. Search result: {str(entry[f'{LDAP_ATTRIBUTE_FOR_USERNAME}'])}",
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(400, detail=str(e))
|
||||
|
||||
|
||||
############################
|
||||
# SignIn
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/signin", response_model=SessionUserResponse)
|
||||
async def signin(request: Request, response: Response, form_data: SigninForm):
|
||||
if WEBUI_AUTH_TRUSTED_EMAIL_HEADER:
|
||||
if WEBUI_AUTH_TRUSTED_EMAIL_HEADER not in request.headers:
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_TRUSTED_HEADER)
|
||||
|
||||
trusted_email = request.headers[WEBUI_AUTH_TRUSTED_EMAIL_HEADER].lower()
|
||||
trusted_name = trusted_email
|
||||
if WEBUI_AUTH_TRUSTED_NAME_HEADER:
|
||||
trusted_name = request.headers.get(
|
||||
WEBUI_AUTH_TRUSTED_NAME_HEADER, trusted_email
|
||||
)
|
||||
if not Users.get_user_by_email(trusted_email.lower()):
|
||||
await signup(
|
||||
request,
|
||||
response,
|
||||
SignupForm(
|
||||
email=trusted_email, password=str(uuid.uuid4()), name=trusted_name
|
||||
),
|
||||
)
|
||||
user = Auths.authenticate_user_by_trusted_header(trusted_email)
|
||||
elif WEBUI_AUTH == False:
|
||||
admin_email = "admin@localhost"
|
||||
admin_password = "admin"
|
||||
|
||||
if Users.get_user_by_email(admin_email.lower()):
|
||||
user = Auths.authenticate_user(admin_email.lower(), admin_password)
|
||||
else:
|
||||
if Users.get_num_users() != 0:
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.EXISTING_USERS)
|
||||
|
||||
await signup(
|
||||
request,
|
||||
response,
|
||||
SignupForm(email=admin_email, password=admin_password, name="User"),
|
||||
)
|
||||
|
||||
user = Auths.authenticate_user(admin_email.lower(), admin_password)
|
||||
else:
|
||||
user = Auths.authenticate_user(form_data.email.lower(), form_data.password)
|
||||
|
||||
if user:
|
||||
|
||||
expires_delta = parse_duration(request.app.state.config.JWT_EXPIRES_IN)
|
||||
expires_at = None
|
||||
if expires_delta:
|
||||
expires_at = int(time.time()) + int(expires_delta.total_seconds())
|
||||
|
||||
token = create_token(
|
||||
data={"id": user.id},
|
||||
expires_delta=expires_delta,
|
||||
)
|
||||
|
||||
datetime_expires_at = (
|
||||
datetime.datetime.fromtimestamp(expires_at, datetime.timezone.utc)
|
||||
if expires_at
|
||||
else None
|
||||
)
|
||||
|
||||
# Set the cookie token
|
||||
response.set_cookie(
|
||||
key="token",
|
||||
value=token,
|
||||
expires=datetime_expires_at,
|
||||
httponly=True, # Ensures the cookie is not accessible via JavaScript
|
||||
samesite=WEBUI_SESSION_COOKIE_SAME_SITE,
|
||||
secure=WEBUI_SESSION_COOKIE_SECURE,
|
||||
)
|
||||
|
||||
user_permissions = get_permissions(
|
||||
user.id, request.app.state.config.USER_PERMISSIONS
|
||||
)
|
||||
|
||||
return {
|
||||
"token": token,
|
||||
"token_type": "Bearer",
|
||||
"expires_at": expires_at,
|
||||
"id": user.id,
|
||||
"email": user.email,
|
||||
"name": user.name,
|
||||
"role": user.role,
|
||||
"profile_image_url": user.profile_image_url,
|
||||
"permissions": user_permissions,
|
||||
}
|
||||
else:
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)
|
||||
|
||||
|
||||
############################
|
||||
# SignUp
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/signup", response_model=SessionUserResponse)
|
||||
async def signup(request: Request, response: Response, form_data: SignupForm):
|
||||
if WEBUI_AUTH:
|
||||
if (
|
||||
not request.app.state.config.ENABLE_SIGNUP
|
||||
or not request.app.state.config.ENABLE_LOGIN_FORM
|
||||
):
|
||||
raise HTTPException(
|
||||
status.HTTP_403_FORBIDDEN, detail=ERROR_MESSAGES.ACCESS_PROHIBITED
|
||||
)
|
||||
else:
|
||||
if Users.get_num_users() != 0:
|
||||
raise HTTPException(
|
||||
status.HTTP_403_FORBIDDEN, detail=ERROR_MESSAGES.ACCESS_PROHIBITED
|
||||
)
|
||||
|
||||
if not validate_email_format(form_data.email.lower()):
|
||||
raise HTTPException(
|
||||
status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.INVALID_EMAIL_FORMAT
|
||||
)
|
||||
|
||||
if Users.get_user_by_email(form_data.email.lower()):
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.EMAIL_TAKEN)
|
||||
|
||||
try:
|
||||
role = (
|
||||
"admin"
|
||||
if Users.get_num_users() == 0
|
||||
else request.app.state.config.DEFAULT_USER_ROLE
|
||||
)
|
||||
|
||||
if Users.get_num_users() == 0:
|
||||
# Disable signup after the first user is created
|
||||
request.app.state.config.ENABLE_SIGNUP = False
|
||||
|
||||
hashed = get_password_hash(form_data.password)
|
||||
user = Auths.insert_new_auth(
|
||||
form_data.email.lower(),
|
||||
hashed,
|
||||
form_data.name,
|
||||
form_data.profile_image_url,
|
||||
role,
|
||||
)
|
||||
|
||||
if user:
|
||||
expires_delta = parse_duration(request.app.state.config.JWT_EXPIRES_IN)
|
||||
expires_at = None
|
||||
if expires_delta:
|
||||
expires_at = int(time.time()) + int(expires_delta.total_seconds())
|
||||
|
||||
token = create_token(
|
||||
data={"id": user.id},
|
||||
expires_delta=expires_delta,
|
||||
)
|
||||
|
||||
datetime_expires_at = (
|
||||
datetime.datetime.fromtimestamp(expires_at, datetime.timezone.utc)
|
||||
if expires_at
|
||||
else None
|
||||
)
|
||||
|
||||
# Set the cookie token
|
||||
response.set_cookie(
|
||||
key="token",
|
||||
value=token,
|
||||
expires=datetime_expires_at,
|
||||
httponly=True, # Ensures the cookie is not accessible via JavaScript
|
||||
samesite=WEBUI_SESSION_COOKIE_SAME_SITE,
|
||||
secure=WEBUI_SESSION_COOKIE_SECURE,
|
||||
)
|
||||
|
||||
if request.app.state.config.WEBHOOK_URL:
|
||||
post_webhook(
|
||||
request.app.state.config.WEBHOOK_URL,
|
||||
WEBHOOK_MESSAGES.USER_SIGNUP(user.name),
|
||||
{
|
||||
"action": "signup",
|
||||
"message": WEBHOOK_MESSAGES.USER_SIGNUP(user.name),
|
||||
"user": user.model_dump_json(exclude_none=True),
|
||||
},
|
||||
)
|
||||
|
||||
user_permissions = get_permissions(
|
||||
user.id, request.app.state.config.USER_PERMISSIONS
|
||||
)
|
||||
|
||||
return {
|
||||
"token": token,
|
||||
"token_type": "Bearer",
|
||||
"expires_at": expires_at,
|
||||
"id": user.id,
|
||||
"email": user.email,
|
||||
"name": user.name,
|
||||
"role": user.role,
|
||||
"profile_image_url": user.profile_image_url,
|
||||
"permissions": user_permissions,
|
||||
}
|
||||
else:
|
||||
raise HTTPException(500, detail=ERROR_MESSAGES.CREATE_USER_ERROR)
|
||||
except Exception as err:
|
||||
raise HTTPException(500, detail=ERROR_MESSAGES.DEFAULT(err))
|
||||
|
||||
|
||||
@router.get("/signout")
|
||||
async def signout(request: Request, response: Response):
|
||||
response.delete_cookie("token")
|
||||
|
||||
if ENABLE_OAUTH_SIGNUP.value:
|
||||
oauth_id_token = request.cookies.get("oauth_id_token")
|
||||
if oauth_id_token:
|
||||
try:
|
||||
async with ClientSession() as session:
|
||||
async with session.get(OPENID_PROVIDER_URL.value) as resp:
|
||||
if resp.status == 200:
|
||||
openid_data = await resp.json()
|
||||
logout_url = openid_data.get("end_session_endpoint")
|
||||
if logout_url:
|
||||
response.delete_cookie("oauth_id_token")
|
||||
return RedirectResponse(
|
||||
url=f"{logout_url}?id_token_hint={oauth_id_token}"
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=resp.status,
|
||||
detail="Failed to fetch OpenID configuration",
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
return {"status": True}
|
||||
|
||||
|
||||
############################
|
||||
# AddUser
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/add", response_model=SigninResponse)
|
||||
async def add_user(form_data: AddUserForm, user=Depends(get_admin_user)):
|
||||
if not validate_email_format(form_data.email.lower()):
|
||||
raise HTTPException(
|
||||
status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.INVALID_EMAIL_FORMAT
|
||||
)
|
||||
|
||||
if Users.get_user_by_email(form_data.email.lower()):
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.EMAIL_TAKEN)
|
||||
|
||||
try:
|
||||
print(form_data)
|
||||
hashed = get_password_hash(form_data.password)
|
||||
user = Auths.insert_new_auth(
|
||||
form_data.email.lower(),
|
||||
hashed,
|
||||
form_data.name,
|
||||
form_data.profile_image_url,
|
||||
form_data.role,
|
||||
)
|
||||
|
||||
if user:
|
||||
token = create_token(data={"id": user.id})
|
||||
return {
|
||||
"token": token,
|
||||
"token_type": "Bearer",
|
||||
"id": user.id,
|
||||
"email": user.email,
|
||||
"name": user.name,
|
||||
"role": user.role,
|
||||
"profile_image_url": user.profile_image_url,
|
||||
}
|
||||
else:
|
||||
raise HTTPException(500, detail=ERROR_MESSAGES.CREATE_USER_ERROR)
|
||||
except Exception as err:
|
||||
raise HTTPException(500, detail=ERROR_MESSAGES.DEFAULT(err))
|
||||
|
||||
|
||||
############################
|
||||
# GetAdminDetails
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/admin/details")
|
||||
async def get_admin_details(request: Request, user=Depends(get_current_user)):
|
||||
if request.app.state.config.SHOW_ADMIN_DETAILS:
|
||||
admin_email = request.app.state.config.ADMIN_EMAIL
|
||||
admin_name = None
|
||||
|
||||
print(admin_email, admin_name)
|
||||
|
||||
if admin_email:
|
||||
admin = Users.get_user_by_email(admin_email)
|
||||
if admin:
|
||||
admin_name = admin.name
|
||||
else:
|
||||
admin = Users.get_first_user()
|
||||
if admin:
|
||||
admin_email = admin.email
|
||||
admin_name = admin.name
|
||||
|
||||
return {
|
||||
"name": admin_name,
|
||||
"email": admin_email,
|
||||
}
|
||||
else:
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.ACTION_PROHIBITED)
|
||||
|
||||
|
||||
############################
|
||||
# ToggleSignUp
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/admin/config")
|
||||
async def get_admin_config(request: Request, user=Depends(get_admin_user)):
|
||||
return {
|
||||
"SHOW_ADMIN_DETAILS": request.app.state.config.SHOW_ADMIN_DETAILS,
|
||||
"ENABLE_SIGNUP": request.app.state.config.ENABLE_SIGNUP,
|
||||
"ENABLE_API_KEY": request.app.state.config.ENABLE_API_KEY,
|
||||
"DEFAULT_USER_ROLE": request.app.state.config.DEFAULT_USER_ROLE,
|
||||
"JWT_EXPIRES_IN": request.app.state.config.JWT_EXPIRES_IN,
|
||||
"ENABLE_COMMUNITY_SHARING": request.app.state.config.ENABLE_COMMUNITY_SHARING,
|
||||
"ENABLE_MESSAGE_RATING": request.app.state.config.ENABLE_MESSAGE_RATING,
|
||||
}
|
||||
|
||||
|
||||
class AdminConfig(BaseModel):
|
||||
SHOW_ADMIN_DETAILS: bool
|
||||
ENABLE_SIGNUP: bool
|
||||
ENABLE_API_KEY: bool
|
||||
DEFAULT_USER_ROLE: str
|
||||
JWT_EXPIRES_IN: str
|
||||
ENABLE_COMMUNITY_SHARING: bool
|
||||
ENABLE_MESSAGE_RATING: bool
|
||||
|
||||
|
||||
@router.post("/admin/config")
|
||||
async def update_admin_config(
|
||||
request: Request, form_data: AdminConfig, user=Depends(get_admin_user)
|
||||
):
|
||||
request.app.state.config.SHOW_ADMIN_DETAILS = form_data.SHOW_ADMIN_DETAILS
|
||||
request.app.state.config.ENABLE_SIGNUP = form_data.ENABLE_SIGNUP
|
||||
request.app.state.config.ENABLE_API_KEY = form_data.ENABLE_API_KEY
|
||||
|
||||
if form_data.DEFAULT_USER_ROLE in ["pending", "user", "admin"]:
|
||||
request.app.state.config.DEFAULT_USER_ROLE = form_data.DEFAULT_USER_ROLE
|
||||
|
||||
pattern = r"^(-1|0|(-?\d+(\.\d+)?)(ms|s|m|h|d|w))$"
|
||||
|
||||
# Check if the input string matches the pattern
|
||||
if re.match(pattern, form_data.JWT_EXPIRES_IN):
|
||||
request.app.state.config.JWT_EXPIRES_IN = form_data.JWT_EXPIRES_IN
|
||||
|
||||
request.app.state.config.ENABLE_COMMUNITY_SHARING = (
|
||||
form_data.ENABLE_COMMUNITY_SHARING
|
||||
)
|
||||
request.app.state.config.ENABLE_MESSAGE_RATING = form_data.ENABLE_MESSAGE_RATING
|
||||
|
||||
return {
|
||||
"SHOW_ADMIN_DETAILS": request.app.state.config.SHOW_ADMIN_DETAILS,
|
||||
"ENABLE_SIGNUP": request.app.state.config.ENABLE_SIGNUP,
|
||||
"ENABLE_API_KEY": request.app.state.config.ENABLE_API_KEY,
|
||||
"DEFAULT_USER_ROLE": request.app.state.config.DEFAULT_USER_ROLE,
|
||||
"JWT_EXPIRES_IN": request.app.state.config.JWT_EXPIRES_IN,
|
||||
"ENABLE_COMMUNITY_SHARING": request.app.state.config.ENABLE_COMMUNITY_SHARING,
|
||||
"ENABLE_MESSAGE_RATING": request.app.state.config.ENABLE_MESSAGE_RATING,
|
||||
}
|
||||
|
||||
|
||||
class LdapServerConfig(BaseModel):
|
||||
label: str
|
||||
host: str
|
||||
port: Optional[int] = None
|
||||
attribute_for_username: str = "uid"
|
||||
app_dn: str
|
||||
app_dn_password: str
|
||||
search_base: str
|
||||
search_filters: str = ""
|
||||
use_tls: bool = True
|
||||
certificate_path: Optional[str] = None
|
||||
ciphers: Optional[str] = "ALL"
|
||||
|
||||
|
||||
@router.get("/admin/config/ldap/server", response_model=LdapServerConfig)
|
||||
async def get_ldap_server(request: Request, user=Depends(get_admin_user)):
|
||||
return {
|
||||
"label": request.app.state.config.LDAP_SERVER_LABEL,
|
||||
"host": request.app.state.config.LDAP_SERVER_HOST,
|
||||
"port": request.app.state.config.LDAP_SERVER_PORT,
|
||||
"attribute_for_username": request.app.state.config.LDAP_ATTRIBUTE_FOR_USERNAME,
|
||||
"app_dn": request.app.state.config.LDAP_APP_DN,
|
||||
"app_dn_password": request.app.state.config.LDAP_APP_PASSWORD,
|
||||
"search_base": request.app.state.config.LDAP_SEARCH_BASE,
|
||||
"search_filters": request.app.state.config.LDAP_SEARCH_FILTERS,
|
||||
"use_tls": request.app.state.config.LDAP_USE_TLS,
|
||||
"certificate_path": request.app.state.config.LDAP_CA_CERT_FILE,
|
||||
"ciphers": request.app.state.config.LDAP_CIPHERS,
|
||||
}
|
||||
|
||||
|
||||
@router.post("/admin/config/ldap/server")
|
||||
async def update_ldap_server(
|
||||
request: Request, form_data: LdapServerConfig, user=Depends(get_admin_user)
|
||||
):
|
||||
required_fields = [
|
||||
"label",
|
||||
"host",
|
||||
"attribute_for_username",
|
||||
"app_dn",
|
||||
"app_dn_password",
|
||||
"search_base",
|
||||
]
|
||||
for key in required_fields:
|
||||
value = getattr(form_data, key)
|
||||
if not value:
|
||||
raise HTTPException(400, detail=f"Required field {key} is empty")
|
||||
|
||||
if form_data.use_tls and not form_data.certificate_path:
|
||||
raise HTTPException(
|
||||
400, detail="TLS is enabled but certificate file path is missing"
|
||||
)
|
||||
|
||||
request.app.state.config.LDAP_SERVER_LABEL = form_data.label
|
||||
request.app.state.config.LDAP_SERVER_HOST = form_data.host
|
||||
request.app.state.config.LDAP_SERVER_PORT = form_data.port
|
||||
request.app.state.config.LDAP_ATTRIBUTE_FOR_USERNAME = (
|
||||
form_data.attribute_for_username
|
||||
)
|
||||
request.app.state.config.LDAP_APP_DN = form_data.app_dn
|
||||
request.app.state.config.LDAP_APP_PASSWORD = form_data.app_dn_password
|
||||
request.app.state.config.LDAP_SEARCH_BASE = form_data.search_base
|
||||
request.app.state.config.LDAP_SEARCH_FILTERS = form_data.search_filters
|
||||
request.app.state.config.LDAP_USE_TLS = form_data.use_tls
|
||||
request.app.state.config.LDAP_CA_CERT_FILE = form_data.certificate_path
|
||||
request.app.state.config.LDAP_CIPHERS = form_data.ciphers
|
||||
|
||||
return {
|
||||
"label": request.app.state.config.LDAP_SERVER_LABEL,
|
||||
"host": request.app.state.config.LDAP_SERVER_HOST,
|
||||
"port": request.app.state.config.LDAP_SERVER_PORT,
|
||||
"attribute_for_username": request.app.state.config.LDAP_ATTRIBUTE_FOR_USERNAME,
|
||||
"app_dn": request.app.state.config.LDAP_APP_DN,
|
||||
"app_dn_password": request.app.state.config.LDAP_APP_PASSWORD,
|
||||
"search_base": request.app.state.config.LDAP_SEARCH_BASE,
|
||||
"search_filters": request.app.state.config.LDAP_SEARCH_FILTERS,
|
||||
"use_tls": request.app.state.config.LDAP_USE_TLS,
|
||||
"certificate_path": request.app.state.config.LDAP_CA_CERT_FILE,
|
||||
"ciphers": request.app.state.config.LDAP_CIPHERS,
|
||||
}
|
||||
|
||||
|
||||
@router.get("/admin/config/ldap")
|
||||
async def get_ldap_config(request: Request, user=Depends(get_admin_user)):
|
||||
return {"ENABLE_LDAP": request.app.state.config.ENABLE_LDAP}
|
||||
|
||||
|
||||
class LdapConfigForm(BaseModel):
|
||||
enable_ldap: Optional[bool] = None
|
||||
|
||||
|
||||
@router.post("/admin/config/ldap")
|
||||
async def update_ldap_config(
|
||||
request: Request, form_data: LdapConfigForm, user=Depends(get_admin_user)
|
||||
):
|
||||
request.app.state.config.ENABLE_LDAP = form_data.enable_ldap
|
||||
return {"ENABLE_LDAP": request.app.state.config.ENABLE_LDAP}
|
||||
|
||||
|
||||
############################
|
||||
# API Key
|
||||
############################
|
||||
|
||||
|
||||
# create api key
|
||||
@router.post("/api_key", response_model=ApiKey)
|
||||
async def generate_api_key(request: Request, user=Depends(get_current_user)):
|
||||
if not request.app.state.config.ENABLE_API_KEY:
|
||||
raise HTTPException(
|
||||
status.HTTP_403_FORBIDDEN,
|
||||
detail=ERROR_MESSAGES.API_KEY_CREATION_NOT_ALLOWED,
|
||||
)
|
||||
|
||||
api_key = create_api_key()
|
||||
success = Users.update_user_api_key_by_id(user.id, api_key)
|
||||
|
||||
if success:
|
||||
return {
|
||||
"api_key": api_key,
|
||||
}
|
||||
else:
|
||||
raise HTTPException(500, detail=ERROR_MESSAGES.CREATE_API_KEY_ERROR)
|
||||
|
||||
|
||||
# delete api key
|
||||
@router.delete("/api_key", response_model=bool)
|
||||
async def delete_api_key(user=Depends(get_current_user)):
|
||||
success = Users.update_user_api_key_by_id(user.id, None)
|
||||
return success
|
||||
|
||||
|
||||
# get api key
|
||||
@router.get("/api_key", response_model=ApiKey)
|
||||
async def get_api_key(user=Depends(get_current_user)):
|
||||
api_key = Users.get_user_api_key_by_id(user.id)
|
||||
if api_key:
|
||||
return {
|
||||
"api_key": api_key,
|
||||
}
|
||||
else:
|
||||
raise HTTPException(404, detail=ERROR_MESSAGES.API_KEY_NOT_FOUND)
|
||||
@@ -0,0 +1,668 @@
|
||||
import json
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from open_webui.models.chats import (
|
||||
ChatForm,
|
||||
ChatImportForm,
|
||||
ChatResponse,
|
||||
Chats,
|
||||
ChatTitleIdResponse,
|
||||
)
|
||||
from open_webui.models.tags import TagModel, Tags
|
||||
from open_webui.models.folders import Folders
|
||||
|
||||
from open_webui.config import ENABLE_ADMIN_CHAT_ACCESS, ENABLE_ADMIN_EXPORT
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
from open_webui.env import SRC_LOG_LEVELS
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, status
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
from open_webui.utils.auth import get_admin_user, get_verified_user
|
||||
from open_webui.utils.access_control import has_permission
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["MODELS"])
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
############################
|
||||
# GetChatList
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/", response_model=list[ChatTitleIdResponse])
|
||||
@router.get("/list", response_model=list[ChatTitleIdResponse])
|
||||
async def get_session_user_chat_list(
|
||||
user=Depends(get_verified_user), page: Optional[int] = None
|
||||
):
|
||||
if page is not None:
|
||||
limit = 60
|
||||
skip = (page - 1) * limit
|
||||
|
||||
return Chats.get_chat_title_id_list_by_user_id(user.id, skip=skip, limit=limit)
|
||||
else:
|
||||
return Chats.get_chat_title_id_list_by_user_id(user.id)
|
||||
|
||||
|
||||
############################
|
||||
# DeleteAllChats
|
||||
############################
|
||||
|
||||
|
||||
@router.delete("/", response_model=bool)
|
||||
async def delete_all_user_chats(request: Request, user=Depends(get_verified_user)):
|
||||
|
||||
if user.role == "user" and not has_permission(
|
||||
user.id, "chat.delete", request.app.state.config.USER_PERMISSIONS
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
|
||||
result = Chats.delete_chats_by_user_id(user.id)
|
||||
return result
|
||||
|
||||
|
||||
############################
|
||||
# GetUserChatList
|
||||
############################
|
||||
|
||||
|
||||
@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,
|
||||
):
|
||||
if not ENABLE_ADMIN_CHAT_ACCESS:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
return Chats.get_chat_list_by_user_id(
|
||||
user_id, include_archived=True, skip=skip, limit=limit
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# CreateNewChat
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/new", response_model=Optional[ChatResponse])
|
||||
async def create_new_chat(form_data: ChatForm, user=Depends(get_verified_user)):
|
||||
try:
|
||||
chat = Chats.insert_new_chat(user.id, form_data)
|
||||
return ChatResponse(**chat.model_dump())
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.DEFAULT()
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# ImportChat
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/import", response_model=Optional[ChatResponse])
|
||||
async def import_chat(form_data: ChatImportForm, user=Depends(get_verified_user)):
|
||||
try:
|
||||
chat = Chats.import_chat(user.id, form_data)
|
||||
if chat:
|
||||
tags = chat.meta.get("tags", [])
|
||||
for tag_id in tags:
|
||||
tag_id = tag_id.replace(" ", "_").lower()
|
||||
tag_name = " ".join([word.capitalize() for word in tag_id.split("_")])
|
||||
if (
|
||||
tag_id != "none"
|
||||
and Tags.get_tag_by_name_and_user_id(tag_name, user.id) is None
|
||||
):
|
||||
Tags.insert_new_tag(tag_name, user.id)
|
||||
|
||||
return ChatResponse(**chat.model_dump())
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.DEFAULT()
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# GetChats
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/search", response_model=list[ChatTitleIdResponse])
|
||||
async def search_user_chats(
|
||||
text: str, page: Optional[int] = None, user=Depends(get_verified_user)
|
||||
):
|
||||
if page is None:
|
||||
page = 1
|
||||
|
||||
limit = 60
|
||||
skip = (page - 1) * limit
|
||||
|
||||
chat_list = [
|
||||
ChatTitleIdResponse(**chat.model_dump())
|
||||
for chat in Chats.get_chats_by_user_id_and_search_text(
|
||||
user.id, text, skip=skip, limit=limit
|
||||
)
|
||||
]
|
||||
|
||||
# Delete tag if no chat is found
|
||||
words = text.strip().split(" ")
|
||||
if page == 1 and len(words) == 1 and words[0].startswith("tag:"):
|
||||
tag_id = words[0].replace("tag:", "")
|
||||
if len(chat_list) == 0:
|
||||
if Tags.get_tag_by_name_and_user_id(tag_id, user.id):
|
||||
log.debug(f"deleting tag: {tag_id}")
|
||||
Tags.delete_tag_by_name_and_user_id(tag_id, user.id)
|
||||
|
||||
return chat_list
|
||||
|
||||
|
||||
############################
|
||||
# GetChatsByFolderId
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/folder/{folder_id}", response_model=list[ChatResponse])
|
||||
async def get_chats_by_folder_id(folder_id: str, user=Depends(get_verified_user)):
|
||||
folder_ids = [folder_id]
|
||||
children_folders = Folders.get_children_folders_by_id_and_user_id(
|
||||
folder_id, user.id
|
||||
)
|
||||
if children_folders:
|
||||
folder_ids.extend([folder.id for folder in children_folders])
|
||||
|
||||
return [
|
||||
ChatResponse(**chat.model_dump())
|
||||
for chat in Chats.get_chats_by_folder_ids_and_user_id(folder_ids, user.id)
|
||||
]
|
||||
|
||||
|
||||
############################
|
||||
# GetPinnedChats
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/pinned", response_model=list[ChatResponse])
|
||||
async def get_user_pinned_chats(user=Depends(get_verified_user)):
|
||||
return [
|
||||
ChatResponse(**chat.model_dump())
|
||||
for chat in Chats.get_pinned_chats_by_user_id(user.id)
|
||||
]
|
||||
|
||||
|
||||
############################
|
||||
# GetChats
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/all", response_model=list[ChatResponse])
|
||||
async def get_user_chats(user=Depends(get_verified_user)):
|
||||
return [
|
||||
ChatResponse(**chat.model_dump())
|
||||
for chat in Chats.get_chats_by_user_id(user.id)
|
||||
]
|
||||
|
||||
|
||||
############################
|
||||
# GetArchivedChats
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/all/archived", response_model=list[ChatResponse])
|
||||
async def get_user_archived_chats(user=Depends(get_verified_user)):
|
||||
return [
|
||||
ChatResponse(**chat.model_dump())
|
||||
for chat in Chats.get_archived_chats_by_user_id(user.id)
|
||||
]
|
||||
|
||||
|
||||
############################
|
||||
# GetAllTags
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/all/tags", response_model=list[TagModel])
|
||||
async def get_all_user_tags(user=Depends(get_verified_user)):
|
||||
try:
|
||||
tags = Tags.get_tags_by_user_id(user.id)
|
||||
return tags
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.DEFAULT()
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# GetAllChatsInDB
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/all/db", response_model=list[ChatResponse])
|
||||
async def get_all_user_chats_in_db(user=Depends(get_admin_user)):
|
||||
if not ENABLE_ADMIN_EXPORT:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
return [ChatResponse(**chat.model_dump()) for chat in Chats.get_chats()]
|
||||
|
||||
|
||||
############################
|
||||
# GetArchivedChats
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/archived", response_model=list[ChatTitleIdResponse])
|
||||
async def get_archived_session_user_chat_list(
|
||||
user=Depends(get_verified_user), skip: int = 0, limit: int = 50
|
||||
):
|
||||
return Chats.get_archived_chat_list_by_user_id(user.id, skip, limit)
|
||||
|
||||
|
||||
############################
|
||||
# ArchiveAllChats
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/archive/all", response_model=bool)
|
||||
async def archive_all_chats(user=Depends(get_verified_user)):
|
||||
return Chats.archive_all_chats_by_user_id(user.id)
|
||||
|
||||
|
||||
############################
|
||||
# GetSharedChatById
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/share/{share_id}", response_model=Optional[ChatResponse])
|
||||
async def get_shared_chat_by_id(share_id: str, user=Depends(get_verified_user)):
|
||||
if user.role == "pending":
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND
|
||||
)
|
||||
|
||||
if user.role == "user" or (user.role == "admin" and not ENABLE_ADMIN_CHAT_ACCESS):
|
||||
chat = Chats.get_chat_by_share_id(share_id)
|
||||
elif user.role == "admin" and ENABLE_ADMIN_CHAT_ACCESS:
|
||||
chat = Chats.get_chat_by_id(share_id)
|
||||
|
||||
if chat:
|
||||
return ChatResponse(**chat.model_dump())
|
||||
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# GetChatsByTags
|
||||
############################
|
||||
|
||||
|
||||
class TagForm(BaseModel):
|
||||
name: str
|
||||
|
||||
|
||||
class TagFilterForm(TagForm):
|
||||
skip: Optional[int] = 0
|
||||
limit: Optional[int] = 50
|
||||
|
||||
|
||||
@router.post("/tags", response_model=list[ChatTitleIdResponse])
|
||||
async def get_user_chat_list_by_tag_name(
|
||||
form_data: TagFilterForm, user=Depends(get_verified_user)
|
||||
):
|
||||
chats = Chats.get_chat_list_by_user_id_and_tag_name(
|
||||
user.id, form_data.name, form_data.skip, form_data.limit
|
||||
)
|
||||
if len(chats) == 0:
|
||||
Tags.delete_tag_by_name_and_user_id(form_data.name, user.id)
|
||||
|
||||
return chats
|
||||
|
||||
|
||||
############################
|
||||
# GetChatById
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/{id}", response_model=Optional[ChatResponse])
|
||||
async def get_chat_by_id(id: str, user=Depends(get_verified_user)):
|
||||
chat = Chats.get_chat_by_id_and_user_id(id, user.id)
|
||||
|
||||
if chat:
|
||||
return ChatResponse(**chat.model_dump())
|
||||
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# UpdateChatById
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/{id}", response_model=Optional[ChatResponse])
|
||||
async def update_chat_by_id(
|
||||
id: str, form_data: ChatForm, user=Depends(get_verified_user)
|
||||
):
|
||||
chat = Chats.get_chat_by_id_and_user_id(id, user.id)
|
||||
if chat:
|
||||
updated_chat = {**chat.chat, **form_data.chat}
|
||||
chat = Chats.update_chat_by_id(id, updated_chat)
|
||||
return ChatResponse(**chat.model_dump())
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# DeleteChatById
|
||||
############################
|
||||
|
||||
|
||||
@router.delete("/{id}", response_model=bool)
|
||||
async def delete_chat_by_id(request: Request, id: str, user=Depends(get_verified_user)):
|
||||
if user.role == "admin":
|
||||
chat = Chats.get_chat_by_id(id)
|
||||
for tag in chat.meta.get("tags", []):
|
||||
if Chats.count_chats_by_tag_name_and_user_id(tag, user.id) == 1:
|
||||
Tags.delete_tag_by_name_and_user_id(tag, user.id)
|
||||
|
||||
result = Chats.delete_chat_by_id(id)
|
||||
|
||||
return result
|
||||
else:
|
||||
if not has_permission(
|
||||
user.id, "chat.delete", request.app.state.config.USER_PERMISSIONS
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
|
||||
chat = Chats.get_chat_by_id(id)
|
||||
for tag in chat.meta.get("tags", []):
|
||||
if Chats.count_chats_by_tag_name_and_user_id(tag, user.id) == 1:
|
||||
Tags.delete_tag_by_name_and_user_id(tag, user.id)
|
||||
|
||||
result = Chats.delete_chat_by_id_and_user_id(id, user.id)
|
||||
return result
|
||||
|
||||
|
||||
############################
|
||||
# GetPinnedStatusById
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/{id}/pinned", response_model=Optional[bool])
|
||||
async def get_pinned_status_by_id(id: str, user=Depends(get_verified_user)):
|
||||
chat = Chats.get_chat_by_id_and_user_id(id, user.id)
|
||||
if chat:
|
||||
return chat.pinned
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.DEFAULT()
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# PinChatById
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/{id}/pin", response_model=Optional[ChatResponse])
|
||||
async def pin_chat_by_id(id: str, user=Depends(get_verified_user)):
|
||||
chat = Chats.get_chat_by_id_and_user_id(id, user.id)
|
||||
if chat:
|
||||
chat = Chats.toggle_chat_pinned_by_id(id)
|
||||
return chat
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.DEFAULT()
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# CloneChat
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/{id}/clone", response_model=Optional[ChatResponse])
|
||||
async def clone_chat_by_id(id: str, user=Depends(get_verified_user)):
|
||||
chat = Chats.get_chat_by_id_and_user_id(id, user.id)
|
||||
if chat:
|
||||
updated_chat = {
|
||||
**chat.chat,
|
||||
"originalChatId": chat.id,
|
||||
"branchPointMessageId": chat.chat["history"]["currentId"],
|
||||
"title": f"Clone of {chat.title}",
|
||||
}
|
||||
|
||||
chat = Chats.insert_new_chat(user.id, ChatForm(**{"chat": updated_chat}))
|
||||
return ChatResponse(**chat.model_dump())
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.DEFAULT()
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# ArchiveChat
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/{id}/archive", response_model=Optional[ChatResponse])
|
||||
async def archive_chat_by_id(id: str, user=Depends(get_verified_user)):
|
||||
chat = Chats.get_chat_by_id_and_user_id(id, user.id)
|
||||
if chat:
|
||||
chat = Chats.toggle_chat_archive_by_id(id)
|
||||
|
||||
# Delete tags if chat is archived
|
||||
if chat.archived:
|
||||
for tag_id in chat.meta.get("tags", []):
|
||||
if Chats.count_chats_by_tag_name_and_user_id(tag_id, user.id) == 0:
|
||||
log.debug(f"deleting tag: {tag_id}")
|
||||
Tags.delete_tag_by_name_and_user_id(tag_id, user.id)
|
||||
else:
|
||||
for tag_id in chat.meta.get("tags", []):
|
||||
tag = Tags.get_tag_by_name_and_user_id(tag_id, user.id)
|
||||
if tag is None:
|
||||
log.debug(f"inserting tag: {tag_id}")
|
||||
tag = Tags.insert_new_tag(tag_id, user.id)
|
||||
|
||||
return ChatResponse(**chat.model_dump())
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.DEFAULT()
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# ShareChatById
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/{id}/share", response_model=Optional[ChatResponse])
|
||||
async def share_chat_by_id(id: str, user=Depends(get_verified_user)):
|
||||
chat = Chats.get_chat_by_id_and_user_id(id, user.id)
|
||||
if chat:
|
||||
if chat.share_id:
|
||||
shared_chat = Chats.update_shared_chat_by_chat_id(chat.id)
|
||||
return ChatResponse(**shared_chat.model_dump())
|
||||
|
||||
shared_chat = Chats.insert_shared_chat_by_chat_id(chat.id)
|
||||
if not shared_chat:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail=ERROR_MESSAGES.DEFAULT(),
|
||||
)
|
||||
return ChatResponse(**shared_chat.model_dump())
|
||||
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# DeletedSharedChatById
|
||||
############################
|
||||
|
||||
|
||||
@router.delete("/{id}/share", response_model=Optional[bool])
|
||||
async def delete_shared_chat_by_id(id: str, user=Depends(get_verified_user)):
|
||||
chat = Chats.get_chat_by_id_and_user_id(id, user.id)
|
||||
if chat:
|
||||
if not chat.share_id:
|
||||
return False
|
||||
|
||||
result = Chats.delete_shared_chat_by_chat_id(id)
|
||||
update_result = Chats.update_chat_share_id_by_id(id, None)
|
||||
|
||||
return result and update_result != None
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# UpdateChatFolderIdById
|
||||
############################
|
||||
|
||||
|
||||
class ChatFolderIdForm(BaseModel):
|
||||
folder_id: Optional[str] = None
|
||||
|
||||
|
||||
@router.post("/{id}/folder", response_model=Optional[ChatResponse])
|
||||
async def update_chat_folder_id_by_id(
|
||||
id: str, form_data: ChatFolderIdForm, user=Depends(get_verified_user)
|
||||
):
|
||||
chat = Chats.get_chat_by_id_and_user_id(id, user.id)
|
||||
if chat:
|
||||
chat = Chats.update_chat_folder_id_by_id_and_user_id(
|
||||
id, user.id, form_data.folder_id
|
||||
)
|
||||
return ChatResponse(**chat.model_dump())
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.DEFAULT()
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# GetChatTagsById
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/{id}/tags", response_model=list[TagModel])
|
||||
async def get_chat_tags_by_id(id: str, user=Depends(get_verified_user)):
|
||||
chat = Chats.get_chat_by_id_and_user_id(id, user.id)
|
||||
if chat:
|
||||
tags = chat.meta.get("tags", [])
|
||||
return Tags.get_tags_by_ids_and_user_id(tags, user.id)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# AddChatTagById
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/{id}/tags", response_model=list[TagModel])
|
||||
async def add_tag_by_id_and_tag_name(
|
||||
id: str, form_data: TagForm, user=Depends(get_verified_user)
|
||||
):
|
||||
chat = Chats.get_chat_by_id_and_user_id(id, user.id)
|
||||
if chat:
|
||||
tags = chat.meta.get("tags", [])
|
||||
tag_id = form_data.name.replace(" ", "_").lower()
|
||||
|
||||
if tag_id == "none":
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Tag name cannot be 'None'"),
|
||||
)
|
||||
|
||||
if tag_id not in tags:
|
||||
Chats.add_chat_tag_by_id_and_user_id_and_tag_name(
|
||||
id, user.id, form_data.name
|
||||
)
|
||||
|
||||
chat = Chats.get_chat_by_id_and_user_id(id, user.id)
|
||||
tags = chat.meta.get("tags", [])
|
||||
return Tags.get_tags_by_ids_and_user_id(tags, user.id)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.DEFAULT()
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# DeleteChatTagById
|
||||
############################
|
||||
|
||||
|
||||
@router.delete("/{id}/tags", response_model=list[TagModel])
|
||||
async def delete_tag_by_id_and_tag_name(
|
||||
id: str, form_data: TagForm, user=Depends(get_verified_user)
|
||||
):
|
||||
chat = Chats.get_chat_by_id_and_user_id(id, user.id)
|
||||
if chat:
|
||||
Chats.delete_tag_by_id_and_user_id_and_tag_name(id, user.id, form_data.name)
|
||||
|
||||
if Chats.count_chats_by_tag_name_and_user_id(form_data.name, user.id) == 0:
|
||||
Tags.delete_tag_by_name_and_user_id(form_data.name, user.id)
|
||||
|
||||
chat = Chats.get_chat_by_id_and_user_id(id, user.id)
|
||||
tags = chat.meta.get("tags", [])
|
||||
return Tags.get_tags_by_ids_and_user_id(tags, user.id)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# DeleteAllTagsById
|
||||
############################
|
||||
|
||||
|
||||
@router.delete("/{id}/tags/all", response_model=Optional[bool])
|
||||
async def delete_all_tags_by_id(id: str, user=Depends(get_verified_user)):
|
||||
chat = Chats.get_chat_by_id_and_user_id(id, user.id)
|
||||
if chat:
|
||||
Chats.delete_all_tags_by_id_and_user_id(id, user.id)
|
||||
|
||||
for tag in chat.meta.get("tags", []):
|
||||
if Chats.count_chats_by_tag_name_and_user_id(tag, user.id) == 0:
|
||||
Tags.delete_tag_by_name_and_user_id(tag, user.id)
|
||||
|
||||
return True
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND
|
||||
)
|
||||
@@ -0,0 +1,112 @@
|
||||
from fastapi import APIRouter, Depends, Request
|
||||
from pydantic import BaseModel
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from open_webui.utils.auth import get_admin_user, get_verified_user
|
||||
from open_webui.config import get_config, save_config
|
||||
from open_webui.config import BannerModel
|
||||
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
############################
|
||||
# ImportConfig
|
||||
############################
|
||||
|
||||
|
||||
class ImportConfigForm(BaseModel):
|
||||
config: dict
|
||||
|
||||
|
||||
@router.post("/import", response_model=dict)
|
||||
async def import_config(form_data: ImportConfigForm, user=Depends(get_admin_user)):
|
||||
save_config(form_data.config)
|
||||
return get_config()
|
||||
|
||||
|
||||
############################
|
||||
# ExportConfig
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/export", response_model=dict)
|
||||
async def export_config(user=Depends(get_admin_user)):
|
||||
return get_config()
|
||||
|
||||
|
||||
############################
|
||||
# SetDefaultModels
|
||||
############################
|
||||
class ModelsConfigForm(BaseModel):
|
||||
DEFAULT_MODELS: Optional[str]
|
||||
MODEL_ORDER_LIST: Optional[list[str]]
|
||||
|
||||
|
||||
@router.get("/models", response_model=ModelsConfigForm)
|
||||
async def get_models_config(request: Request, user=Depends(get_admin_user)):
|
||||
return {
|
||||
"DEFAULT_MODELS": request.app.state.config.DEFAULT_MODELS,
|
||||
"MODEL_ORDER_LIST": request.app.state.config.MODEL_ORDER_LIST,
|
||||
}
|
||||
|
||||
|
||||
@router.post("/models", response_model=ModelsConfigForm)
|
||||
async def set_models_config(
|
||||
request: Request, form_data: ModelsConfigForm, user=Depends(get_admin_user)
|
||||
):
|
||||
request.app.state.config.DEFAULT_MODELS = form_data.DEFAULT_MODELS
|
||||
request.app.state.config.MODEL_ORDER_LIST = form_data.MODEL_ORDER_LIST
|
||||
return {
|
||||
"DEFAULT_MODELS": request.app.state.config.DEFAULT_MODELS,
|
||||
"MODEL_ORDER_LIST": request.app.state.config.MODEL_ORDER_LIST,
|
||||
}
|
||||
|
||||
|
||||
class PromptSuggestion(BaseModel):
|
||||
title: list[str]
|
||||
content: str
|
||||
|
||||
|
||||
class SetDefaultSuggestionsForm(BaseModel):
|
||||
suggestions: list[PromptSuggestion]
|
||||
|
||||
|
||||
@router.post("/suggestions", response_model=list[PromptSuggestion])
|
||||
async def set_default_suggestions(
|
||||
request: Request,
|
||||
form_data: SetDefaultSuggestionsForm,
|
||||
user=Depends(get_admin_user),
|
||||
):
|
||||
data = form_data.model_dump()
|
||||
request.app.state.config.DEFAULT_PROMPT_SUGGESTIONS = data["suggestions"]
|
||||
return request.app.state.config.DEFAULT_PROMPT_SUGGESTIONS
|
||||
|
||||
|
||||
############################
|
||||
# SetBanners
|
||||
############################
|
||||
|
||||
|
||||
class SetBannersForm(BaseModel):
|
||||
banners: list[BannerModel]
|
||||
|
||||
|
||||
@router.post("/banners", response_model=list[BannerModel])
|
||||
async def set_banners(
|
||||
request: Request,
|
||||
form_data: SetBannersForm,
|
||||
user=Depends(get_admin_user),
|
||||
):
|
||||
data = form_data.model_dump()
|
||||
request.app.state.config.BANNERS = data["banners"]
|
||||
return request.app.state.config.BANNERS
|
||||
|
||||
|
||||
@router.get("/banners", response_model=list[BannerModel])
|
||||
async def get_banners(
|
||||
request: Request,
|
||||
user=Depends(get_verified_user),
|
||||
):
|
||||
return request.app.state.config.BANNERS
|
||||
@@ -0,0 +1,159 @@
|
||||
from typing import Optional
|
||||
from fastapi import APIRouter, Depends, HTTPException, status, Request
|
||||
from pydantic import BaseModel
|
||||
|
||||
from open_webui.models.users import Users, UserModel
|
||||
from open_webui.models.feedbacks import (
|
||||
FeedbackModel,
|
||||
FeedbackResponse,
|
||||
FeedbackForm,
|
||||
Feedbacks,
|
||||
)
|
||||
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
from open_webui.utils.auth import get_admin_user, get_verified_user
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
############################
|
||||
# GetConfig
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/config")
|
||||
async def get_config(request: Request, user=Depends(get_admin_user)):
|
||||
return {
|
||||
"ENABLE_EVALUATION_ARENA_MODELS": request.app.state.config.ENABLE_EVALUATION_ARENA_MODELS,
|
||||
"EVALUATION_ARENA_MODELS": request.app.state.config.EVALUATION_ARENA_MODELS,
|
||||
}
|
||||
|
||||
|
||||
############################
|
||||
# UpdateConfig
|
||||
############################
|
||||
|
||||
|
||||
class UpdateConfigForm(BaseModel):
|
||||
ENABLE_EVALUATION_ARENA_MODELS: Optional[bool] = None
|
||||
EVALUATION_ARENA_MODELS: Optional[list[dict]] = None
|
||||
|
||||
|
||||
@router.post("/config")
|
||||
async def update_config(
|
||||
request: Request,
|
||||
form_data: UpdateConfigForm,
|
||||
user=Depends(get_admin_user),
|
||||
):
|
||||
config = request.app.state.config
|
||||
if form_data.ENABLE_EVALUATION_ARENA_MODELS is not None:
|
||||
config.ENABLE_EVALUATION_ARENA_MODELS = form_data.ENABLE_EVALUATION_ARENA_MODELS
|
||||
if form_data.EVALUATION_ARENA_MODELS is not None:
|
||||
config.EVALUATION_ARENA_MODELS = form_data.EVALUATION_ARENA_MODELS
|
||||
return {
|
||||
"ENABLE_EVALUATION_ARENA_MODELS": config.ENABLE_EVALUATION_ARENA_MODELS,
|
||||
"EVALUATION_ARENA_MODELS": config.EVALUATION_ARENA_MODELS,
|
||||
}
|
||||
|
||||
|
||||
class FeedbackUserResponse(FeedbackResponse):
|
||||
user: Optional[UserModel] = None
|
||||
|
||||
|
||||
@router.get("/feedbacks/all", response_model=list[FeedbackUserResponse])
|
||||
async def get_all_feedbacks(user=Depends(get_admin_user)):
|
||||
feedbacks = Feedbacks.get_all_feedbacks()
|
||||
return [
|
||||
FeedbackUserResponse(
|
||||
**feedback.model_dump(), user=Users.get_user_by_id(feedback.user_id)
|
||||
)
|
||||
for feedback in feedbacks
|
||||
]
|
||||
|
||||
|
||||
@router.delete("/feedbacks/all")
|
||||
async def delete_all_feedbacks(user=Depends(get_admin_user)):
|
||||
success = Feedbacks.delete_all_feedbacks()
|
||||
return success
|
||||
|
||||
|
||||
@router.get("/feedbacks/all/export", response_model=list[FeedbackModel])
|
||||
async def get_all_feedbacks(user=Depends(get_admin_user)):
|
||||
feedbacks = Feedbacks.get_all_feedbacks()
|
||||
return [
|
||||
FeedbackModel(
|
||||
**feedback.model_dump(), user=Users.get_user_by_id(feedback.user_id)
|
||||
)
|
||||
for feedback in feedbacks
|
||||
]
|
||||
|
||||
|
||||
@router.get("/feedbacks/user", response_model=list[FeedbackUserResponse])
|
||||
async def get_feedbacks(user=Depends(get_verified_user)):
|
||||
feedbacks = Feedbacks.get_feedbacks_by_user_id(user.id)
|
||||
return feedbacks
|
||||
|
||||
|
||||
@router.delete("/feedbacks", response_model=bool)
|
||||
async def delete_feedbacks(user=Depends(get_verified_user)):
|
||||
success = Feedbacks.delete_feedbacks_by_user_id(user.id)
|
||||
return success
|
||||
|
||||
|
||||
@router.post("/feedback", response_model=FeedbackModel)
|
||||
async def create_feedback(
|
||||
request: Request,
|
||||
form_data: FeedbackForm,
|
||||
user=Depends(get_verified_user),
|
||||
):
|
||||
feedback = Feedbacks.insert_new_feedback(user_id=user.id, form_data=form_data)
|
||||
if not feedback:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(),
|
||||
)
|
||||
|
||||
return feedback
|
||||
|
||||
|
||||
@router.get("/feedback/{id}", response_model=FeedbackModel)
|
||||
async def get_feedback_by_id(id: str, user=Depends(get_verified_user)):
|
||||
feedback = Feedbacks.get_feedback_by_id_and_user_id(id=id, user_id=user.id)
|
||||
|
||||
if not feedback:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND, detail=ERROR_MESSAGES.NOT_FOUND
|
||||
)
|
||||
|
||||
return feedback
|
||||
|
||||
|
||||
@router.post("/feedback/{id}", response_model=FeedbackModel)
|
||||
async def update_feedback_by_id(
|
||||
id: str, form_data: FeedbackForm, user=Depends(get_verified_user)
|
||||
):
|
||||
feedback = Feedbacks.update_feedback_by_id_and_user_id(
|
||||
id=id, user_id=user.id, form_data=form_data
|
||||
)
|
||||
|
||||
if not feedback:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND, detail=ERROR_MESSAGES.NOT_FOUND
|
||||
)
|
||||
|
||||
return feedback
|
||||
|
||||
|
||||
@router.delete("/feedback/{id}")
|
||||
async def delete_feedback_by_id(id: str, user=Depends(get_verified_user)):
|
||||
if user.role == "admin":
|
||||
success = Feedbacks.delete_feedback_by_id(id=id)
|
||||
else:
|
||||
success = Feedbacks.delete_feedback_by_id_and_user_id(id=id, user_id=user.id)
|
||||
|
||||
if not success:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND, detail=ERROR_MESSAGES.NOT_FOUND
|
||||
)
|
||||
|
||||
return success
|
||||
@@ -0,0 +1,353 @@
|
||||
import logging
|
||||
import os
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
from pydantic import BaseModel
|
||||
import mimetypes
|
||||
|
||||
from open_webui.storage.provider import Storage
|
||||
|
||||
from open_webui.models.files import (
|
||||
FileForm,
|
||||
FileModel,
|
||||
FileModelResponse,
|
||||
Files,
|
||||
)
|
||||
from open_webui.routers.retrieval import process_file, ProcessFileForm
|
||||
|
||||
from open_webui.config import UPLOAD_DIR
|
||||
from open_webui.env import SRC_LOG_LEVELS
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
|
||||
|
||||
from fastapi import APIRouter, Depends, File, HTTPException, UploadFile, status, Request
|
||||
from fastapi.responses import FileResponse, StreamingResponse
|
||||
|
||||
|
||||
from open_webui.utils.auth import get_admin_user, get_verified_user
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["MODELS"])
|
||||
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
############################
|
||||
# Upload File
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/", response_model=FileModelResponse)
|
||||
def upload_file(
|
||||
request: Request, file: UploadFile = File(...), user=Depends(get_verified_user)
|
||||
):
|
||||
log.info(f"file.content_type: {file.content_type}")
|
||||
try:
|
||||
unsanitized_filename = file.filename
|
||||
filename = os.path.basename(unsanitized_filename)
|
||||
|
||||
# replace filename with uuid
|
||||
id = str(uuid.uuid4())
|
||||
name = filename
|
||||
filename = f"{id}_{filename}"
|
||||
contents, file_path = Storage.upload_file(file.file, filename)
|
||||
|
||||
file_item = Files.insert_new_file(
|
||||
user.id,
|
||||
FileForm(
|
||||
**{
|
||||
"id": id,
|
||||
"filename": name,
|
||||
"path": file_path,
|
||||
"meta": {
|
||||
"name": name,
|
||||
"content_type": file.content_type,
|
||||
"size": len(contents),
|
||||
},
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
try:
|
||||
process_file(request, ProcessFileForm(file_id=id))
|
||||
file_item = Files.get_file_by_id(id=id)
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
log.error(f"Error processing file: {file_item.id}")
|
||||
file_item = FileModelResponse(
|
||||
**{
|
||||
**file_item.model_dump(),
|
||||
"error": str(e.detail) if hasattr(e, "detail") else str(e),
|
||||
}
|
||||
)
|
||||
|
||||
if file_item:
|
||||
return file_item
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error uploading file"),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(e),
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# List Files
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/", response_model=list[FileModelResponse])
|
||||
async def list_files(user=Depends(get_verified_user)):
|
||||
if user.role == "admin":
|
||||
files = Files.get_files()
|
||||
else:
|
||||
files = Files.get_files_by_user_id(user.id)
|
||||
return files
|
||||
|
||||
|
||||
############################
|
||||
# Delete All Files
|
||||
############################
|
||||
|
||||
|
||||
@router.delete("/all")
|
||||
async def delete_all_files(user=Depends(get_admin_user)):
|
||||
result = Files.delete_all_files()
|
||||
if result:
|
||||
try:
|
||||
Storage.delete_all_files()
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
log.error(f"Error deleting files")
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error deleting files"),
|
||||
)
|
||||
return {"message": "All files deleted successfully"}
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error deleting files"),
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# Get File By Id
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/{id}", response_model=Optional[FileModel])
|
||||
async def get_file_by_id(id: str, user=Depends(get_verified_user)):
|
||||
file = Files.get_file_by_id(id)
|
||||
|
||||
if file and (file.user_id == user.id or user.role == "admin"):
|
||||
return file
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# Get File Data Content By Id
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/{id}/data/content")
|
||||
async def get_file_data_content_by_id(id: str, user=Depends(get_verified_user)):
|
||||
file = Files.get_file_by_id(id)
|
||||
|
||||
if file and (file.user_id == user.id or user.role == "admin"):
|
||||
return {"content": file.data.get("content", "")}
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# Update File Data Content By Id
|
||||
############################
|
||||
|
||||
|
||||
class ContentForm(BaseModel):
|
||||
content: str
|
||||
|
||||
|
||||
@router.post("/{id}/data/content/update")
|
||||
async def update_file_data_content_by_id(
|
||||
request: Request, id: str, form_data: ContentForm, user=Depends(get_verified_user)
|
||||
):
|
||||
file = Files.get_file_by_id(id)
|
||||
|
||||
if file and (file.user_id == user.id or user.role == "admin"):
|
||||
try:
|
||||
process_file(
|
||||
request, ProcessFileForm(file_id=id, content=form_data.content)
|
||||
)
|
||||
file = Files.get_file_by_id(id=id)
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
log.error(f"Error processing file: {file.id}")
|
||||
|
||||
return {"content": file.data.get("content", "")}
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# Get File Content By Id
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/{id}/content")
|
||||
async def get_file_content_by_id(id: str, user=Depends(get_verified_user)):
|
||||
file = Files.get_file_by_id(id)
|
||||
if file and (file.user_id == user.id or user.role == "admin"):
|
||||
try:
|
||||
file_path = Storage.get_file(file.path)
|
||||
file_path = Path(file_path)
|
||||
|
||||
# Check if the file already exists in the cache
|
||||
if file_path.is_file():
|
||||
print(f"file_path: {file_path}")
|
||||
headers = {
|
||||
"Content-Disposition": f'attachment; filename="{file.meta.get("name", file.filename)}"'
|
||||
}
|
||||
return FileResponse(file_path, headers=headers)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
log.error(f"Error getting file content")
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error getting file content"),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{id}/content/html")
|
||||
async def get_html_file_content_by_id(id: str, user=Depends(get_verified_user)):
|
||||
file = Files.get_file_by_id(id)
|
||||
if file and (file.user_id == user.id or user.role == "admin"):
|
||||
try:
|
||||
file_path = Storage.get_file(file.path)
|
||||
file_path = Path(file_path)
|
||||
|
||||
# Check if the file already exists in the cache
|
||||
if file_path.is_file():
|
||||
print(f"file_path: {file_path}")
|
||||
return FileResponse(file_path)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
log.error(f"Error getting file content")
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error getting file content"),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{id}/content/{file_name}")
|
||||
async def get_file_content_by_id(id: str, user=Depends(get_verified_user)):
|
||||
file = Files.get_file_by_id(id)
|
||||
|
||||
if file and (file.user_id == user.id or user.role == "admin"):
|
||||
file_path = file.path
|
||||
if file_path:
|
||||
file_path = Storage.get_file(file_path)
|
||||
file_path = Path(file_path)
|
||||
|
||||
# Check if the file already exists in the cache
|
||||
if file_path.is_file():
|
||||
print(f"file_path: {file_path}")
|
||||
headers = {
|
||||
"Content-Disposition": f'attachment; filename="{file.meta.get("name", file.filename)}"'
|
||||
}
|
||||
return FileResponse(file_path, headers=headers)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
else:
|
||||
# File path doesn’t exist, return the content as .txt if possible
|
||||
file_content = file.content.get("content", "")
|
||||
file_name = file.filename
|
||||
|
||||
# Create a generator that encodes the file content
|
||||
def generator():
|
||||
yield file_content.encode("utf-8")
|
||||
|
||||
return StreamingResponse(
|
||||
generator(),
|
||||
media_type="text/plain",
|
||||
headers={"Content-Disposition": f"attachment; filename={file_name}"},
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# Delete File By Id
|
||||
############################
|
||||
|
||||
|
||||
@router.delete("/{id}")
|
||||
async def delete_file_by_id(id: str, user=Depends(get_verified_user)):
|
||||
file = Files.get_file_by_id(id)
|
||||
if file and (file.user_id == user.id or user.role == "admin"):
|
||||
result = Files.delete_file_by_id(id)
|
||||
if result:
|
||||
try:
|
||||
Storage.delete_file(file.filename)
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
log.error(f"Error deleting files")
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error deleting files"),
|
||||
)
|
||||
return {"message": "File deleted successfully"}
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error deleting file"),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
@@ -0,0 +1,251 @@
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
from pydantic import BaseModel
|
||||
import mimetypes
|
||||
|
||||
|
||||
from open_webui.models.folders import (
|
||||
FolderForm,
|
||||
FolderModel,
|
||||
Folders,
|
||||
)
|
||||
from open_webui.models.chats import Chats
|
||||
|
||||
from open_webui.config import UPLOAD_DIR
|
||||
from open_webui.env import SRC_LOG_LEVELS
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
|
||||
|
||||
from fastapi import APIRouter, Depends, File, HTTPException, UploadFile, status
|
||||
from fastapi.responses import FileResponse, StreamingResponse
|
||||
|
||||
|
||||
from open_webui.utils.auth import get_admin_user, get_verified_user
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["MODELS"])
|
||||
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
############################
|
||||
# Get Folders
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/", response_model=list[FolderModel])
|
||||
async def get_folders(user=Depends(get_verified_user)):
|
||||
folders = Folders.get_folders_by_user_id(user.id)
|
||||
|
||||
return [
|
||||
{
|
||||
**folder.model_dump(),
|
||||
"items": {
|
||||
"chats": [
|
||||
{"title": chat.title, "id": chat.id}
|
||||
for chat in Chats.get_chats_by_folder_id_and_user_id(
|
||||
folder.id, user.id
|
||||
)
|
||||
]
|
||||
},
|
||||
}
|
||||
for folder in folders
|
||||
]
|
||||
|
||||
|
||||
############################
|
||||
# Create Folder
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/")
|
||||
def create_folder(form_data: FolderForm, user=Depends(get_verified_user)):
|
||||
folder = Folders.get_folder_by_parent_id_and_user_id_and_name(
|
||||
None, user.id, form_data.name
|
||||
)
|
||||
|
||||
if folder:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Folder already exists"),
|
||||
)
|
||||
|
||||
try:
|
||||
folder = Folders.insert_new_folder(user.id, form_data.name)
|
||||
return folder
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
log.error("Error creating folder")
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error creating folder"),
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# Get Folders By Id
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/{id}", response_model=Optional[FolderModel])
|
||||
async def get_folder_by_id(id: str, user=Depends(get_verified_user)):
|
||||
folder = Folders.get_folder_by_id_and_user_id(id, user.id)
|
||||
if folder:
|
||||
return folder
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# Update Folder Name By Id
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/{id}/update")
|
||||
async def update_folder_name_by_id(
|
||||
id: str, form_data: FolderForm, user=Depends(get_verified_user)
|
||||
):
|
||||
folder = Folders.get_folder_by_id_and_user_id(id, user.id)
|
||||
if folder:
|
||||
existing_folder = Folders.get_folder_by_parent_id_and_user_id_and_name(
|
||||
folder.parent_id, user.id, form_data.name
|
||||
)
|
||||
if existing_folder:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Folder already exists"),
|
||||
)
|
||||
|
||||
try:
|
||||
folder = Folders.update_folder_name_by_id_and_user_id(
|
||||
id, user.id, form_data.name
|
||||
)
|
||||
|
||||
return folder
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
log.error(f"Error updating folder: {id}")
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error updating folder"),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# Update Folder Parent Id By Id
|
||||
############################
|
||||
|
||||
|
||||
class FolderParentIdForm(BaseModel):
|
||||
parent_id: Optional[str] = None
|
||||
|
||||
|
||||
@router.post("/{id}/update/parent")
|
||||
async def update_folder_parent_id_by_id(
|
||||
id: str, form_data: FolderParentIdForm, user=Depends(get_verified_user)
|
||||
):
|
||||
folder = Folders.get_folder_by_id_and_user_id(id, user.id)
|
||||
if folder:
|
||||
existing_folder = Folders.get_folder_by_parent_id_and_user_id_and_name(
|
||||
form_data.parent_id, user.id, folder.name
|
||||
)
|
||||
|
||||
if existing_folder:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Folder already exists"),
|
||||
)
|
||||
|
||||
try:
|
||||
folder = Folders.update_folder_parent_id_by_id_and_user_id(
|
||||
id, user.id, form_data.parent_id
|
||||
)
|
||||
return folder
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
log.error(f"Error updating folder: {id}")
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error updating folder"),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# Update Folder Is Expanded By Id
|
||||
############################
|
||||
|
||||
|
||||
class FolderIsExpandedForm(BaseModel):
|
||||
is_expanded: bool
|
||||
|
||||
|
||||
@router.post("/{id}/update/expanded")
|
||||
async def update_folder_is_expanded_by_id(
|
||||
id: str, form_data: FolderIsExpandedForm, user=Depends(get_verified_user)
|
||||
):
|
||||
folder = Folders.get_folder_by_id_and_user_id(id, user.id)
|
||||
if folder:
|
||||
try:
|
||||
folder = Folders.update_folder_is_expanded_by_id_and_user_id(
|
||||
id, user.id, form_data.is_expanded
|
||||
)
|
||||
return folder
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
log.error(f"Error updating folder: {id}")
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error updating folder"),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# Delete Folder By Id
|
||||
############################
|
||||
|
||||
|
||||
@router.delete("/{id}")
|
||||
async def delete_folder_by_id(id: str, user=Depends(get_verified_user)):
|
||||
folder = Folders.get_folder_by_id_and_user_id(id, user.id)
|
||||
if folder:
|
||||
try:
|
||||
result = Folders.delete_folder_by_id_and_user_id(id, user.id)
|
||||
if result:
|
||||
return result
|
||||
else:
|
||||
raise Exception("Error deleting folder")
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
log.error(f"Error deleting folder: {id}")
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error deleting folder"),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
@@ -0,0 +1,405 @@
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from open_webui.models.functions import (
|
||||
FunctionForm,
|
||||
FunctionModel,
|
||||
FunctionResponse,
|
||||
Functions,
|
||||
)
|
||||
from open_webui.utils.plugin import load_function_module_by_id, replace_imports
|
||||
from open_webui.config import CACHE_DIR
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, status
|
||||
from open_webui.utils.auth import get_admin_user, get_verified_user
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
############################
|
||||
# GetFunctions
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/", response_model=list[FunctionResponse])
|
||||
async def get_functions(user=Depends(get_verified_user)):
|
||||
return Functions.get_functions()
|
||||
|
||||
|
||||
############################
|
||||
# ExportFunctions
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/export", response_model=list[FunctionModel])
|
||||
async def get_functions(user=Depends(get_admin_user)):
|
||||
return Functions.get_functions()
|
||||
|
||||
|
||||
############################
|
||||
# CreateNewFunction
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/create", response_model=Optional[FunctionResponse])
|
||||
async def create_new_function(
|
||||
request: Request, form_data: FunctionForm, user=Depends(get_admin_user)
|
||||
):
|
||||
if not form_data.id.isidentifier():
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Only alphanumeric characters and underscores are allowed in the id",
|
||||
)
|
||||
|
||||
form_data.id = form_data.id.lower()
|
||||
|
||||
function = Functions.get_function_by_id(form_data.id)
|
||||
if function is None:
|
||||
try:
|
||||
form_data.content = replace_imports(form_data.content)
|
||||
function_module, function_type, frontmatter = load_function_module_by_id(
|
||||
form_data.id,
|
||||
content=form_data.content,
|
||||
)
|
||||
form_data.meta.manifest = frontmatter
|
||||
|
||||
FUNCTIONS = request.app.state.FUNCTIONS
|
||||
FUNCTIONS[form_data.id] = function_module
|
||||
|
||||
function = Functions.insert_new_function(user.id, function_type, form_data)
|
||||
|
||||
function_cache_dir = Path(CACHE_DIR) / "functions" / form_data.id
|
||||
function_cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if function:
|
||||
return function
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error creating function"),
|
||||
)
|
||||
except Exception as e:
|
||||
print(e)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(e),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.ID_TAKEN,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# GetFunctionById
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/id/{id}", response_model=Optional[FunctionModel])
|
||||
async def get_function_by_id(id: str, user=Depends(get_admin_user)):
|
||||
function = Functions.get_function_by_id(id)
|
||||
|
||||
if function:
|
||||
return function
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# ToggleFunctionById
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/id/{id}/toggle", response_model=Optional[FunctionModel])
|
||||
async def toggle_function_by_id(id: str, user=Depends(get_admin_user)):
|
||||
function = Functions.get_function_by_id(id)
|
||||
if function:
|
||||
function = Functions.update_function_by_id(
|
||||
id, {"is_active": not function.is_active}
|
||||
)
|
||||
|
||||
if function:
|
||||
return function
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error updating function"),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# ToggleGlobalById
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/id/{id}/toggle/global", response_model=Optional[FunctionModel])
|
||||
async def toggle_global_by_id(id: str, user=Depends(get_admin_user)):
|
||||
function = Functions.get_function_by_id(id)
|
||||
if function:
|
||||
function = Functions.update_function_by_id(
|
||||
id, {"is_global": not function.is_global}
|
||||
)
|
||||
|
||||
if function:
|
||||
return function
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error updating function"),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# UpdateFunctionById
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/id/{id}/update", response_model=Optional[FunctionModel])
|
||||
async def update_function_by_id(
|
||||
request: Request, id: str, form_data: FunctionForm, user=Depends(get_admin_user)
|
||||
):
|
||||
try:
|
||||
form_data.content = replace_imports(form_data.content)
|
||||
function_module, function_type, frontmatter = load_function_module_by_id(
|
||||
id, content=form_data.content
|
||||
)
|
||||
form_data.meta.manifest = frontmatter
|
||||
|
||||
FUNCTIONS = request.app.state.FUNCTIONS
|
||||
FUNCTIONS[id] = function_module
|
||||
|
||||
updated = {**form_data.model_dump(exclude={"id"}), "type": function_type}
|
||||
print(updated)
|
||||
|
||||
function = Functions.update_function_by_id(id, updated)
|
||||
|
||||
if function:
|
||||
return function
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error updating function"),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(e),
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# DeleteFunctionById
|
||||
############################
|
||||
|
||||
|
||||
@router.delete("/id/{id}/delete", response_model=bool)
|
||||
async def delete_function_by_id(
|
||||
request: Request, id: str, user=Depends(get_admin_user)
|
||||
):
|
||||
result = Functions.delete_function_by_id(id)
|
||||
|
||||
if result:
|
||||
FUNCTIONS = request.app.state.FUNCTIONS
|
||||
if id in FUNCTIONS:
|
||||
del FUNCTIONS[id]
|
||||
|
||||
return result
|
||||
|
||||
|
||||
############################
|
||||
# GetFunctionValves
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/id/{id}/valves", response_model=Optional[dict])
|
||||
async def get_function_valves_by_id(id: str, user=Depends(get_admin_user)):
|
||||
function = Functions.get_function_by_id(id)
|
||||
if function:
|
||||
try:
|
||||
valves = Functions.get_function_valves_by_id(id)
|
||||
return valves
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(e),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# GetFunctionValvesSpec
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/id/{id}/valves/spec", response_model=Optional[dict])
|
||||
async def get_function_valves_spec_by_id(
|
||||
request: Request, id: str, user=Depends(get_admin_user)
|
||||
):
|
||||
function = Functions.get_function_by_id(id)
|
||||
if function:
|
||||
if id in request.app.state.FUNCTIONS:
|
||||
function_module = request.app.state.FUNCTIONS[id]
|
||||
else:
|
||||
function_module, function_type, frontmatter = load_function_module_by_id(id)
|
||||
request.app.state.FUNCTIONS[id] = function_module
|
||||
|
||||
if hasattr(function_module, "Valves"):
|
||||
Valves = function_module.Valves
|
||||
return Valves.schema()
|
||||
return None
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# UpdateFunctionValves
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/id/{id}/valves/update", response_model=Optional[dict])
|
||||
async def update_function_valves_by_id(
|
||||
request: Request, id: str, form_data: dict, user=Depends(get_admin_user)
|
||||
):
|
||||
function = Functions.get_function_by_id(id)
|
||||
if function:
|
||||
if id in request.app.state.FUNCTIONS:
|
||||
function_module = request.app.state.FUNCTIONS[id]
|
||||
else:
|
||||
function_module, function_type, frontmatter = load_function_module_by_id(id)
|
||||
request.app.state.FUNCTIONS[id] = function_module
|
||||
|
||||
if hasattr(function_module, "Valves"):
|
||||
Valves = function_module.Valves
|
||||
|
||||
try:
|
||||
form_data = {k: v for k, v in form_data.items() if v is not None}
|
||||
valves = Valves(**form_data)
|
||||
Functions.update_function_valves_by_id(id, valves.model_dump())
|
||||
return valves.model_dump()
|
||||
except Exception as e:
|
||||
print(e)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(e),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# FunctionUserValves
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/id/{id}/valves/user", response_model=Optional[dict])
|
||||
async def get_function_user_valves_by_id(id: str, user=Depends(get_verified_user)):
|
||||
function = Functions.get_function_by_id(id)
|
||||
if function:
|
||||
try:
|
||||
user_valves = Functions.get_user_valves_by_id_and_user_id(id, user.id)
|
||||
return user_valves
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(e),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
@router.get("/id/{id}/valves/user/spec", response_model=Optional[dict])
|
||||
async def get_function_user_valves_spec_by_id(
|
||||
request: Request, id: str, user=Depends(get_verified_user)
|
||||
):
|
||||
function = Functions.get_function_by_id(id)
|
||||
if function:
|
||||
if id in request.app.state.FUNCTIONS:
|
||||
function_module = request.app.state.FUNCTIONS[id]
|
||||
else:
|
||||
function_module, function_type, frontmatter = load_function_module_by_id(id)
|
||||
request.app.state.FUNCTIONS[id] = function_module
|
||||
|
||||
if hasattr(function_module, "UserValves"):
|
||||
UserValves = function_module.UserValves
|
||||
return UserValves.schema()
|
||||
return None
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
@router.post("/id/{id}/valves/user/update", response_model=Optional[dict])
|
||||
async def update_function_user_valves_by_id(
|
||||
request: Request, id: str, form_data: dict, user=Depends(get_verified_user)
|
||||
):
|
||||
function = Functions.get_function_by_id(id)
|
||||
|
||||
if function:
|
||||
if id in request.app.state.FUNCTIONS:
|
||||
function_module = request.app.state.FUNCTIONS[id]
|
||||
else:
|
||||
function_module, function_type, frontmatter = load_function_module_by_id(id)
|
||||
request.app.state.FUNCTIONS[id] = function_module
|
||||
|
||||
if hasattr(function_module, "UserValves"):
|
||||
UserValves = function_module.UserValves
|
||||
|
||||
try:
|
||||
form_data = {k: v for k, v in form_data.items() if v is not None}
|
||||
user_valves = UserValves(**form_data)
|
||||
Functions.update_user_valves_by_id_and_user_id(
|
||||
id, user.id, user_valves.model_dump()
|
||||
)
|
||||
return user_valves.model_dump()
|
||||
except Exception as e:
|
||||
print(e)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(e),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
@@ -0,0 +1,120 @@
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from open_webui.models.groups import (
|
||||
Groups,
|
||||
GroupForm,
|
||||
GroupUpdateForm,
|
||||
GroupResponse,
|
||||
)
|
||||
|
||||
from open_webui.config import CACHE_DIR
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, status
|
||||
from open_webui.utils.auth import get_admin_user, get_verified_user
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
############################
|
||||
# GetFunctions
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/", response_model=list[GroupResponse])
|
||||
async def get_groups(user=Depends(get_verified_user)):
|
||||
if user.role == "admin":
|
||||
return Groups.get_groups()
|
||||
else:
|
||||
return Groups.get_groups_by_member_id(user.id)
|
||||
|
||||
|
||||
############################
|
||||
# CreateNewGroup
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/create", response_model=Optional[GroupResponse])
|
||||
async def create_new_function(form_data: GroupForm, user=Depends(get_admin_user)):
|
||||
try:
|
||||
group = Groups.insert_new_group(user.id, form_data)
|
||||
if group:
|
||||
return group
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error creating group"),
|
||||
)
|
||||
except Exception as e:
|
||||
print(e)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(e),
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# GetGroupById
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/id/{id}", response_model=Optional[GroupResponse])
|
||||
async def get_group_by_id(id: str, user=Depends(get_admin_user)):
|
||||
group = Groups.get_group_by_id(id)
|
||||
if group:
|
||||
return group
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# UpdateGroupById
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/id/{id}/update", response_model=Optional[GroupResponse])
|
||||
async def update_group_by_id(
|
||||
id: str, form_data: GroupUpdateForm, user=Depends(get_admin_user)
|
||||
):
|
||||
try:
|
||||
group = Groups.update_group_by_id(id, form_data)
|
||||
if group:
|
||||
return group
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error updating group"),
|
||||
)
|
||||
except Exception as e:
|
||||
print(e)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(e),
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# DeleteGroupById
|
||||
############################
|
||||
|
||||
|
||||
@router.delete("/id/{id}/delete", response_model=bool)
|
||||
async def delete_group_by_id(id: str, user=Depends(get_admin_user)):
|
||||
try:
|
||||
result = Groups.delete_group_by_id(id)
|
||||
if result:
|
||||
return result
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error deleting group"),
|
||||
)
|
||||
except Exception as e:
|
||||
print(e)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(e),
|
||||
)
|
||||
@@ -0,0 +1,596 @@
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import mimetypes
|
||||
import re
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
from fastapi import Depends, FastAPI, HTTPException, Request, APIRouter
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
from open_webui.config import CACHE_DIR
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
from open_webui.env import ENV, SRC_LOG_LEVELS, ENABLE_FORWARD_USER_INFO_HEADERS
|
||||
|
||||
from open_webui.utils.auth import get_admin_user, get_verified_user
|
||||
from open_webui.utils.images.comfyui import (
|
||||
ComfyUIGenerateImageForm,
|
||||
ComfyUIWorkflow,
|
||||
comfyui_generate_image,
|
||||
)
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["IMAGES"])
|
||||
|
||||
IMAGE_CACHE_DIR = Path(CACHE_DIR).joinpath("./image/generations/")
|
||||
IMAGE_CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("/config")
|
||||
async def get_config(request: Request, user=Depends(get_admin_user)):
|
||||
return {
|
||||
"enabled": request.app.state.config.ENABLE_IMAGE_GENERATION,
|
||||
"engine": request.app.state.config.IMAGE_GENERATION_ENGINE,
|
||||
"openai": {
|
||||
"OPENAI_API_BASE_URL": request.app.state.config.IMAGES_OPENAI_API_BASE_URL,
|
||||
"OPENAI_API_KEY": request.app.state.config.IMAGES_OPENAI_API_KEY,
|
||||
},
|
||||
"automatic1111": {
|
||||
"AUTOMATIC1111_BASE_URL": request.app.state.config.AUTOMATIC1111_BASE_URL,
|
||||
"AUTOMATIC1111_API_AUTH": request.app.state.config.AUTOMATIC1111_API_AUTH,
|
||||
"AUTOMATIC1111_CFG_SCALE": request.app.state.config.AUTOMATIC1111_CFG_SCALE,
|
||||
"AUTOMATIC1111_SAMPLER": request.app.state.config.AUTOMATIC1111_SAMPLER,
|
||||
"AUTOMATIC1111_SCHEDULER": request.app.state.config.AUTOMATIC1111_SCHEDULER,
|
||||
},
|
||||
"comfyui": {
|
||||
"COMFYUI_BASE_URL": request.app.state.config.COMFYUI_BASE_URL,
|
||||
"COMFYUI_WORKFLOW": request.app.state.config.COMFYUI_WORKFLOW,
|
||||
"COMFYUI_WORKFLOW_NODES": request.app.state.config.COMFYUI_WORKFLOW_NODES,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class OpenAIConfigForm(BaseModel):
|
||||
OPENAI_API_BASE_URL: str
|
||||
OPENAI_API_KEY: str
|
||||
|
||||
|
||||
class Automatic1111ConfigForm(BaseModel):
|
||||
AUTOMATIC1111_BASE_URL: str
|
||||
AUTOMATIC1111_API_AUTH: str
|
||||
AUTOMATIC1111_CFG_SCALE: Optional[str | float | int]
|
||||
AUTOMATIC1111_SAMPLER: Optional[str]
|
||||
AUTOMATIC1111_SCHEDULER: Optional[str]
|
||||
|
||||
|
||||
class ComfyUIConfigForm(BaseModel):
|
||||
COMFYUI_BASE_URL: str
|
||||
COMFYUI_WORKFLOW: str
|
||||
COMFYUI_WORKFLOW_NODES: list[dict]
|
||||
|
||||
|
||||
class ConfigForm(BaseModel):
|
||||
enabled: bool
|
||||
engine: str
|
||||
openai: OpenAIConfigForm
|
||||
automatic1111: Automatic1111ConfigForm
|
||||
comfyui: ComfyUIConfigForm
|
||||
|
||||
|
||||
@router.post("/config/update")
|
||||
async def update_config(
|
||||
request: Request, form_data: ConfigForm, user=Depends(get_admin_user)
|
||||
):
|
||||
request.app.state.config.IMAGE_GENERATION_ENGINE = form_data.engine
|
||||
request.app.state.config.ENABLE_IMAGE_GENERATION = form_data.enabled
|
||||
|
||||
request.app.state.config.IMAGES_OPENAI_API_BASE_URL = (
|
||||
form_data.openai.OPENAI_API_BASE_URL
|
||||
)
|
||||
request.app.state.config.IMAGES_OPENAI_API_KEY = form_data.openai.OPENAI_API_KEY
|
||||
|
||||
request.app.state.config.AUTOMATIC1111_BASE_URL = (
|
||||
form_data.automatic1111.AUTOMATIC1111_BASE_URL
|
||||
)
|
||||
request.app.state.config.AUTOMATIC1111_API_AUTH = (
|
||||
form_data.automatic1111.AUTOMATIC1111_API_AUTH
|
||||
)
|
||||
|
||||
request.app.state.config.AUTOMATIC1111_CFG_SCALE = (
|
||||
float(form_data.automatic1111.AUTOMATIC1111_CFG_SCALE)
|
||||
if form_data.automatic1111.AUTOMATIC1111_CFG_SCALE
|
||||
else None
|
||||
)
|
||||
request.app.state.config.AUTOMATIC1111_SAMPLER = (
|
||||
form_data.automatic1111.AUTOMATIC1111_SAMPLER
|
||||
if form_data.automatic1111.AUTOMATIC1111_SAMPLER
|
||||
else None
|
||||
)
|
||||
request.app.state.config.AUTOMATIC1111_SCHEDULER = (
|
||||
form_data.automatic1111.AUTOMATIC1111_SCHEDULER
|
||||
if form_data.automatic1111.AUTOMATIC1111_SCHEDULER
|
||||
else None
|
||||
)
|
||||
|
||||
request.app.state.config.COMFYUI_BASE_URL = (
|
||||
form_data.comfyui.COMFYUI_BASE_URL.strip("/")
|
||||
)
|
||||
request.app.state.config.COMFYUI_WORKFLOW = form_data.comfyui.COMFYUI_WORKFLOW
|
||||
request.app.state.config.COMFYUI_WORKFLOW_NODES = (
|
||||
form_data.comfyui.COMFYUI_WORKFLOW_NODES
|
||||
)
|
||||
|
||||
return {
|
||||
"enabled": request.app.state.config.ENABLE_IMAGE_GENERATION,
|
||||
"engine": request.app.state.config.IMAGE_GENERATION_ENGINE,
|
||||
"openai": {
|
||||
"OPENAI_API_BASE_URL": request.app.state.config.IMAGES_OPENAI_API_BASE_URL,
|
||||
"OPENAI_API_KEY": request.app.state.config.IMAGES_OPENAI_API_KEY,
|
||||
},
|
||||
"automatic1111": {
|
||||
"AUTOMATIC1111_BASE_URL": request.app.state.config.AUTOMATIC1111_BASE_URL,
|
||||
"AUTOMATIC1111_API_AUTH": request.app.state.config.AUTOMATIC1111_API_AUTH,
|
||||
"AUTOMATIC1111_CFG_SCALE": request.app.state.config.AUTOMATIC1111_CFG_SCALE,
|
||||
"AUTOMATIC1111_SAMPLER": request.app.state.config.AUTOMATIC1111_SAMPLER,
|
||||
"AUTOMATIC1111_SCHEDULER": request.app.state.config.AUTOMATIC1111_SCHEDULER,
|
||||
},
|
||||
"comfyui": {
|
||||
"COMFYUI_BASE_URL": request.app.state.config.COMFYUI_BASE_URL,
|
||||
"COMFYUI_WORKFLOW": request.app.state.config.COMFYUI_WORKFLOW,
|
||||
"COMFYUI_WORKFLOW_NODES": request.app.state.config.COMFYUI_WORKFLOW_NODES,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def get_automatic1111_api_auth(request: Request):
|
||||
if request.app.state.config.AUTOMATIC1111_API_AUTH is None:
|
||||
return ""
|
||||
else:
|
||||
auth1111_byte_string = request.app.state.config.AUTOMATIC1111_API_AUTH.encode(
|
||||
"utf-8"
|
||||
)
|
||||
auth1111_base64_encoded_bytes = base64.b64encode(auth1111_byte_string)
|
||||
auth1111_base64_encoded_string = auth1111_base64_encoded_bytes.decode("utf-8")
|
||||
return f"Basic {auth1111_base64_encoded_string}"
|
||||
|
||||
|
||||
@router.get("/config/url/verify")
|
||||
async def verify_url(request: Request, user=Depends(get_admin_user)):
|
||||
if request.app.state.config.IMAGE_GENERATION_ENGINE == "automatic1111":
|
||||
try:
|
||||
r = requests.get(
|
||||
url=f"{request.app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/options",
|
||||
headers={"authorization": get_automatic1111_api_auth(request)},
|
||||
)
|
||||
r.raise_for_status()
|
||||
return True
|
||||
except Exception:
|
||||
request.app.state.config.ENABLE_IMAGE_GENERATION = False
|
||||
raise HTTPException(status_code=400, detail=ERROR_MESSAGES.INVALID_URL)
|
||||
elif request.app.state.config.IMAGE_GENERATION_ENGINE == "comfyui":
|
||||
try:
|
||||
r = requests.get(
|
||||
url=f"{request.app.state.config.COMFYUI_BASE_URL}/object_info"
|
||||
)
|
||||
r.raise_for_status()
|
||||
return True
|
||||
except Exception:
|
||||
request.app.state.config.ENABLE_IMAGE_GENERATION = False
|
||||
raise HTTPException(status_code=400, detail=ERROR_MESSAGES.INVALID_URL)
|
||||
else:
|
||||
return True
|
||||
|
||||
|
||||
def set_image_model(request: Request, model: str):
|
||||
log.info(f"Setting image model to {model}")
|
||||
request.app.state.config.IMAGE_GENERATION_MODEL = model
|
||||
if request.app.state.config.IMAGE_GENERATION_ENGINE in ["", "automatic1111"]:
|
||||
api_auth = get_automatic1111_api_auth()
|
||||
r = requests.get(
|
||||
url=f"{request.app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/options",
|
||||
headers={"authorization": api_auth},
|
||||
)
|
||||
options = r.json()
|
||||
if model != options["sd_model_checkpoint"]:
|
||||
options["sd_model_checkpoint"] = model
|
||||
r = requests.post(
|
||||
url=f"{request.app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/options",
|
||||
json=options,
|
||||
headers={"authorization": api_auth},
|
||||
)
|
||||
return request.app.state.config.IMAGE_GENERATION_MODEL
|
||||
|
||||
|
||||
def get_image_model(request):
|
||||
if request.app.state.config.IMAGE_GENERATION_ENGINE == "openai":
|
||||
return (
|
||||
request.app.state.config.IMAGE_GENERATION_MODEL
|
||||
if request.app.state.config.IMAGE_GENERATION_MODEL
|
||||
else "dall-e-2"
|
||||
)
|
||||
elif request.app.state.config.IMAGE_GENERATION_ENGINE == "comfyui":
|
||||
return (
|
||||
request.app.state.config.IMAGE_GENERATION_MODEL
|
||||
if request.app.state.config.IMAGE_GENERATION_MODEL
|
||||
else ""
|
||||
)
|
||||
elif (
|
||||
request.app.state.config.IMAGE_GENERATION_ENGINE == "automatic1111"
|
||||
or request.app.state.config.IMAGE_GENERATION_ENGINE == ""
|
||||
):
|
||||
try:
|
||||
r = requests.get(
|
||||
url=f"{request.app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/options",
|
||||
headers={"authorization": get_automatic1111_api_auth()},
|
||||
)
|
||||
options = r.json()
|
||||
return options["sd_model_checkpoint"]
|
||||
except Exception as e:
|
||||
request.app.state.config.ENABLE_IMAGE_GENERATION = False
|
||||
raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))
|
||||
|
||||
|
||||
class ImageConfigForm(BaseModel):
|
||||
MODEL: str
|
||||
IMAGE_SIZE: str
|
||||
IMAGE_STEPS: int
|
||||
|
||||
|
||||
@router.get("/image/config")
|
||||
async def get_image_config(request: Request, user=Depends(get_admin_user)):
|
||||
return {
|
||||
"MODEL": request.app.state.config.IMAGE_GENERATION_MODEL,
|
||||
"IMAGE_SIZE": request.app.state.config.IMAGE_SIZE,
|
||||
"IMAGE_STEPS": request.app.state.config.IMAGE_STEPS,
|
||||
}
|
||||
|
||||
|
||||
@router.post("/image/config/update")
|
||||
async def update_image_config(
|
||||
request: Request, form_data: ImageConfigForm, user=Depends(get_admin_user)
|
||||
):
|
||||
|
||||
set_image_model(request, form_data.MODEL)
|
||||
|
||||
pattern = r"^\d+x\d+$"
|
||||
if re.match(pattern, form_data.IMAGE_SIZE):
|
||||
request.app.state.config.IMAGE_SIZE = form_data.IMAGE_SIZE
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=ERROR_MESSAGES.INCORRECT_FORMAT(" (e.g., 512x512)."),
|
||||
)
|
||||
|
||||
if form_data.IMAGE_STEPS >= 0:
|
||||
request.app.state.config.IMAGE_STEPS = form_data.IMAGE_STEPS
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=ERROR_MESSAGES.INCORRECT_FORMAT(" (e.g., 50)."),
|
||||
)
|
||||
|
||||
return {
|
||||
"MODEL": request.app.state.config.IMAGE_GENERATION_MODEL,
|
||||
"IMAGE_SIZE": request.app.state.config.IMAGE_SIZE,
|
||||
"IMAGE_STEPS": request.app.state.config.IMAGE_STEPS,
|
||||
}
|
||||
|
||||
|
||||
@router.get("/models")
|
||||
def get_models(request: Request, user=Depends(get_verified_user)):
|
||||
try:
|
||||
if request.app.state.config.IMAGE_GENERATION_ENGINE == "openai":
|
||||
return [
|
||||
{"id": "dall-e-2", "name": "DALL·E 2"},
|
||||
{"id": "dall-e-3", "name": "DALL·E 3"},
|
||||
]
|
||||
elif request.app.state.config.IMAGE_GENERATION_ENGINE == "comfyui":
|
||||
# TODO - get models from comfyui
|
||||
r = requests.get(
|
||||
url=f"{request.app.state.config.COMFYUI_BASE_URL}/object_info"
|
||||
)
|
||||
info = r.json()
|
||||
|
||||
workflow = json.loads(request.app.state.config.COMFYUI_WORKFLOW)
|
||||
model_node_id = None
|
||||
|
||||
for node in request.app.state.config.COMFYUI_WORKFLOW_NODES:
|
||||
if node["type"] == "model":
|
||||
if node["node_ids"]:
|
||||
model_node_id = node["node_ids"][0]
|
||||
break
|
||||
|
||||
if model_node_id:
|
||||
model_list_key = None
|
||||
|
||||
print(workflow[model_node_id]["class_type"])
|
||||
for key in info[workflow[model_node_id]["class_type"]]["input"][
|
||||
"required"
|
||||
]:
|
||||
if "_name" in key:
|
||||
model_list_key = key
|
||||
break
|
||||
|
||||
if model_list_key:
|
||||
return list(
|
||||
map(
|
||||
lambda model: {"id": model, "name": model},
|
||||
info[workflow[model_node_id]["class_type"]]["input"][
|
||||
"required"
|
||||
][model_list_key][0],
|
||||
)
|
||||
)
|
||||
else:
|
||||
return list(
|
||||
map(
|
||||
lambda model: {"id": model, "name": model},
|
||||
info["CheckpointLoaderSimple"]["input"]["required"][
|
||||
"ckpt_name"
|
||||
][0],
|
||||
)
|
||||
)
|
||||
elif (
|
||||
request.app.state.config.IMAGE_GENERATION_ENGINE == "automatic1111"
|
||||
or request.app.state.config.IMAGE_GENERATION_ENGINE == ""
|
||||
):
|
||||
r = requests.get(
|
||||
url=f"{request.app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/sd-models",
|
||||
headers={"authorization": get_automatic1111_api_auth()},
|
||||
)
|
||||
models = r.json()
|
||||
return list(
|
||||
map(
|
||||
lambda model: {"id": model["title"], "name": model["model_name"]},
|
||||
models,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
request.app.state.config.ENABLE_IMAGE_GENERATION = False
|
||||
raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))
|
||||
|
||||
|
||||
class GenerateImageForm(BaseModel):
|
||||
model: Optional[str] = None
|
||||
prompt: str
|
||||
size: Optional[str] = None
|
||||
n: int = 1
|
||||
negative_prompt: Optional[str] = None
|
||||
|
||||
|
||||
def save_b64_image(b64_str):
|
||||
try:
|
||||
image_id = str(uuid.uuid4())
|
||||
|
||||
if "," in b64_str:
|
||||
header, encoded = b64_str.split(",", 1)
|
||||
mime_type = header.split(";")[0]
|
||||
|
||||
img_data = base64.b64decode(encoded)
|
||||
image_format = mimetypes.guess_extension(mime_type)
|
||||
|
||||
image_filename = f"{image_id}{image_format}"
|
||||
file_path = IMAGE_CACHE_DIR / f"{image_filename}"
|
||||
with open(file_path, "wb") as f:
|
||||
f.write(img_data)
|
||||
return image_filename
|
||||
else:
|
||||
image_filename = f"{image_id}.png"
|
||||
file_path = IMAGE_CACHE_DIR.joinpath(image_filename)
|
||||
|
||||
img_data = base64.b64decode(b64_str)
|
||||
|
||||
# Write the image data to a file
|
||||
with open(file_path, "wb") as f:
|
||||
f.write(img_data)
|
||||
return image_filename
|
||||
|
||||
except Exception as e:
|
||||
log.exception(f"Error saving image: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def save_url_image(url):
|
||||
image_id = str(uuid.uuid4())
|
||||
try:
|
||||
r = requests.get(url)
|
||||
r.raise_for_status()
|
||||
if r.headers["content-type"].split("/")[0] == "image":
|
||||
mime_type = r.headers["content-type"]
|
||||
image_format = mimetypes.guess_extension(mime_type)
|
||||
|
||||
if not image_format:
|
||||
raise ValueError("Could not determine image type from MIME type")
|
||||
|
||||
image_filename = f"{image_id}{image_format}"
|
||||
|
||||
file_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}")
|
||||
with open(file_path, "wb") as image_file:
|
||||
for chunk in r.iter_content(chunk_size=8192):
|
||||
image_file.write(chunk)
|
||||
return image_filename
|
||||
else:
|
||||
log.error("Url does not point to an image.")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
log.exception(f"Error saving image: {e}")
|
||||
return None
|
||||
|
||||
|
||||
@router.post("/generations")
|
||||
async def image_generations(
|
||||
request: Request,
|
||||
form_data: GenerateImageForm,
|
||||
user=Depends(get_verified_user),
|
||||
):
|
||||
width, height = tuple(map(int, request.app.state.config.IMAGE_SIZE.split("x")))
|
||||
|
||||
r = None
|
||||
try:
|
||||
if request.app.state.config.IMAGE_GENERATION_ENGINE == "openai":
|
||||
headers = {}
|
||||
headers["Authorization"] = (
|
||||
f"Bearer {request.app.state.config.IMAGES_OPENAI_API_KEY}"
|
||||
)
|
||||
headers["Content-Type"] = "application/json"
|
||||
|
||||
if ENABLE_FORWARD_USER_INFO_HEADERS:
|
||||
headers["X-OpenWebUI-User-Name"] = user.name
|
||||
headers["X-OpenWebUI-User-Id"] = user.id
|
||||
headers["X-OpenWebUI-User-Email"] = user.email
|
||||
headers["X-OpenWebUI-User-Role"] = user.role
|
||||
|
||||
data = {
|
||||
"model": (
|
||||
request.app.state.config.IMAGE_GENERATION_MODEL
|
||||
if request.app.state.config.IMAGE_GENERATION_MODEL != ""
|
||||
else "dall-e-2"
|
||||
),
|
||||
"prompt": form_data.prompt,
|
||||
"n": form_data.n,
|
||||
"size": (
|
||||
form_data.size
|
||||
if form_data.size
|
||||
else request.app.state.config.IMAGE_SIZE
|
||||
),
|
||||
"response_format": "b64_json",
|
||||
}
|
||||
|
||||
# Use asyncio.to_thread for the requests.post call
|
||||
r = await asyncio.to_thread(
|
||||
requests.post,
|
||||
url=f"{request.app.state.config.IMAGES_OPENAI_API_BASE_URL}/images/generations",
|
||||
json=data,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
r.raise_for_status()
|
||||
res = r.json()
|
||||
|
||||
images = []
|
||||
|
||||
for image in res["data"]:
|
||||
image_filename = save_b64_image(image["b64_json"])
|
||||
images.append({"url": f"/cache/image/generations/{image_filename}"})
|
||||
file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}.json")
|
||||
|
||||
with open(file_body_path, "w") as f:
|
||||
json.dump(data, f)
|
||||
|
||||
return images
|
||||
|
||||
elif request.app.state.config.IMAGE_GENERATION_ENGINE == "comfyui":
|
||||
data = {
|
||||
"prompt": form_data.prompt,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"n": form_data.n,
|
||||
}
|
||||
|
||||
if request.app.state.config.IMAGE_STEPS is not None:
|
||||
data["steps"] = request.app.state.config.IMAGE_STEPS
|
||||
|
||||
if form_data.negative_prompt is not None:
|
||||
data["negative_prompt"] = form_data.negative_prompt
|
||||
|
||||
form_data = ComfyUIGenerateImageForm(
|
||||
**{
|
||||
"workflow": ComfyUIWorkflow(
|
||||
**{
|
||||
"workflow": request.app.state.config.COMFYUI_WORKFLOW,
|
||||
"nodes": request.app.state.config.COMFYUI_WORKFLOW_NODES,
|
||||
}
|
||||
),
|
||||
**data,
|
||||
}
|
||||
)
|
||||
res = await comfyui_generate_image(
|
||||
request.app.state.config.IMAGE_GENERATION_MODEL,
|
||||
form_data,
|
||||
user.id,
|
||||
request.app.state.config.COMFYUI_BASE_URL,
|
||||
)
|
||||
log.debug(f"res: {res}")
|
||||
|
||||
images = []
|
||||
|
||||
for image in res["data"]:
|
||||
image_filename = save_url_image(image["url"])
|
||||
images.append({"url": f"/cache/image/generations/{image_filename}"})
|
||||
file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}.json")
|
||||
|
||||
with open(file_body_path, "w") as f:
|
||||
json.dump(form_data.model_dump(exclude_none=True), f)
|
||||
|
||||
log.debug(f"images: {images}")
|
||||
return images
|
||||
elif (
|
||||
request.app.state.config.IMAGE_GENERATION_ENGINE == "automatic1111"
|
||||
or request.app.state.config.IMAGE_GENERATION_ENGINE == ""
|
||||
):
|
||||
if form_data.model:
|
||||
set_image_model(form_data.model)
|
||||
|
||||
data = {
|
||||
"prompt": form_data.prompt,
|
||||
"batch_size": form_data.n,
|
||||
"width": width,
|
||||
"height": height,
|
||||
}
|
||||
|
||||
if request.app.state.config.IMAGE_STEPS is not None:
|
||||
data["steps"] = request.app.state.config.IMAGE_STEPS
|
||||
|
||||
if form_data.negative_prompt is not None:
|
||||
data["negative_prompt"] = form_data.negative_prompt
|
||||
|
||||
if request.app.state.config.AUTOMATIC1111_CFG_SCALE:
|
||||
data["cfg_scale"] = request.app.state.config.AUTOMATIC1111_CFG_SCALE
|
||||
|
||||
if request.app.state.config.AUTOMATIC1111_SAMPLER:
|
||||
data["sampler_name"] = request.app.state.config.AUTOMATIC1111_SAMPLER
|
||||
|
||||
if request.app.state.config.AUTOMATIC1111_SCHEDULER:
|
||||
data["scheduler"] = request.app.state.config.AUTOMATIC1111_SCHEDULER
|
||||
|
||||
# Use asyncio.to_thread for the requests.post call
|
||||
r = await asyncio.to_thread(
|
||||
requests.post,
|
||||
url=f"{request.app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/txt2img",
|
||||
json=data,
|
||||
headers={"authorization": get_automatic1111_api_auth()},
|
||||
)
|
||||
|
||||
res = r.json()
|
||||
log.debug(f"res: {res}")
|
||||
|
||||
images = []
|
||||
|
||||
for image in res["images"]:
|
||||
image_filename = save_b64_image(image)
|
||||
images.append({"url": f"/cache/image/generations/{image_filename}"})
|
||||
file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}.json")
|
||||
|
||||
with open(file_body_path, "w") as f:
|
||||
json.dump({**data, "info": res["info"]}, f)
|
||||
|
||||
return images
|
||||
except Exception as e:
|
||||
error = e
|
||||
if r != None:
|
||||
data = r.json()
|
||||
if "error" in data:
|
||||
error = data["error"]["message"]
|
||||
raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(error))
|
||||
@@ -0,0 +1,597 @@
|
||||
from typing import List, Optional
|
||||
from pydantic import BaseModel
|
||||
from fastapi import APIRouter, Depends, HTTPException, status, Request
|
||||
import logging
|
||||
|
||||
from open_webui.models.knowledge import (
|
||||
Knowledges,
|
||||
KnowledgeForm,
|
||||
KnowledgeResponse,
|
||||
KnowledgeUserResponse,
|
||||
)
|
||||
from open_webui.models.files import Files, FileModel
|
||||
from open_webui.retrieval.vector.connector import VECTOR_DB_CLIENT
|
||||
from open_webui.routers.retrieval import process_file, ProcessFileForm, process_files_batch, BatchProcessFilesForm
|
||||
|
||||
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
from open_webui.utils.auth import get_verified_user
|
||||
from open_webui.utils.access_control import has_access, has_permission
|
||||
|
||||
|
||||
from open_webui.env import SRC_LOG_LEVELS
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["MODELS"])
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
############################
|
||||
# getKnowledgeBases
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/", response_model=list[KnowledgeUserResponse])
|
||||
async def get_knowledge(user=Depends(get_verified_user)):
|
||||
knowledge_bases = []
|
||||
|
||||
if user.role == "admin":
|
||||
knowledge_bases = Knowledges.get_knowledge_bases()
|
||||
else:
|
||||
knowledge_bases = Knowledges.get_knowledge_bases_by_user_id(user.id, "read")
|
||||
|
||||
# Get files for each knowledge base
|
||||
knowledge_with_files = []
|
||||
for knowledge_base in knowledge_bases:
|
||||
files = []
|
||||
if knowledge_base.data:
|
||||
files = Files.get_file_metadatas_by_ids(
|
||||
knowledge_base.data.get("file_ids", [])
|
||||
)
|
||||
|
||||
# Check if all files exist
|
||||
if len(files) != len(knowledge_base.data.get("file_ids", [])):
|
||||
missing_files = list(
|
||||
set(knowledge_base.data.get("file_ids", []))
|
||||
- set([file.id for file in files])
|
||||
)
|
||||
if missing_files:
|
||||
data = knowledge_base.data or {}
|
||||
file_ids = data.get("file_ids", [])
|
||||
|
||||
for missing_file in missing_files:
|
||||
file_ids.remove(missing_file)
|
||||
|
||||
data["file_ids"] = file_ids
|
||||
Knowledges.update_knowledge_data_by_id(
|
||||
id=knowledge_base.id, data=data
|
||||
)
|
||||
|
||||
files = Files.get_file_metadatas_by_ids(file_ids)
|
||||
|
||||
knowledge_with_files.append(
|
||||
KnowledgeUserResponse(
|
||||
**knowledge_base.model_dump(),
|
||||
files=files,
|
||||
)
|
||||
)
|
||||
|
||||
return knowledge_with_files
|
||||
|
||||
|
||||
@router.get("/list", response_model=list[KnowledgeUserResponse])
|
||||
async def get_knowledge_list(user=Depends(get_verified_user)):
|
||||
knowledge_bases = []
|
||||
|
||||
if user.role == "admin":
|
||||
knowledge_bases = Knowledges.get_knowledge_bases()
|
||||
else:
|
||||
knowledge_bases = Knowledges.get_knowledge_bases_by_user_id(user.id, "write")
|
||||
|
||||
# Get files for each knowledge base
|
||||
knowledge_with_files = []
|
||||
for knowledge_base in knowledge_bases:
|
||||
files = []
|
||||
if knowledge_base.data:
|
||||
files = Files.get_file_metadatas_by_ids(
|
||||
knowledge_base.data.get("file_ids", [])
|
||||
)
|
||||
|
||||
# Check if all files exist
|
||||
if len(files) != len(knowledge_base.data.get("file_ids", [])):
|
||||
missing_files = list(
|
||||
set(knowledge_base.data.get("file_ids", []))
|
||||
- set([file.id for file in files])
|
||||
)
|
||||
if missing_files:
|
||||
data = knowledge_base.data or {}
|
||||
file_ids = data.get("file_ids", [])
|
||||
|
||||
for missing_file in missing_files:
|
||||
file_ids.remove(missing_file)
|
||||
|
||||
data["file_ids"] = file_ids
|
||||
Knowledges.update_knowledge_data_by_id(
|
||||
id=knowledge_base.id, data=data
|
||||
)
|
||||
|
||||
files = Files.get_file_metadatas_by_ids(file_ids)
|
||||
|
||||
knowledge_with_files.append(
|
||||
KnowledgeUserResponse(
|
||||
**knowledge_base.model_dump(),
|
||||
files=files,
|
||||
)
|
||||
)
|
||||
return knowledge_with_files
|
||||
|
||||
|
||||
############################
|
||||
# CreateNewKnowledge
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/create", response_model=Optional[KnowledgeResponse])
|
||||
async def create_new_knowledge(
|
||||
request: Request, form_data: KnowledgeForm, user=Depends(get_verified_user)
|
||||
):
|
||||
if user.role != "admin" and not has_permission(
|
||||
user.id, "workspace.knowledge", request.app.state.config.USER_PERMISSIONS
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.UNAUTHORIZED,
|
||||
)
|
||||
|
||||
knowledge = Knowledges.insert_new_knowledge(user.id, form_data)
|
||||
|
||||
if knowledge:
|
||||
return knowledge
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.FILE_EXISTS,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# GetKnowledgeById
|
||||
############################
|
||||
|
||||
|
||||
class KnowledgeFilesResponse(KnowledgeResponse):
|
||||
files: list[FileModel]
|
||||
|
||||
|
||||
@router.get("/{id}", response_model=Optional[KnowledgeFilesResponse])
|
||||
async def get_knowledge_by_id(id: str, user=Depends(get_verified_user)):
|
||||
knowledge = Knowledges.get_knowledge_by_id(id=id)
|
||||
|
||||
if knowledge:
|
||||
|
||||
if (
|
||||
user.role == "admin"
|
||||
or knowledge.user_id == user.id
|
||||
or has_access(user.id, "read", knowledge.access_control)
|
||||
):
|
||||
|
||||
file_ids = knowledge.data.get("file_ids", []) if knowledge.data else []
|
||||
files = Files.get_files_by_ids(file_ids)
|
||||
|
||||
return KnowledgeFilesResponse(
|
||||
**knowledge.model_dump(),
|
||||
files=files,
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# UpdateKnowledgeById
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/{id}/update", response_model=Optional[KnowledgeFilesResponse])
|
||||
async def update_knowledge_by_id(
|
||||
id: str,
|
||||
form_data: KnowledgeForm,
|
||||
user=Depends(get_verified_user),
|
||||
):
|
||||
knowledge = Knowledges.get_knowledge_by_id(id=id)
|
||||
if not knowledge:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
if knowledge.user_id != user.id and user.role != "admin":
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
|
||||
knowledge = Knowledges.update_knowledge_by_id(id=id, form_data=form_data)
|
||||
if knowledge:
|
||||
file_ids = knowledge.data.get("file_ids", []) if knowledge.data else []
|
||||
files = Files.get_files_by_ids(file_ids)
|
||||
|
||||
return KnowledgeFilesResponse(
|
||||
**knowledge.model_dump(),
|
||||
files=files,
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.ID_TAKEN,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# AddFileToKnowledge
|
||||
############################
|
||||
|
||||
|
||||
class KnowledgeFileIdForm(BaseModel):
|
||||
file_id: str
|
||||
|
||||
|
||||
@router.post("/{id}/file/add", response_model=Optional[KnowledgeFilesResponse])
|
||||
def add_file_to_knowledge_by_id(
|
||||
request: Request,
|
||||
id: str,
|
||||
form_data: KnowledgeFileIdForm,
|
||||
user=Depends(get_verified_user),
|
||||
):
|
||||
knowledge = Knowledges.get_knowledge_by_id(id=id)
|
||||
|
||||
if not knowledge:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
if knowledge.user_id != user.id and user.role != "admin":
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
|
||||
file = Files.get_file_by_id(form_data.file_id)
|
||||
if not file:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
if not file.data:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.FILE_NOT_PROCESSED,
|
||||
)
|
||||
|
||||
# Add content to the vector database
|
||||
try:
|
||||
process_file(
|
||||
request, ProcessFileForm(file_id=form_data.file_id, collection_name=id)
|
||||
)
|
||||
except Exception as e:
|
||||
log.debug(e)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=str(e),
|
||||
)
|
||||
|
||||
if knowledge:
|
||||
data = knowledge.data or {}
|
||||
file_ids = data.get("file_ids", [])
|
||||
|
||||
if form_data.file_id not in file_ids:
|
||||
file_ids.append(form_data.file_id)
|
||||
data["file_ids"] = file_ids
|
||||
|
||||
knowledge = Knowledges.update_knowledge_data_by_id(id=id, data=data)
|
||||
|
||||
if knowledge:
|
||||
files = Files.get_files_by_ids(file_ids)
|
||||
|
||||
return KnowledgeFilesResponse(
|
||||
**knowledge.model_dump(),
|
||||
files=files,
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("knowledge"),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("file_id"),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
@router.post("/{id}/file/update", response_model=Optional[KnowledgeFilesResponse])
|
||||
def update_file_from_knowledge_by_id(
|
||||
request: Request,
|
||||
id: str,
|
||||
form_data: KnowledgeFileIdForm,
|
||||
user=Depends(get_verified_user),
|
||||
):
|
||||
knowledge = Knowledges.get_knowledge_by_id(id=id)
|
||||
if not knowledge:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
if knowledge.user_id != user.id and user.role != "admin":
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
|
||||
file = Files.get_file_by_id(form_data.file_id)
|
||||
if not file:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
# Remove content from the vector database
|
||||
VECTOR_DB_CLIENT.delete(
|
||||
collection_name=knowledge.id, filter={"file_id": form_data.file_id}
|
||||
)
|
||||
|
||||
# Add content to the vector database
|
||||
try:
|
||||
process_file(
|
||||
request, ProcessFileForm(file_id=form_data.file_id, collection_name=id)
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=str(e),
|
||||
)
|
||||
|
||||
if knowledge:
|
||||
data = knowledge.data or {}
|
||||
file_ids = data.get("file_ids", [])
|
||||
|
||||
files = Files.get_files_by_ids(file_ids)
|
||||
|
||||
return KnowledgeFilesResponse(
|
||||
**knowledge.model_dump(),
|
||||
files=files,
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# RemoveFileFromKnowledge
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/{id}/file/remove", response_model=Optional[KnowledgeFilesResponse])
|
||||
def remove_file_from_knowledge_by_id(
|
||||
id: str,
|
||||
form_data: KnowledgeFileIdForm,
|
||||
user=Depends(get_verified_user),
|
||||
):
|
||||
knowledge = Knowledges.get_knowledge_by_id(id=id)
|
||||
if not knowledge:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
if knowledge.user_id != user.id and user.role != "admin":
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
|
||||
file = Files.get_file_by_id(form_data.file_id)
|
||||
if not file:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
# Remove content from the vector database
|
||||
VECTOR_DB_CLIENT.delete(
|
||||
collection_name=knowledge.id, filter={"file_id": form_data.file_id}
|
||||
)
|
||||
|
||||
result = VECTOR_DB_CLIENT.query(
|
||||
collection_name=knowledge.id,
|
||||
filter={"file_id": form_data.file_id},
|
||||
)
|
||||
|
||||
Files.delete_file_by_id(form_data.file_id)
|
||||
|
||||
if knowledge:
|
||||
data = knowledge.data or {}
|
||||
file_ids = data.get("file_ids", [])
|
||||
|
||||
if form_data.file_id in file_ids:
|
||||
file_ids.remove(form_data.file_id)
|
||||
data["file_ids"] = file_ids
|
||||
|
||||
knowledge = Knowledges.update_knowledge_data_by_id(id=id, data=data)
|
||||
|
||||
if knowledge:
|
||||
files = Files.get_files_by_ids(file_ids)
|
||||
|
||||
return KnowledgeFilesResponse(
|
||||
**knowledge.model_dump(),
|
||||
files=files,
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("knowledge"),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("file_id"),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# DeleteKnowledgeById
|
||||
############################
|
||||
|
||||
|
||||
@router.delete("/{id}/delete", response_model=bool)
|
||||
async def delete_knowledge_by_id(id: str, user=Depends(get_verified_user)):
|
||||
knowledge = Knowledges.get_knowledge_by_id(id=id)
|
||||
if not knowledge:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
if knowledge.user_id != user.id and user.role != "admin":
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
|
||||
try:
|
||||
VECTOR_DB_CLIENT.delete_collection(collection_name=id)
|
||||
except Exception as e:
|
||||
log.debug(e)
|
||||
pass
|
||||
result = Knowledges.delete_knowledge_by_id(id=id)
|
||||
return result
|
||||
|
||||
|
||||
############################
|
||||
# ResetKnowledgeById
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/{id}/reset", response_model=Optional[KnowledgeResponse])
|
||||
async def reset_knowledge_by_id(id: str, user=Depends(get_verified_user)):
|
||||
knowledge = Knowledges.get_knowledge_by_id(id=id)
|
||||
if not knowledge:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
if knowledge.user_id != user.id and user.role != "admin":
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
|
||||
try:
|
||||
VECTOR_DB_CLIENT.delete_collection(collection_name=id)
|
||||
except Exception as e:
|
||||
log.debug(e)
|
||||
pass
|
||||
|
||||
knowledge = Knowledges.update_knowledge_data_by_id(id=id, data={"file_ids": []})
|
||||
|
||||
return knowledge
|
||||
|
||||
|
||||
############################
|
||||
# AddFilesToKnowledge
|
||||
############################
|
||||
|
||||
@router.post("/{id}/files/batch/add", response_model=Optional[KnowledgeFilesResponse])
|
||||
def add_files_to_knowledge_batch(
|
||||
id: str,
|
||||
form_data: list[KnowledgeFileIdForm],
|
||||
user=Depends(get_verified_user),
|
||||
):
|
||||
"""
|
||||
Add multiple files to a knowledge base
|
||||
"""
|
||||
knowledge = Knowledges.get_knowledge_by_id(id=id)
|
||||
if not knowledge:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
if knowledge.user_id != user.id and user.role != "admin":
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
|
||||
# Get files content
|
||||
print(f"files/batch/add - {len(form_data)} files")
|
||||
files: List[FileModel] = []
|
||||
for form in form_data:
|
||||
file = Files.get_file_by_id(form.file_id)
|
||||
if not file:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=f"File {form.file_id} not found",
|
||||
)
|
||||
files.append(file)
|
||||
|
||||
# Process files
|
||||
try:
|
||||
result = process_files_batch(BatchProcessFilesForm(
|
||||
files=files,
|
||||
collection_name=id
|
||||
))
|
||||
except Exception as e:
|
||||
log.error(f"add_files_to_knowledge_batch: Exception occurred: {e}", exc_info=True)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=str(e)
|
||||
)
|
||||
|
||||
# Add successful files to knowledge base
|
||||
data = knowledge.data or {}
|
||||
existing_file_ids = data.get("file_ids", [])
|
||||
|
||||
# Only add files that were successfully processed
|
||||
successful_file_ids = [r.file_id for r in result.results if r.status == "completed"]
|
||||
for file_id in successful_file_ids:
|
||||
if file_id not in existing_file_ids:
|
||||
existing_file_ids.append(file_id)
|
||||
|
||||
data["file_ids"] = existing_file_ids
|
||||
knowledge = Knowledges.update_knowledge_data_by_id(id=id, data=data)
|
||||
|
||||
# If there were any errors, include them in the response
|
||||
if result.errors:
|
||||
error_details = [f"{err.file_id}: {err.error}" for err in result.errors]
|
||||
return KnowledgeFilesResponse(
|
||||
**knowledge.model_dump(),
|
||||
files=Files.get_files_by_ids(existing_file_ids),
|
||||
warnings={
|
||||
"message": "Some files failed to process",
|
||||
"errors": error_details
|
||||
}
|
||||
)
|
||||
|
||||
return KnowledgeFilesResponse(
|
||||
**knowledge.model_dump(),
|
||||
files=Files.get_files_by_ids(existing_file_ids)
|
||||
)
|
||||
@@ -0,0 +1,190 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request
|
||||
from pydantic import BaseModel
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from open_webui.models.memories import Memories, MemoryModel
|
||||
from open_webui.retrieval.vector.connector import VECTOR_DB_CLIENT
|
||||
from open_webui.utils.auth import get_verified_user
|
||||
from open_webui.env import SRC_LOG_LEVELS
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["MODELS"])
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("/ef")
|
||||
async def get_embeddings(request: Request):
|
||||
return {"result": request.app.state.EMBEDDING_FUNCTION("hello world")}
|
||||
|
||||
|
||||
############################
|
||||
# GetMemories
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/", response_model=list[MemoryModel])
|
||||
async def get_memories(user=Depends(get_verified_user)):
|
||||
return Memories.get_memories_by_user_id(user.id)
|
||||
|
||||
|
||||
############################
|
||||
# AddMemory
|
||||
############################
|
||||
|
||||
|
||||
class AddMemoryForm(BaseModel):
|
||||
content: str
|
||||
|
||||
|
||||
class MemoryUpdateModel(BaseModel):
|
||||
content: Optional[str] = None
|
||||
|
||||
|
||||
@router.post("/add", response_model=Optional[MemoryModel])
|
||||
async def add_memory(
|
||||
request: Request,
|
||||
form_data: AddMemoryForm,
|
||||
user=Depends(get_verified_user),
|
||||
):
|
||||
memory = Memories.insert_new_memory(user.id, form_data.content)
|
||||
|
||||
VECTOR_DB_CLIENT.upsert(
|
||||
collection_name=f"user-memory-{user.id}",
|
||||
items=[
|
||||
{
|
||||
"id": memory.id,
|
||||
"text": memory.content,
|
||||
"vector": request.app.state.EMBEDDING_FUNCTION(memory.content),
|
||||
"metadata": {"created_at": memory.created_at},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
return memory
|
||||
|
||||
|
||||
############################
|
||||
# QueryMemory
|
||||
############################
|
||||
|
||||
|
||||
class QueryMemoryForm(BaseModel):
|
||||
content: str
|
||||
k: Optional[int] = 1
|
||||
|
||||
|
||||
@router.post("/query")
|
||||
async def query_memory(
|
||||
request: Request, form_data: QueryMemoryForm, user=Depends(get_verified_user)
|
||||
):
|
||||
results = VECTOR_DB_CLIENT.search(
|
||||
collection_name=f"user-memory-{user.id}",
|
||||
vectors=[request.app.state.EMBEDDING_FUNCTION(form_data.content)],
|
||||
limit=form_data.k,
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
############################
|
||||
# ResetMemoryFromVectorDB
|
||||
############################
|
||||
@router.post("/reset", response_model=bool)
|
||||
async def reset_memory_from_vector_db(
|
||||
request: Request, user=Depends(get_verified_user)
|
||||
):
|
||||
VECTOR_DB_CLIENT.delete_collection(f"user-memory-{user.id}")
|
||||
|
||||
memories = Memories.get_memories_by_user_id(user.id)
|
||||
VECTOR_DB_CLIENT.upsert(
|
||||
collection_name=f"user-memory-{user.id}",
|
||||
items=[
|
||||
{
|
||||
"id": memory.id,
|
||||
"text": memory.content,
|
||||
"vector": request.app.state.EMBEDDING_FUNCTION(memory.content),
|
||||
"metadata": {
|
||||
"created_at": memory.created_at,
|
||||
"updated_at": memory.updated_at,
|
||||
},
|
||||
}
|
||||
for memory in memories
|
||||
],
|
||||
)
|
||||
|
||||
return True
|
||||
|
||||
|
||||
############################
|
||||
# DeleteMemoriesByUserId
|
||||
############################
|
||||
|
||||
|
||||
@router.delete("/delete/user", response_model=bool)
|
||||
async def delete_memory_by_user_id(user=Depends(get_verified_user)):
|
||||
result = Memories.delete_memories_by_user_id(user.id)
|
||||
|
||||
if result:
|
||||
try:
|
||||
VECTOR_DB_CLIENT.delete_collection(f"user-memory-{user.id}")
|
||||
except Exception as e:
|
||||
log.error(e)
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
############################
|
||||
# UpdateMemoryById
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/{memory_id}/update", response_model=Optional[MemoryModel])
|
||||
async def update_memory_by_id(
|
||||
memory_id: str,
|
||||
request: Request,
|
||||
form_data: MemoryUpdateModel,
|
||||
user=Depends(get_verified_user),
|
||||
):
|
||||
memory = Memories.update_memory_by_id(memory_id, form_data.content)
|
||||
if memory is None:
|
||||
raise HTTPException(status_code=404, detail="Memory not found")
|
||||
|
||||
if form_data.content is not None:
|
||||
VECTOR_DB_CLIENT.upsert(
|
||||
collection_name=f"user-memory-{user.id}",
|
||||
items=[
|
||||
{
|
||||
"id": memory.id,
|
||||
"text": memory.content,
|
||||
"vector": request.app.state.EMBEDDING_FUNCTION(memory.content),
|
||||
"metadata": {
|
||||
"created_at": memory.created_at,
|
||||
"updated_at": memory.updated_at,
|
||||
},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
return memory
|
||||
|
||||
|
||||
############################
|
||||
# DeleteMemoryById
|
||||
############################
|
||||
|
||||
|
||||
@router.delete("/{memory_id}", response_model=bool)
|
||||
async def delete_memory_by_id(memory_id: str, user=Depends(get_verified_user)):
|
||||
result = Memories.delete_memory_by_id_and_user_id(memory_id, user.id)
|
||||
|
||||
if result:
|
||||
VECTOR_DB_CLIENT.delete(
|
||||
collection_name=f"user-memory-{user.id}", ids=[memory_id]
|
||||
)
|
||||
return True
|
||||
|
||||
return False
|
||||
@@ -0,0 +1,189 @@
|
||||
from typing import Optional
|
||||
|
||||
from open_webui.models.models import (
|
||||
ModelForm,
|
||||
ModelModel,
|
||||
ModelResponse,
|
||||
ModelUserResponse,
|
||||
Models,
|
||||
)
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, status
|
||||
|
||||
|
||||
from open_webui.utils.auth import get_admin_user, get_verified_user
|
||||
from open_webui.utils.access_control import has_access, has_permission
|
||||
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
###########################
|
||||
# GetModels
|
||||
###########################
|
||||
|
||||
|
||||
@router.get("/", response_model=list[ModelUserResponse])
|
||||
async def get_models(id: Optional[str] = None, user=Depends(get_verified_user)):
|
||||
if user.role == "admin":
|
||||
return Models.get_models()
|
||||
else:
|
||||
return Models.get_models_by_user_id(user.id)
|
||||
|
||||
|
||||
###########################
|
||||
# GetBaseModels
|
||||
###########################
|
||||
|
||||
|
||||
@router.get("/base", response_model=list[ModelResponse])
|
||||
async def get_base_models(user=Depends(get_admin_user)):
|
||||
return Models.get_base_models()
|
||||
|
||||
|
||||
############################
|
||||
# CreateNewModel
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/create", response_model=Optional[ModelModel])
|
||||
async def create_new_model(
|
||||
request: Request,
|
||||
form_data: ModelForm,
|
||||
user=Depends(get_verified_user),
|
||||
):
|
||||
if user.role != "admin" and not has_permission(
|
||||
user.id, "workspace.models", request.app.state.config.USER_PERMISSIONS
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.UNAUTHORIZED,
|
||||
)
|
||||
|
||||
model = Models.get_model_by_id(form_data.id)
|
||||
if model:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.MODEL_ID_TAKEN,
|
||||
)
|
||||
|
||||
else:
|
||||
model = Models.insert_new_model(form_data, user.id)
|
||||
if model:
|
||||
return model
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.DEFAULT(),
|
||||
)
|
||||
|
||||
|
||||
###########################
|
||||
# GetModelById
|
||||
###########################
|
||||
|
||||
|
||||
# Note: We're not using the typical url path param here, but instead using a query parameter to allow '/' in the id
|
||||
@router.get("/model", response_model=Optional[ModelResponse])
|
||||
async def get_model_by_id(id: str, user=Depends(get_verified_user)):
|
||||
model = Models.get_model_by_id(id)
|
||||
if model:
|
||||
if (
|
||||
user.role == "admin"
|
||||
or model.user_id == user.id
|
||||
or has_access(user.id, "read", model.access_control)
|
||||
):
|
||||
return model
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# ToggelModelById
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/model/toggle", response_model=Optional[ModelResponse])
|
||||
async def toggle_model_by_id(id: str, user=Depends(get_verified_user)):
|
||||
model = Models.get_model_by_id(id)
|
||||
if model:
|
||||
if (
|
||||
user.role == "admin"
|
||||
or model.user_id == user.id
|
||||
or has_access(user.id, "write", model.access_control)
|
||||
):
|
||||
model = Models.toggle_model_by_id(id)
|
||||
|
||||
if model:
|
||||
return model
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error updating function"),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.UNAUTHORIZED,
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# UpdateModelById
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/model/update", response_model=Optional[ModelModel])
|
||||
async def update_model_by_id(
|
||||
id: str,
|
||||
form_data: ModelForm,
|
||||
user=Depends(get_verified_user),
|
||||
):
|
||||
model = Models.get_model_by_id(id)
|
||||
|
||||
if not model:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
model = Models.update_model_by_id(id, form_data)
|
||||
return model
|
||||
|
||||
|
||||
############################
|
||||
# DeleteModelById
|
||||
############################
|
||||
|
||||
|
||||
@router.delete("/model/delete", response_model=bool)
|
||||
async def delete_model_by_id(id: str, user=Depends(get_verified_user)):
|
||||
model = Models.get_model_by_id(id)
|
||||
if not model:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
if model.user_id != user.id and user.role != "admin":
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.UNAUTHORIZED,
|
||||
)
|
||||
|
||||
result = Models.delete_model_by_id(id)
|
||||
return result
|
||||
|
||||
|
||||
@router.delete("/delete/all", response_model=bool)
|
||||
async def delete_all_models(user=Depends(get_admin_user)):
|
||||
result = Models.delete_all_models()
|
||||
return result
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,772 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Literal, Optional, overload
|
||||
|
||||
import aiohttp
|
||||
from aiocache import cached
|
||||
import requests
|
||||
|
||||
|
||||
from fastapi import Depends, FastAPI, HTTPException, Request, APIRouter
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import FileResponse, StreamingResponse
|
||||
from pydantic import BaseModel
|
||||
from starlette.background import BackgroundTask
|
||||
|
||||
from open_webui.models.models import Models
|
||||
from open_webui.config import (
|
||||
CACHE_DIR,
|
||||
)
|
||||
from open_webui.env import (
|
||||
AIOHTTP_CLIENT_TIMEOUT,
|
||||
AIOHTTP_CLIENT_TIMEOUT_OPENAI_MODEL_LIST,
|
||||
ENABLE_FORWARD_USER_INFO_HEADERS,
|
||||
BYPASS_MODEL_ACCESS_CONTROL,
|
||||
)
|
||||
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
from open_webui.env import ENV, SRC_LOG_LEVELS
|
||||
|
||||
|
||||
from open_webui.utils.payload import (
|
||||
apply_model_params_to_body_openai,
|
||||
apply_model_system_prompt_to_body,
|
||||
)
|
||||
|
||||
from open_webui.utils.auth import get_admin_user, get_verified_user
|
||||
from open_webui.utils.access_control import has_access
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["OPENAI"])
|
||||
|
||||
|
||||
##########################################
|
||||
#
|
||||
# Utility functions
|
||||
#
|
||||
##########################################
|
||||
|
||||
|
||||
async def send_get_request(url, key=None):
|
||||
timeout = aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT_OPENAI_MODEL_LIST)
|
||||
try:
|
||||
async with aiohttp.ClientSession(timeout=timeout, trust_env=True) as session:
|
||||
async with session.get(
|
||||
url, headers={**({"Authorization": f"Bearer {key}"} if key else {})}
|
||||
) as response:
|
||||
return await response.json()
|
||||
except Exception as e:
|
||||
# Handle connection error here
|
||||
log.error(f"Connection error: {e}")
|
||||
return None
|
||||
|
||||
|
||||
async def cleanup_response(
|
||||
response: Optional[aiohttp.ClientResponse],
|
||||
session: Optional[aiohttp.ClientSession],
|
||||
):
|
||||
if response:
|
||||
response.close()
|
||||
if session:
|
||||
await session.close()
|
||||
|
||||
|
||||
def openai_o1_handler(payload):
|
||||
"""
|
||||
Handle O1 specific parameters
|
||||
"""
|
||||
if "max_tokens" in payload:
|
||||
# Remove "max_tokens" from the payload
|
||||
payload["max_completion_tokens"] = payload["max_tokens"]
|
||||
del payload["max_tokens"]
|
||||
|
||||
# Fix: O1 does not support the "system" parameter, Modify "system" to "user"
|
||||
if payload["messages"][0]["role"] == "system":
|
||||
payload["messages"][0]["role"] = "user"
|
||||
|
||||
return payload
|
||||
|
||||
|
||||
##########################################
|
||||
#
|
||||
# API routes
|
||||
#
|
||||
##########################################
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("/config")
|
||||
async def get_config(request: Request, user=Depends(get_admin_user)):
|
||||
return {
|
||||
"ENABLE_OPENAI_API": request.app.state.config.ENABLE_OPENAI_API,
|
||||
"OPENAI_API_BASE_URLS": request.app.state.config.OPENAI_API_BASE_URLS,
|
||||
"OPENAI_API_KEYS": request.app.state.config.OPENAI_API_KEYS,
|
||||
"OPENAI_API_CONFIGS": request.app.state.config.OPENAI_API_CONFIGS,
|
||||
}
|
||||
|
||||
|
||||
class OpenAIConfigForm(BaseModel):
|
||||
ENABLE_OPENAI_API: Optional[bool] = None
|
||||
OPENAI_API_BASE_URLS: list[str]
|
||||
OPENAI_API_KEYS: list[str]
|
||||
OPENAI_API_CONFIGS: dict
|
||||
|
||||
|
||||
@router.post("/config/update")
|
||||
async def update_config(
|
||||
request: Request, form_data: OpenAIConfigForm, user=Depends(get_admin_user)
|
||||
):
|
||||
request.app.state.config.ENABLE_OPENAI_API = form_data.ENABLE_OPENAI_API
|
||||
request.app.state.config.OPENAI_API_BASE_URLS = form_data.OPENAI_API_BASE_URLS
|
||||
request.app.state.config.OPENAI_API_KEYS = form_data.OPENAI_API_KEYS
|
||||
|
||||
# Check if API KEYS length is same than API URLS length
|
||||
if len(request.app.state.config.OPENAI_API_KEYS) != len(
|
||||
request.app.state.config.OPENAI_API_BASE_URLS
|
||||
):
|
||||
if len(request.app.state.config.OPENAI_API_KEYS) > len(
|
||||
request.app.state.config.OPENAI_API_BASE_URLS
|
||||
):
|
||||
request.app.state.config.OPENAI_API_KEYS = (
|
||||
request.app.state.config.OPENAI_API_KEYS[
|
||||
: len(request.app.state.config.OPENAI_API_BASE_URLS)
|
||||
]
|
||||
)
|
||||
else:
|
||||
request.app.state.config.OPENAI_API_KEYS += [""] * (
|
||||
len(request.app.state.config.OPENAI_API_BASE_URLS)
|
||||
- len(request.app.state.config.OPENAI_API_KEYS)
|
||||
)
|
||||
|
||||
request.app.state.config.OPENAI_API_CONFIGS = form_data.OPENAI_API_CONFIGS
|
||||
|
||||
# Remove any extra configs
|
||||
config_urls = request.app.state.config.OPENAI_API_CONFIGS.keys()
|
||||
for idx, url in enumerate(request.app.state.config.OPENAI_API_BASE_URLS):
|
||||
if url not in config_urls:
|
||||
request.app.state.config.OPENAI_API_CONFIGS.pop(url, None)
|
||||
|
||||
return {
|
||||
"ENABLE_OPENAI_API": request.app.state.config.ENABLE_OPENAI_API,
|
||||
"OPENAI_API_BASE_URLS": request.app.state.config.OPENAI_API_BASE_URLS,
|
||||
"OPENAI_API_KEYS": request.app.state.config.OPENAI_API_KEYS,
|
||||
"OPENAI_API_CONFIGS": request.app.state.config.OPENAI_API_CONFIGS,
|
||||
}
|
||||
|
||||
|
||||
@router.post("/audio/speech")
|
||||
async def speech(request: Request, user=Depends(get_verified_user)):
|
||||
idx = None
|
||||
try:
|
||||
idx = request.app.state.config.OPENAI_API_BASE_URLS.index(
|
||||
"https://api.openai.com/v1"
|
||||
)
|
||||
|
||||
body = await request.body()
|
||||
name = hashlib.sha256(body).hexdigest()
|
||||
|
||||
SPEECH_CACHE_DIR = Path(CACHE_DIR).joinpath("./audio/speech/")
|
||||
SPEECH_CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
file_path = SPEECH_CACHE_DIR.joinpath(f"{name}.mp3")
|
||||
file_body_path = SPEECH_CACHE_DIR.joinpath(f"{name}.json")
|
||||
|
||||
# Check if the file already exists in the cache
|
||||
if file_path.is_file():
|
||||
return FileResponse(file_path)
|
||||
|
||||
url = request.app.state.config.OPENAI_API_BASE_URLS[idx]
|
||||
|
||||
r = None
|
||||
try:
|
||||
r = requests.post(
|
||||
url=f"{url}/audio/speech",
|
||||
data=body,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {request.app.state.config.OPENAI_API_KEYS[idx]}",
|
||||
**(
|
||||
{
|
||||
"HTTP-Referer": "https://openwebui.com/",
|
||||
"X-Title": "Open WebUI",
|
||||
}
|
||||
if "openrouter.ai" in url
|
||||
else {}
|
||||
),
|
||||
**(
|
||||
{
|
||||
"X-OpenWebUI-User-Name": user.name,
|
||||
"X-OpenWebUI-User-Id": user.id,
|
||||
"X-OpenWebUI-User-Email": user.email,
|
||||
"X-OpenWebUI-User-Role": user.role,
|
||||
}
|
||||
if ENABLE_FORWARD_USER_INFO_HEADERS
|
||||
else {}
|
||||
),
|
||||
},
|
||||
stream=True,
|
||||
)
|
||||
|
||||
r.raise_for_status()
|
||||
|
||||
# Save the streaming content to a file
|
||||
with open(file_path, "wb") as f:
|
||||
for chunk in r.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
|
||||
with open(file_body_path, "w") as f:
|
||||
json.dump(json.loads(body.decode("utf-8")), f)
|
||||
|
||||
# Return the saved file
|
||||
return FileResponse(file_path)
|
||||
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
|
||||
detail = None
|
||||
if r is not None:
|
||||
try:
|
||||
res = r.json()
|
||||
if "error" in res:
|
||||
detail = f"External: {res['error']}"
|
||||
except Exception:
|
||||
detail = f"External: {e}"
|
||||
|
||||
raise HTTPException(
|
||||
status_code=r.status_code if r else 500,
|
||||
detail=detail if detail else "Open WebUI: Server Connection Error",
|
||||
)
|
||||
|
||||
except ValueError:
|
||||
raise HTTPException(status_code=401, detail=ERROR_MESSAGES.OPENAI_NOT_FOUND)
|
||||
|
||||
|
||||
async def get_all_models_responses(request: Request) -> list:
|
||||
if not request.app.state.config.ENABLE_OPENAI_API:
|
||||
return []
|
||||
|
||||
# Check if API KEYS length is same than API URLS length
|
||||
num_urls = len(request.app.state.config.OPENAI_API_BASE_URLS)
|
||||
num_keys = len(request.app.state.config.OPENAI_API_KEYS)
|
||||
|
||||
if num_keys != num_urls:
|
||||
# if there are more keys than urls, remove the extra keys
|
||||
if num_keys > num_urls:
|
||||
new_keys = request.app.state.config.OPENAI_API_KEYS[:num_urls]
|
||||
request.app.state.config.OPENAI_API_KEYS = new_keys
|
||||
# if there are more urls than keys, add empty keys
|
||||
else:
|
||||
request.app.state.config.OPENAI_API_KEYS += [""] * (num_urls - num_keys)
|
||||
|
||||
request_tasks = []
|
||||
for idx, url in enumerate(request.app.state.config.OPENAI_API_BASE_URLS):
|
||||
if url not in request.app.state.config.OPENAI_API_CONFIGS:
|
||||
request_tasks.append(
|
||||
send_get_request(
|
||||
f"{url}/models", request.app.state.config.OPENAI_API_KEYS[idx]
|
||||
)
|
||||
)
|
||||
else:
|
||||
api_config = request.app.state.config.OPENAI_API_CONFIGS.get(url, {})
|
||||
|
||||
enable = api_config.get("enable", True)
|
||||
model_ids = api_config.get("model_ids", [])
|
||||
|
||||
if enable:
|
||||
if len(model_ids) == 0:
|
||||
request_tasks.append(
|
||||
send_get_request(
|
||||
f"{url}/models",
|
||||
request.app.state.config.OPENAI_API_KEYS[idx],
|
||||
)
|
||||
)
|
||||
else:
|
||||
model_list = {
|
||||
"object": "list",
|
||||
"data": [
|
||||
{
|
||||
"id": model_id,
|
||||
"name": model_id,
|
||||
"owned_by": "openai",
|
||||
"openai": {"id": model_id},
|
||||
"urlIdx": idx,
|
||||
}
|
||||
for model_id in model_ids
|
||||
],
|
||||
}
|
||||
|
||||
request_tasks.append(
|
||||
asyncio.ensure_future(asyncio.sleep(0, model_list))
|
||||
)
|
||||
else:
|
||||
request_tasks.append(asyncio.ensure_future(asyncio.sleep(0, None)))
|
||||
|
||||
responses = await asyncio.gather(*request_tasks)
|
||||
|
||||
for idx, response in enumerate(responses):
|
||||
if response:
|
||||
url = request.app.state.config.OPENAI_API_BASE_URLS[idx]
|
||||
api_config = request.app.state.config.OPENAI_API_CONFIGS.get(url, {})
|
||||
|
||||
prefix_id = api_config.get("prefix_id", None)
|
||||
|
||||
if prefix_id:
|
||||
for model in (
|
||||
response if isinstance(response, list) else response.get("data", [])
|
||||
):
|
||||
model["id"] = f"{prefix_id}.{model['id']}"
|
||||
|
||||
log.debug(f"get_all_models:responses() {responses}")
|
||||
return responses
|
||||
|
||||
|
||||
async def get_filtered_models(models, user):
|
||||
# Filter models based on user access control
|
||||
filtered_models = []
|
||||
for model in models.get("data", []):
|
||||
model_info = Models.get_model_by_id(model["id"])
|
||||
if model_info:
|
||||
if user.id == model_info.user_id or has_access(
|
||||
user.id, type="read", access_control=model_info.access_control
|
||||
):
|
||||
filtered_models.append(model)
|
||||
return filtered_models
|
||||
|
||||
|
||||
@cached(ttl=3)
|
||||
async def get_all_models(request: Request) -> dict[str, list]:
|
||||
log.info("get_all_models()")
|
||||
|
||||
if not request.app.state.config.ENABLE_OPENAI_API:
|
||||
return {"data": []}
|
||||
|
||||
responses = await get_all_models_responses(request)
|
||||
|
||||
def extract_data(response):
|
||||
if response and "data" in response:
|
||||
return response["data"]
|
||||
if isinstance(response, list):
|
||||
return response
|
||||
return None
|
||||
|
||||
def merge_models_lists(model_lists):
|
||||
log.debug(f"merge_models_lists {model_lists}")
|
||||
merged_list = []
|
||||
|
||||
for idx, models in enumerate(model_lists):
|
||||
if models is not None and "error" not in models:
|
||||
merged_list.extend(
|
||||
[
|
||||
{
|
||||
**model,
|
||||
"name": model.get("name", model["id"]),
|
||||
"owned_by": "openai",
|
||||
"openai": model,
|
||||
"urlIdx": idx,
|
||||
}
|
||||
for model in models
|
||||
if "api.openai.com"
|
||||
not in request.app.state.config.OPENAI_API_BASE_URLS[idx]
|
||||
or not any(
|
||||
name in model["id"]
|
||||
for name in [
|
||||
"babbage",
|
||||
"dall-e",
|
||||
"davinci",
|
||||
"embedding",
|
||||
"tts",
|
||||
"whisper",
|
||||
]
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
return merged_list
|
||||
|
||||
models = {"data": merge_models_lists(map(extract_data, responses))}
|
||||
log.debug(f"models: {models}")
|
||||
|
||||
request.app.state.OPENAI_MODELS = {model["id"]: model for model in models["data"]}
|
||||
return models
|
||||
|
||||
|
||||
@router.get("/models")
|
||||
@router.get("/models/{url_idx}")
|
||||
async def get_models(
|
||||
request: Request, url_idx: Optional[int] = None, user=Depends(get_verified_user)
|
||||
):
|
||||
models = {
|
||||
"data": [],
|
||||
}
|
||||
|
||||
if url_idx is None:
|
||||
models = await get_all_models(request)
|
||||
else:
|
||||
url = request.app.state.config.OPENAI_API_BASE_URLS[url_idx]
|
||||
key = request.app.state.config.OPENAI_API_KEYS[url_idx]
|
||||
|
||||
r = None
|
||||
async with aiohttp.ClientSession(
|
||||
timeout=aiohttp.ClientTimeout(
|
||||
total=AIOHTTP_CLIENT_TIMEOUT_OPENAI_MODEL_LIST
|
||||
)
|
||||
) as session:
|
||||
try:
|
||||
async with session.get(
|
||||
f"{url}/models",
|
||||
headers={
|
||||
"Authorization": f"Bearer {key}",
|
||||
"Content-Type": "application/json",
|
||||
**(
|
||||
{
|
||||
"X-OpenWebUI-User-Name": user.name,
|
||||
"X-OpenWebUI-User-Id": user.id,
|
||||
"X-OpenWebUI-User-Email": user.email,
|
||||
"X-OpenWebUI-User-Role": user.role,
|
||||
}
|
||||
if ENABLE_FORWARD_USER_INFO_HEADERS
|
||||
else {}
|
||||
),
|
||||
},
|
||||
) as r:
|
||||
if r.status != 200:
|
||||
# Extract response error details if available
|
||||
error_detail = f"HTTP Error: {r.status}"
|
||||
res = await r.json()
|
||||
if "error" in res:
|
||||
error_detail = f"External Error: {res['error']}"
|
||||
raise Exception(error_detail)
|
||||
|
||||
response_data = await r.json()
|
||||
|
||||
# Check if we're calling OpenAI API based on the URL
|
||||
if "api.openai.com" in url:
|
||||
# Filter models according to the specified conditions
|
||||
response_data["data"] = [
|
||||
model
|
||||
for model in response_data.get("data", [])
|
||||
if not any(
|
||||
name in model["id"]
|
||||
for name in [
|
||||
"babbage",
|
||||
"dall-e",
|
||||
"davinci",
|
||||
"embedding",
|
||||
"tts",
|
||||
"whisper",
|
||||
]
|
||||
)
|
||||
]
|
||||
|
||||
models = response_data
|
||||
except aiohttp.ClientError as e:
|
||||
# ClientError covers all aiohttp requests issues
|
||||
log.exception(f"Client error: {str(e)}")
|
||||
raise HTTPException(
|
||||
status_code=500, detail="Open WebUI: Server Connection Error"
|
||||
)
|
||||
except Exception as e:
|
||||
log.exception(f"Unexpected error: {e}")
|
||||
error_detail = f"Unexpected error: {str(e)}"
|
||||
raise HTTPException(status_code=500, detail=error_detail)
|
||||
|
||||
if user.role == "user" and not BYPASS_MODEL_ACCESS_CONTROL:
|
||||
models["data"] = get_filtered_models(models, user)
|
||||
|
||||
return models
|
||||
|
||||
|
||||
class ConnectionVerificationForm(BaseModel):
|
||||
url: str
|
||||
key: str
|
||||
|
||||
|
||||
@router.post("/verify")
|
||||
async def verify_connection(
|
||||
form_data: ConnectionVerificationForm, user=Depends(get_admin_user)
|
||||
):
|
||||
url = form_data.url
|
||||
key = form_data.key
|
||||
|
||||
async with aiohttp.ClientSession(
|
||||
timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT_OPENAI_MODEL_LIST)
|
||||
) as session:
|
||||
try:
|
||||
async with session.get(
|
||||
f"{url}/models",
|
||||
headers={
|
||||
"Authorization": f"Bearer {key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
) as r:
|
||||
if r.status != 200:
|
||||
# Extract response error details if available
|
||||
error_detail = f"HTTP Error: {r.status}"
|
||||
res = await r.json()
|
||||
if "error" in res:
|
||||
error_detail = f"External Error: {res['error']}"
|
||||
raise Exception(error_detail)
|
||||
|
||||
response_data = await r.json()
|
||||
return response_data
|
||||
|
||||
except aiohttp.ClientError as e:
|
||||
# ClientError covers all aiohttp requests issues
|
||||
log.exception(f"Client error: {str(e)}")
|
||||
raise HTTPException(
|
||||
status_code=500, detail="Open WebUI: Server Connection Error"
|
||||
)
|
||||
except Exception as e:
|
||||
log.exception(f"Unexpected error: {e}")
|
||||
error_detail = f"Unexpected error: {str(e)}"
|
||||
raise HTTPException(status_code=500, detail=error_detail)
|
||||
|
||||
|
||||
@router.post("/chat/completions")
|
||||
async def generate_chat_completion(
|
||||
request: Request,
|
||||
form_data: dict,
|
||||
user=Depends(get_verified_user),
|
||||
bypass_filter: Optional[bool] = False,
|
||||
):
|
||||
idx = 0
|
||||
payload = {**form_data}
|
||||
if "metadata" in payload:
|
||||
del payload["metadata"]
|
||||
|
||||
model_id = form_data.get("model")
|
||||
model_info = Models.get_model_by_id(model_id)
|
||||
|
||||
# Check model info and override the payload
|
||||
if model_info:
|
||||
if model_info.base_model_id:
|
||||
payload["model"] = model_info.base_model_id
|
||||
|
||||
params = model_info.params.model_dump()
|
||||
payload = apply_model_params_to_body_openai(params, payload)
|
||||
payload = apply_model_system_prompt_to_body(params, payload, user)
|
||||
|
||||
# Check if user has access to the model
|
||||
if not bypass_filter and user.role == "user":
|
||||
if not (
|
||||
user.id == model_info.user_id
|
||||
or has_access(
|
||||
user.id, type="read", access_control=model_info.access_control
|
||||
)
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail="Model not found",
|
||||
)
|
||||
elif not bypass_filter:
|
||||
if user.role != "admin":
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail="Model not found",
|
||||
)
|
||||
|
||||
model = request.app.state.OPENAI_MODELS.get(model_id)
|
||||
if model:
|
||||
idx = model["urlIdx"]
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail="Model not found",
|
||||
)
|
||||
|
||||
# Get the API config for the model
|
||||
api_config = request.app.state.config.OPENAI_API_CONFIGS.get(
|
||||
request.app.state.config.OPENAI_API_BASE_URLS[idx], {}
|
||||
)
|
||||
|
||||
prefix_id = api_config.get("prefix_id", None)
|
||||
if prefix_id:
|
||||
payload["model"] = payload["model"].replace(f"{prefix_id}.", "")
|
||||
|
||||
# Add user info to the payload if the model is a pipeline
|
||||
if "pipeline" in model and model.get("pipeline"):
|
||||
payload["user"] = {
|
||||
"name": user.name,
|
||||
"id": user.id,
|
||||
"email": user.email,
|
||||
"role": user.role,
|
||||
}
|
||||
|
||||
url = request.app.state.config.OPENAI_API_BASE_URLS[idx]
|
||||
key = request.app.state.config.OPENAI_API_KEYS[idx]
|
||||
|
||||
# Fix: O1 does not support the "max_tokens" parameter, Modify "max_tokens" to "max_completion_tokens"
|
||||
is_o1 = payload["model"].lower().startswith("o1-")
|
||||
if is_o1:
|
||||
payload = openai_o1_handler(payload)
|
||||
elif "api.openai.com" not in url:
|
||||
# Remove "max_tokens" from the payload for backward compatibility
|
||||
if "max_tokens" in payload:
|
||||
payload["max_completion_tokens"] = payload["max_tokens"]
|
||||
del payload["max_tokens"]
|
||||
|
||||
# TODO: check if below is needed
|
||||
# if "max_tokens" in payload and "max_completion_tokens" in payload:
|
||||
# del payload["max_tokens"]
|
||||
|
||||
# Convert the modified body back to JSON
|
||||
payload = json.dumps(payload)
|
||||
|
||||
r = None
|
||||
session = None
|
||||
streaming = False
|
||||
response = None
|
||||
|
||||
try:
|
||||
session = aiohttp.ClientSession(
|
||||
trust_env=True, timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT)
|
||||
)
|
||||
|
||||
r = await session.request(
|
||||
method="POST",
|
||||
url=f"{url}/chat/completions",
|
||||
data=payload,
|
||||
headers={
|
||||
"Authorization": f"Bearer {key}",
|
||||
"Content-Type": "application/json",
|
||||
**(
|
||||
{
|
||||
"HTTP-Referer": "https://openwebui.com/",
|
||||
"X-Title": "Open WebUI",
|
||||
}
|
||||
if "openrouter.ai" in url
|
||||
else {}
|
||||
),
|
||||
**(
|
||||
{
|
||||
"X-OpenWebUI-User-Name": user.name,
|
||||
"X-OpenWebUI-User-Id": user.id,
|
||||
"X-OpenWebUI-User-Email": user.email,
|
||||
"X-OpenWebUI-User-Role": user.role,
|
||||
}
|
||||
if ENABLE_FORWARD_USER_INFO_HEADERS
|
||||
else {}
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
# Check if response is SSE
|
||||
if "text/event-stream" in r.headers.get("Content-Type", ""):
|
||||
streaming = True
|
||||
return StreamingResponse(
|
||||
r.content,
|
||||
status_code=r.status,
|
||||
headers=dict(r.headers),
|
||||
background=BackgroundTask(
|
||||
cleanup_response, response=r, session=session
|
||||
),
|
||||
)
|
||||
else:
|
||||
try:
|
||||
response = await r.json()
|
||||
except Exception as e:
|
||||
log.error(e)
|
||||
response = await r.text()
|
||||
|
||||
r.raise_for_status()
|
||||
return response
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
|
||||
detail = None
|
||||
if isinstance(response, dict):
|
||||
if "error" in response:
|
||||
detail = f"{response['error']['message'] if 'message' in response['error'] else response['error']}"
|
||||
elif isinstance(response, str):
|
||||
detail = response
|
||||
|
||||
raise HTTPException(
|
||||
status_code=r.status if r else 500,
|
||||
detail=detail if detail else "Open WebUI: Server Connection Error",
|
||||
)
|
||||
finally:
|
||||
if not streaming and session:
|
||||
if r:
|
||||
r.close()
|
||||
await session.close()
|
||||
|
||||
|
||||
@router.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
|
||||
async def proxy(path: str, request: Request, user=Depends(get_verified_user)):
|
||||
"""
|
||||
Deprecated: proxy all requests to OpenAI API
|
||||
"""
|
||||
|
||||
body = await request.body()
|
||||
|
||||
idx = 0
|
||||
url = request.app.state.config.OPENAI_API_BASE_URLS[idx]
|
||||
key = request.app.state.config.OPENAI_API_KEYS[idx]
|
||||
|
||||
r = None
|
||||
session = None
|
||||
streaming = False
|
||||
|
||||
try:
|
||||
session = aiohttp.ClientSession(trust_env=True)
|
||||
r = await session.request(
|
||||
method=request.method,
|
||||
url=f"{url}/{path}",
|
||||
data=body,
|
||||
headers={
|
||||
"Authorization": f"Bearer {key}",
|
||||
"Content-Type": "application/json",
|
||||
**(
|
||||
{
|
||||
"X-OpenWebUI-User-Name": user.name,
|
||||
"X-OpenWebUI-User-Id": user.id,
|
||||
"X-OpenWebUI-User-Email": user.email,
|
||||
"X-OpenWebUI-User-Role": user.role,
|
||||
}
|
||||
if ENABLE_FORWARD_USER_INFO_HEADERS
|
||||
else {}
|
||||
),
|
||||
},
|
||||
)
|
||||
r.raise_for_status()
|
||||
|
||||
# Check if response is SSE
|
||||
if "text/event-stream" in r.headers.get("Content-Type", ""):
|
||||
streaming = True
|
||||
return StreamingResponse(
|
||||
r.content,
|
||||
status_code=r.status,
|
||||
headers=dict(r.headers),
|
||||
background=BackgroundTask(
|
||||
cleanup_response, response=r, session=session
|
||||
),
|
||||
)
|
||||
else:
|
||||
response_data = await r.json()
|
||||
return response_data
|
||||
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
|
||||
detail = None
|
||||
if r is not None:
|
||||
try:
|
||||
res = await r.json()
|
||||
print(res)
|
||||
if "error" in res:
|
||||
detail = f"External: {res['error']['message'] if 'message' in res['error'] else res['error']}"
|
||||
except Exception:
|
||||
detail = f"External: {e}"
|
||||
raise HTTPException(
|
||||
status_code=r.status if r else 500,
|
||||
detail=detail if detail else "Open WebUI: Server Connection Error",
|
||||
)
|
||||
finally:
|
||||
if not streaming and session:
|
||||
if r:
|
||||
r.close()
|
||||
await session.close()
|
||||
@@ -0,0 +1,496 @@
|
||||
from fastapi import (
|
||||
Depends,
|
||||
FastAPI,
|
||||
File,
|
||||
Form,
|
||||
HTTPException,
|
||||
Request,
|
||||
UploadFile,
|
||||
status,
|
||||
APIRouter,
|
||||
)
|
||||
import os
|
||||
import logging
|
||||
import shutil
|
||||
import requests
|
||||
from pydantic import BaseModel
|
||||
from starlette.responses import FileResponse
|
||||
from typing import Optional
|
||||
|
||||
from open_webui.env import SRC_LOG_LEVELS
|
||||
from open_webui.config import CACHE_DIR
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
|
||||
|
||||
from open_webui.routers.openai import get_all_models_responses
|
||||
|
||||
from open_webui.utils.auth import get_admin_user
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["MAIN"])
|
||||
|
||||
|
||||
##################################
|
||||
#
|
||||
# Pipeline Middleware
|
||||
#
|
||||
##################################
|
||||
|
||||
|
||||
def get_sorted_filters(model_id, models):
|
||||
filters = [
|
||||
model
|
||||
for model in 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"])
|
||||
return sorted_filters
|
||||
|
||||
|
||||
def process_pipeline_inlet_filter(request, payload, user, models):
|
||||
user = {"id": user.id, "email": user.email, "name": user.name, "role": user.role}
|
||||
model_id = payload["model"]
|
||||
|
||||
sorted_filters = get_sorted_filters(model_id, models)
|
||||
model = models[model_id]
|
||||
|
||||
if "pipeline" in model:
|
||||
sorted_filters.append(model)
|
||||
|
||||
for filter in sorted_filters:
|
||||
r = None
|
||||
try:
|
||||
urlIdx = filter["urlIdx"]
|
||||
|
||||
url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx]
|
||||
key = request.app.state.config.OPENAI_API_KEYS[urlIdx]
|
||||
|
||||
if key == "":
|
||||
continue
|
||||
|
||||
headers = {"Authorization": f"Bearer {key}"}
|
||||
r = requests.post(
|
||||
f"{url}/{filter['id']}/filter/inlet",
|
||||
headers=headers,
|
||||
json={
|
||||
"user": user,
|
||||
"body": payload,
|
||||
},
|
||||
)
|
||||
|
||||
r.raise_for_status()
|
||||
payload = r.json()
|
||||
except Exception as e:
|
||||
# Handle connection error here
|
||||
print(f"Connection error: {e}")
|
||||
|
||||
if r is not None:
|
||||
res = r.json()
|
||||
if "detail" in res:
|
||||
raise Exception(r.status_code, res["detail"])
|
||||
|
||||
return payload
|
||||
|
||||
|
||||
def process_pipeline_outlet_filter(request, payload, user, models):
|
||||
user = {"id": user.id, "email": user.email, "name": user.name, "role": user.role}
|
||||
model_id = payload["model"]
|
||||
|
||||
sorted_filters = get_sorted_filters(model_id, models)
|
||||
model = models[model_id]
|
||||
|
||||
if "pipeline" in model:
|
||||
sorted_filters = [model] + sorted_filters
|
||||
|
||||
for filter in sorted_filters:
|
||||
r = None
|
||||
try:
|
||||
urlIdx = filter["urlIdx"]
|
||||
|
||||
url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx]
|
||||
key = request.app.state.config.OPENAI_API_KEYS[urlIdx]
|
||||
|
||||
if key != "":
|
||||
r = requests.post(
|
||||
f"{url}/{filter['id']}/filter/outlet",
|
||||
headers={"Authorization": f"Bearer {key}"},
|
||||
json={
|
||||
"user": {
|
||||
"id": user.id,
|
||||
"name": user.name,
|
||||
"email": user.email,
|
||||
"role": user.role,
|
||||
},
|
||||
"body": data,
|
||||
},
|
||||
)
|
||||
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
except Exception as e:
|
||||
# Handle connection error here
|
||||
print(f"Connection error: {e}")
|
||||
|
||||
if r is not None:
|
||||
try:
|
||||
res = r.json()
|
||||
if "detail" in res:
|
||||
return Exception(r.status_code, res)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
else:
|
||||
pass
|
||||
|
||||
return payload
|
||||
|
||||
|
||||
##################################
|
||||
#
|
||||
# Pipelines Endpoints
|
||||
#
|
||||
##################################
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("/list")
|
||||
async def get_pipelines_list(request: Request, user=Depends(get_admin_user)):
|
||||
responses = await get_all_models_responses(request)
|
||||
log.debug(f"get_pipelines_list: get_openai_models_responses returned {responses}")
|
||||
|
||||
urlIdxs = [
|
||||
idx
|
||||
for idx, response in enumerate(responses)
|
||||
if response is not None and "pipelines" in response
|
||||
]
|
||||
|
||||
return {
|
||||
"data": [
|
||||
{
|
||||
"url": request.app.state.config.OPENAI_API_BASE_URLS[urlIdx],
|
||||
"idx": urlIdx,
|
||||
}
|
||||
for urlIdx in urlIdxs
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
@router.post("/upload")
|
||||
async def upload_pipeline(
|
||||
request: Request,
|
||||
urlIdx: int = Form(...),
|
||||
file: UploadFile = File(...),
|
||||
user=Depends(get_admin_user),
|
||||
):
|
||||
print("upload_pipeline", urlIdx, file.filename)
|
||||
# Check if the uploaded file is a python file
|
||||
if not (file.filename and file.filename.endswith(".py")):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Only Python (.py) files are allowed.",
|
||||
)
|
||||
|
||||
upload_folder = f"{CACHE_DIR}/pipelines"
|
||||
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:
|
||||
shutil.copyfileobj(file.file, buffer)
|
||||
|
||||
url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx]
|
||||
key = request.app.state.config.OPENAI_API_KEYS[urlIdx]
|
||||
|
||||
with open(file_path, "rb") as f:
|
||||
files = {"file": f}
|
||||
r = requests.post(
|
||||
f"{url}/pipelines/upload",
|
||||
headers={"Authorization": f"Bearer {key}"},
|
||||
files=files,
|
||||
)
|
||||
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
|
||||
return {**data}
|
||||
except Exception as e:
|
||||
# Handle connection error here
|
||||
print(f"Connection error: {e}")
|
||||
|
||||
detail = None
|
||||
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:
|
||||
detail = res["detail"]
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
raise HTTPException(
|
||||
status_code=status_code,
|
||||
detail=detail if detail else "Pipeline not found",
|
||||
)
|
||||
finally:
|
||||
# Ensure the file is deleted after the upload is completed or on failure
|
||||
if os.path.exists(file_path):
|
||||
os.remove(file_path)
|
||||
|
||||
|
||||
class AddPipelineForm(BaseModel):
|
||||
url: str
|
||||
urlIdx: int
|
||||
|
||||
|
||||
@router.post("/add")
|
||||
async def add_pipeline(
|
||||
request: Request, form_data: AddPipelineForm, user=Depends(get_admin_user)
|
||||
):
|
||||
r = None
|
||||
try:
|
||||
urlIdx = form_data.urlIdx
|
||||
|
||||
url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx]
|
||||
key = request.app.state.config.OPENAI_API_KEYS[urlIdx]
|
||||
|
||||
r = requests.post(
|
||||
f"{url}/pipelines/add",
|
||||
headers={"Authorization": f"Bearer {key}"},
|
||||
json={"url": form_data.url},
|
||||
)
|
||||
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
|
||||
return {**data}
|
||||
except Exception as e:
|
||||
# Handle connection error here
|
||||
print(f"Connection error: {e}")
|
||||
|
||||
detail = None
|
||||
if r is not None:
|
||||
try:
|
||||
res = r.json()
|
||||
if "detail" in res:
|
||||
detail = res["detail"]
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
raise HTTPException(
|
||||
status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
|
||||
detail=detail if detail else "Pipeline not found",
|
||||
)
|
||||
|
||||
|
||||
class DeletePipelineForm(BaseModel):
|
||||
id: str
|
||||
urlIdx: int
|
||||
|
||||
|
||||
@router.delete("/delete")
|
||||
async def delete_pipeline(
|
||||
request: Request, form_data: DeletePipelineForm, user=Depends(get_admin_user)
|
||||
):
|
||||
r = None
|
||||
try:
|
||||
urlIdx = form_data.urlIdx
|
||||
|
||||
url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx]
|
||||
key = request.app.state.config.OPENAI_API_KEYS[urlIdx]
|
||||
|
||||
r = requests.delete(
|
||||
f"{url}/pipelines/delete",
|
||||
headers={"Authorization": f"Bearer {key}"},
|
||||
json={"id": form_data.id},
|
||||
)
|
||||
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
|
||||
return {**data}
|
||||
except Exception as e:
|
||||
# Handle connection error here
|
||||
print(f"Connection error: {e}")
|
||||
|
||||
detail = None
|
||||
if r is not None:
|
||||
try:
|
||||
res = r.json()
|
||||
if "detail" in res:
|
||||
detail = res["detail"]
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
raise HTTPException(
|
||||
status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
|
||||
detail=detail if detail else "Pipeline not found",
|
||||
)
|
||||
|
||||
|
||||
@router.get("/")
|
||||
async def get_pipelines(
|
||||
request: Request, urlIdx: Optional[int] = None, user=Depends(get_admin_user)
|
||||
):
|
||||
r = None
|
||||
try:
|
||||
url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx]
|
||||
key = request.app.state.config.OPENAI_API_KEYS[urlIdx]
|
||||
|
||||
r = requests.get(f"{url}/pipelines", headers={"Authorization": f"Bearer {key}"})
|
||||
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
|
||||
return {**data}
|
||||
except Exception as e:
|
||||
# Handle connection error here
|
||||
print(f"Connection error: {e}")
|
||||
|
||||
detail = None
|
||||
if r is not None:
|
||||
try:
|
||||
res = r.json()
|
||||
if "detail" in res:
|
||||
detail = res["detail"]
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
raise HTTPException(
|
||||
status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
|
||||
detail=detail if detail else "Pipeline not found",
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{pipeline_id}/valves")
|
||||
async def get_pipeline_valves(
|
||||
request: Request,
|
||||
urlIdx: Optional[int],
|
||||
pipeline_id: str,
|
||||
user=Depends(get_admin_user),
|
||||
):
|
||||
r = None
|
||||
try:
|
||||
url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx]
|
||||
key = request.app.state.config.OPENAI_API_KEYS[urlIdx]
|
||||
|
||||
r = requests.get(
|
||||
f"{url}/{pipeline_id}/valves", headers={"Authorization": f"Bearer {key}"}
|
||||
)
|
||||
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
|
||||
return {**data}
|
||||
except Exception as e:
|
||||
# Handle connection error here
|
||||
print(f"Connection error: {e}")
|
||||
|
||||
detail = None
|
||||
if r is not None:
|
||||
try:
|
||||
res = r.json()
|
||||
if "detail" in res:
|
||||
detail = res["detail"]
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
raise HTTPException(
|
||||
status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
|
||||
detail=detail if detail else "Pipeline not found",
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{pipeline_id}/valves/spec")
|
||||
async def get_pipeline_valves_spec(
|
||||
request: Request,
|
||||
urlIdx: Optional[int],
|
||||
pipeline_id: str,
|
||||
user=Depends(get_admin_user),
|
||||
):
|
||||
r = None
|
||||
try:
|
||||
url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx]
|
||||
key = request.app.state.config.OPENAI_API_KEYS[urlIdx]
|
||||
|
||||
r = requests.get(
|
||||
f"{url}/{pipeline_id}/valves/spec",
|
||||
headers={"Authorization": f"Bearer {key}"},
|
||||
)
|
||||
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
|
||||
return {**data}
|
||||
except Exception as e:
|
||||
# Handle connection error here
|
||||
print(f"Connection error: {e}")
|
||||
|
||||
detail = None
|
||||
if r is not None:
|
||||
try:
|
||||
res = r.json()
|
||||
if "detail" in res:
|
||||
detail = res["detail"]
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
raise HTTPException(
|
||||
status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
|
||||
detail=detail if detail else "Pipeline not found",
|
||||
)
|
||||
|
||||
|
||||
@router.post("/{pipeline_id}/valves/update")
|
||||
async def update_pipeline_valves(
|
||||
request: Request,
|
||||
urlIdx: Optional[int],
|
||||
pipeline_id: str,
|
||||
form_data: dict,
|
||||
user=Depends(get_admin_user),
|
||||
):
|
||||
r = None
|
||||
try:
|
||||
url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx]
|
||||
key = request.app.state.config.OPENAI_API_KEYS[urlIdx]
|
||||
|
||||
r = requests.post(
|
||||
f"{url}/{pipeline_id}/valves/update",
|
||||
headers={"Authorization": f"Bearer {key}"},
|
||||
json={**form_data},
|
||||
)
|
||||
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
|
||||
return {**data}
|
||||
except Exception as e:
|
||||
# Handle connection error here
|
||||
print(f"Connection error: {e}")
|
||||
|
||||
detail = None
|
||||
|
||||
if r is not None:
|
||||
try:
|
||||
res = r.json()
|
||||
if "detail" in res:
|
||||
detail = res["detail"]
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
raise HTTPException(
|
||||
status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
|
||||
detail=detail if detail else "Pipeline not found",
|
||||
)
|
||||
@@ -0,0 +1,152 @@
|
||||
from typing import Optional
|
||||
|
||||
from open_webui.models.prompts import (
|
||||
PromptForm,
|
||||
PromptUserResponse,
|
||||
PromptModel,
|
||||
Prompts,
|
||||
)
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
from fastapi import APIRouter, Depends, HTTPException, status, Request
|
||||
from open_webui.utils.auth import get_admin_user, get_verified_user
|
||||
from open_webui.utils.access_control import has_access, has_permission
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
############################
|
||||
# GetPrompts
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/", response_model=list[PromptModel])
|
||||
async def get_prompts(user=Depends(get_verified_user)):
|
||||
if user.role == "admin":
|
||||
prompts = Prompts.get_prompts()
|
||||
else:
|
||||
prompts = Prompts.get_prompts_by_user_id(user.id, "read")
|
||||
|
||||
return prompts
|
||||
|
||||
|
||||
@router.get("/list", response_model=list[PromptUserResponse])
|
||||
async def get_prompt_list(user=Depends(get_verified_user)):
|
||||
if user.role == "admin":
|
||||
prompts = Prompts.get_prompts()
|
||||
else:
|
||||
prompts = Prompts.get_prompts_by_user_id(user.id, "write")
|
||||
|
||||
return prompts
|
||||
|
||||
|
||||
############################
|
||||
# CreateNewPrompt
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/create", response_model=Optional[PromptModel])
|
||||
async def create_new_prompt(
|
||||
request: Request, form_data: PromptForm, user=Depends(get_verified_user)
|
||||
):
|
||||
if user.role != "admin" and not has_permission(
|
||||
user.id, "workspace.prompts", request.app.state.config.USER_PERMISSIONS
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.UNAUTHORIZED,
|
||||
)
|
||||
|
||||
prompt = Prompts.get_prompt_by_command(form_data.command)
|
||||
if prompt is None:
|
||||
prompt = Prompts.insert_new_prompt(user.id, form_data)
|
||||
|
||||
if prompt:
|
||||
return prompt
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(),
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.COMMAND_TAKEN,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# GetPromptByCommand
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/command/{command}", response_model=Optional[PromptModel])
|
||||
async def get_prompt_by_command(command: str, user=Depends(get_verified_user)):
|
||||
prompt = Prompts.get_prompt_by_command(f"/{command}")
|
||||
|
||||
if prompt:
|
||||
if (
|
||||
user.role == "admin"
|
||||
or prompt.user_id == user.id
|
||||
or has_access(user.id, "read", prompt.access_control)
|
||||
):
|
||||
return prompt
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# UpdatePromptByCommand
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/command/{command}/update", response_model=Optional[PromptModel])
|
||||
async def update_prompt_by_command(
|
||||
command: str,
|
||||
form_data: PromptForm,
|
||||
user=Depends(get_verified_user),
|
||||
):
|
||||
prompt = Prompts.get_prompt_by_command(f"/{command}")
|
||||
if not prompt:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
if prompt.user_id != user.id and user.role != "admin":
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
|
||||
prompt = Prompts.update_prompt_by_command(f"/{command}", form_data)
|
||||
if prompt:
|
||||
return prompt
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# DeletePromptByCommand
|
||||
############################
|
||||
|
||||
|
||||
@router.delete("/command/{command}/delete", response_model=bool)
|
||||
async def delete_prompt_by_command(command: str, user=Depends(get_verified_user)):
|
||||
prompt = Prompts.get_prompt_by_command(f"/{command}")
|
||||
if not prompt:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
if prompt.user_id != user.id and user.role != "admin":
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
|
||||
result = Prompts.delete_prompt_by_command(f"/{command}")
|
||||
return result
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,512 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, Response, status, Request
|
||||
from fastapi.responses import JSONResponse, RedirectResponse
|
||||
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional
|
||||
import logging
|
||||
|
||||
from open_webui.utils.chat import generate_chat_completion
|
||||
from open_webui.utils.task import (
|
||||
title_generation_template,
|
||||
query_generation_template,
|
||||
autocomplete_generation_template,
|
||||
tags_generation_template,
|
||||
emoji_generation_template,
|
||||
moa_response_generation_template,
|
||||
)
|
||||
from open_webui.utils.auth import get_admin_user, get_verified_user
|
||||
from open_webui.constants import TASKS
|
||||
|
||||
from open_webui.routers.pipelines import process_pipeline_inlet_filter
|
||||
from open_webui.utils.task import get_task_model_id
|
||||
|
||||
from open_webui.config import (
|
||||
DEFAULT_TITLE_GENERATION_PROMPT_TEMPLATE,
|
||||
DEFAULT_TAGS_GENERATION_PROMPT_TEMPLATE,
|
||||
DEFAULT_QUERY_GENERATION_PROMPT_TEMPLATE,
|
||||
DEFAULT_AUTOCOMPLETE_GENERATION_PROMPT_TEMPLATE,
|
||||
DEFAULT_EMOJI_GENERATION_PROMPT_TEMPLATE,
|
||||
DEFAULT_MOA_GENERATION_PROMPT_TEMPLATE,
|
||||
)
|
||||
from open_webui.env import SRC_LOG_LEVELS
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["MODELS"])
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
##################################
|
||||
#
|
||||
# Task Endpoints
|
||||
#
|
||||
##################################
|
||||
|
||||
|
||||
@router.get("/config")
|
||||
async def get_task_config(request: Request, user=Depends(get_verified_user)):
|
||||
return {
|
||||
"TASK_MODEL": request.app.state.config.TASK_MODEL,
|
||||
"TASK_MODEL_EXTERNAL": request.app.state.config.TASK_MODEL_EXTERNAL,
|
||||
"TITLE_GENERATION_PROMPT_TEMPLATE": request.app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE,
|
||||
"ENABLE_AUTOCOMPLETE_GENERATION": request.app.state.config.ENABLE_AUTOCOMPLETE_GENERATION,
|
||||
"AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH": request.app.state.config.AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH,
|
||||
"TAGS_GENERATION_PROMPT_TEMPLATE": request.app.state.config.TAGS_GENERATION_PROMPT_TEMPLATE,
|
||||
"ENABLE_TAGS_GENERATION": request.app.state.config.ENABLE_TAGS_GENERATION,
|
||||
"ENABLE_SEARCH_QUERY_GENERATION": request.app.state.config.ENABLE_SEARCH_QUERY_GENERATION,
|
||||
"ENABLE_RETRIEVAL_QUERY_GENERATION": request.app.state.config.ENABLE_RETRIEVAL_QUERY_GENERATION,
|
||||
"QUERY_GENERATION_PROMPT_TEMPLATE": request.app.state.config.QUERY_GENERATION_PROMPT_TEMPLATE,
|
||||
"TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE": request.app.state.config.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE,
|
||||
}
|
||||
|
||||
|
||||
class TaskConfigForm(BaseModel):
|
||||
TASK_MODEL: Optional[str]
|
||||
TASK_MODEL_EXTERNAL: Optional[str]
|
||||
TITLE_GENERATION_PROMPT_TEMPLATE: str
|
||||
ENABLE_AUTOCOMPLETE_GENERATION: bool
|
||||
AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH: int
|
||||
TAGS_GENERATION_PROMPT_TEMPLATE: str
|
||||
ENABLE_TAGS_GENERATION: bool
|
||||
ENABLE_SEARCH_QUERY_GENERATION: bool
|
||||
ENABLE_RETRIEVAL_QUERY_GENERATION: bool
|
||||
QUERY_GENERATION_PROMPT_TEMPLATE: str
|
||||
TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE: str
|
||||
|
||||
|
||||
@router.post("/config/update")
|
||||
async def update_task_config(
|
||||
request: Request, form_data: TaskConfigForm, user=Depends(get_admin_user)
|
||||
):
|
||||
request.app.state.config.TASK_MODEL = form_data.TASK_MODEL
|
||||
request.app.state.config.TASK_MODEL_EXTERNAL = form_data.TASK_MODEL_EXTERNAL
|
||||
request.app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE = (
|
||||
form_data.TITLE_GENERATION_PROMPT_TEMPLATE
|
||||
)
|
||||
|
||||
request.app.state.config.ENABLE_AUTOCOMPLETE_GENERATION = (
|
||||
form_data.ENABLE_AUTOCOMPLETE_GENERATION
|
||||
)
|
||||
request.app.state.config.AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH = (
|
||||
form_data.AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH
|
||||
)
|
||||
|
||||
request.app.state.config.TAGS_GENERATION_PROMPT_TEMPLATE = (
|
||||
form_data.TAGS_GENERATION_PROMPT_TEMPLATE
|
||||
)
|
||||
request.app.state.config.ENABLE_TAGS_GENERATION = form_data.ENABLE_TAGS_GENERATION
|
||||
request.app.state.config.ENABLE_SEARCH_QUERY_GENERATION = (
|
||||
form_data.ENABLE_SEARCH_QUERY_GENERATION
|
||||
)
|
||||
request.app.state.config.ENABLE_RETRIEVAL_QUERY_GENERATION = (
|
||||
form_data.ENABLE_RETRIEVAL_QUERY_GENERATION
|
||||
)
|
||||
|
||||
request.app.state.config.QUERY_GENERATION_PROMPT_TEMPLATE = (
|
||||
form_data.QUERY_GENERATION_PROMPT_TEMPLATE
|
||||
)
|
||||
request.app.state.config.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE = (
|
||||
form_data.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE
|
||||
)
|
||||
|
||||
return {
|
||||
"TASK_MODEL": request.app.state.config.TASK_MODEL,
|
||||
"TASK_MODEL_EXTERNAL": request.app.state.config.TASK_MODEL_EXTERNAL,
|
||||
"TITLE_GENERATION_PROMPT_TEMPLATE": request.app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE,
|
||||
"ENABLE_AUTOCOMPLETE_GENERATION": request.app.state.config.ENABLE_AUTOCOMPLETE_GENERATION,
|
||||
"AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH": request.app.state.config.AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH,
|
||||
"TAGS_GENERATION_PROMPT_TEMPLATE": request.app.state.config.TAGS_GENERATION_PROMPT_TEMPLATE,
|
||||
"ENABLE_TAGS_GENERATION": request.app.state.config.ENABLE_TAGS_GENERATION,
|
||||
"ENABLE_SEARCH_QUERY_GENERATION": request.app.state.config.ENABLE_SEARCH_QUERY_GENERATION,
|
||||
"ENABLE_RETRIEVAL_QUERY_GENERATION": request.app.state.config.ENABLE_RETRIEVAL_QUERY_GENERATION,
|
||||
"QUERY_GENERATION_PROMPT_TEMPLATE": request.app.state.config.QUERY_GENERATION_PROMPT_TEMPLATE,
|
||||
"TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE": request.app.state.config.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE,
|
||||
}
|
||||
|
||||
|
||||
@router.post("/title/completions")
|
||||
async def generate_title(
|
||||
request: Request, form_data: dict, user=Depends(get_verified_user)
|
||||
):
|
||||
models = request.app.state.MODELS
|
||||
|
||||
model_id = form_data["model"]
|
||||
if model_id not in models:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Model not found",
|
||||
)
|
||||
|
||||
# Check if the user has a custom task model
|
||||
# If the user has a custom task model, use that model
|
||||
task_model_id = get_task_model_id(
|
||||
model_id,
|
||||
request.app.state.config.TASK_MODEL,
|
||||
request.app.state.config.TASK_MODEL_EXTERNAL,
|
||||
models,
|
||||
)
|
||||
|
||||
log.debug(
|
||||
f"generating chat title using model {task_model_id} for user {user.email} "
|
||||
)
|
||||
|
||||
if request.app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE != "":
|
||||
template = request.app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE
|
||||
else:
|
||||
template = DEFAULT_TITLE_GENERATION_PROMPT_TEMPLATE
|
||||
|
||||
content = title_generation_template(
|
||||
template,
|
||||
form_data["messages"],
|
||||
{
|
||||
"name": user.name,
|
||||
"location": user.info.get("location") if user.info else None,
|
||||
},
|
||||
)
|
||||
|
||||
payload = {
|
||||
"model": task_model_id,
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
"stream": False,
|
||||
**(
|
||||
{"max_tokens": 50}
|
||||
if models[task_model_id]["owned_by"] == "ollama"
|
||||
else {
|
||||
"max_completion_tokens": 50,
|
||||
}
|
||||
),
|
||||
"metadata": {
|
||||
"task": str(TASKS.TITLE_GENERATION),
|
||||
"task_body": form_data,
|
||||
"chat_id": form_data.get("chat_id", None),
|
||||
},
|
||||
}
|
||||
|
||||
try:
|
||||
return await generate_chat_completion(request, form_data=payload, user=user)
|
||||
except Exception as e:
|
||||
return JSONResponse(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
content={"detail": str(e)},
|
||||
)
|
||||
|
||||
|
||||
@router.post("/tags/completions")
|
||||
async def generate_chat_tags(
|
||||
request: Request, form_data: dict, user=Depends(get_verified_user)
|
||||
):
|
||||
|
||||
if not request.app.state.config.ENABLE_TAGS_GENERATION:
|
||||
return JSONResponse(
|
||||
status_code=status.HTTP_200_OK,
|
||||
content={"detail": "Tags generation is disabled"},
|
||||
)
|
||||
|
||||
models = request.app.state.MODELS
|
||||
|
||||
model_id = form_data["model"]
|
||||
if model_id not in models:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Model not found",
|
||||
)
|
||||
|
||||
# Check if the user has a custom task model
|
||||
# If the user has a custom task model, use that model
|
||||
task_model_id = get_task_model_id(
|
||||
model_id,
|
||||
request.app.state.config.TASK_MODEL,
|
||||
request.app.state.config.TASK_MODEL_EXTERNAL,
|
||||
models,
|
||||
)
|
||||
|
||||
log.debug(
|
||||
f"generating chat tags using model {task_model_id} for user {user.email} "
|
||||
)
|
||||
|
||||
if request.app.state.config.TAGS_GENERATION_PROMPT_TEMPLATE != "":
|
||||
template = request.app.state.config.TAGS_GENERATION_PROMPT_TEMPLATE
|
||||
else:
|
||||
template = DEFAULT_TAGS_GENERATION_PROMPT_TEMPLATE
|
||||
|
||||
content = tags_generation_template(
|
||||
template, form_data["messages"], {"name": user.name}
|
||||
)
|
||||
|
||||
payload = {
|
||||
"model": task_model_id,
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
"stream": False,
|
||||
"metadata": {
|
||||
"task": str(TASKS.TAGS_GENERATION),
|
||||
"task_body": form_data,
|
||||
"chat_id": form_data.get("chat_id", None),
|
||||
},
|
||||
}
|
||||
|
||||
try:
|
||||
return await generate_chat_completion(request, form_data=payload, user=user)
|
||||
except Exception as e:
|
||||
return JSONResponse(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
content={"detail": str(e)},
|
||||
)
|
||||
|
||||
|
||||
@router.post("/queries/completions")
|
||||
async def generate_queries(
|
||||
request: Request, form_data: dict, user=Depends(get_verified_user)
|
||||
):
|
||||
|
||||
type = form_data.get("type")
|
||||
if type == "web_search":
|
||||
if not request.app.state.config.ENABLE_SEARCH_QUERY_GENERATION:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=f"Search query generation is disabled",
|
||||
)
|
||||
elif type == "retrieval":
|
||||
if not request.app.state.config.ENABLE_RETRIEVAL_QUERY_GENERATION:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=f"Query generation is disabled",
|
||||
)
|
||||
|
||||
models = request.app.state.MODELS
|
||||
|
||||
model_id = form_data["model"]
|
||||
if model_id not in models:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Model not found",
|
||||
)
|
||||
|
||||
# Check if the user has a custom task model
|
||||
# If the user has a custom task model, use that model
|
||||
task_model_id = get_task_model_id(
|
||||
model_id,
|
||||
request.app.state.config.TASK_MODEL,
|
||||
request.app.state.config.TASK_MODEL_EXTERNAL,
|
||||
models,
|
||||
)
|
||||
|
||||
log.debug(
|
||||
f"generating {type} queries using model {task_model_id} for user {user.email}"
|
||||
)
|
||||
|
||||
if (request.app.state.config.QUERY_GENERATION_PROMPT_TEMPLATE).strip() != "":
|
||||
template = request.app.state.config.QUERY_GENERATION_PROMPT_TEMPLATE
|
||||
else:
|
||||
template = DEFAULT_QUERY_GENERATION_PROMPT_TEMPLATE
|
||||
|
||||
content = query_generation_template(
|
||||
template, form_data["messages"], {"name": user.name}
|
||||
)
|
||||
|
||||
payload = {
|
||||
"model": task_model_id,
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
"stream": False,
|
||||
"metadata": {
|
||||
"task": str(TASKS.QUERY_GENERATION),
|
||||
"task_body": form_data,
|
||||
"chat_id": form_data.get("chat_id", None),
|
||||
},
|
||||
}
|
||||
|
||||
try:
|
||||
return await generate_chat_completion(request, form_data=payload, user=user)
|
||||
except Exception as e:
|
||||
return JSONResponse(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
content={"detail": str(e)},
|
||||
)
|
||||
|
||||
|
||||
@router.post("/auto/completions")
|
||||
async def generate_autocompletion(
|
||||
request: Request, form_data: dict, user=Depends(get_verified_user)
|
||||
):
|
||||
if not request.app.state.config.ENABLE_AUTOCOMPLETE_GENERATION:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=f"Autocompletion generation is disabled",
|
||||
)
|
||||
|
||||
type = form_data.get("type")
|
||||
prompt = form_data.get("prompt")
|
||||
messages = form_data.get("messages")
|
||||
|
||||
if request.app.state.config.AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH > 0:
|
||||
if (
|
||||
len(prompt)
|
||||
> request.app.state.config.AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=f"Input prompt exceeds maximum length of {request.app.state.config.AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH}",
|
||||
)
|
||||
|
||||
models = request.app.state.MODELS
|
||||
|
||||
model_id = form_data["model"]
|
||||
if model_id not in models:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Model not found",
|
||||
)
|
||||
|
||||
# Check if the user has a custom task model
|
||||
# If the user has a custom task model, use that model
|
||||
task_model_id = get_task_model_id(
|
||||
model_id,
|
||||
request.app.state.config.TASK_MODEL,
|
||||
request.app.state.config.TASK_MODEL_EXTERNAL,
|
||||
models,
|
||||
)
|
||||
|
||||
log.debug(
|
||||
f"generating autocompletion using model {task_model_id} for user {user.email}"
|
||||
)
|
||||
|
||||
if (request.app.state.config.AUTOCOMPLETE_GENERATION_PROMPT_TEMPLATE).strip() != "":
|
||||
template = request.app.state.config.AUTOCOMPLETE_GENERATION_PROMPT_TEMPLATE
|
||||
else:
|
||||
template = DEFAULT_AUTOCOMPLETE_GENERATION_PROMPT_TEMPLATE
|
||||
|
||||
content = autocomplete_generation_template(
|
||||
template, prompt, messages, type, {"name": user.name}
|
||||
)
|
||||
|
||||
payload = {
|
||||
"model": task_model_id,
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
"stream": False,
|
||||
"metadata": {
|
||||
"task": str(TASKS.AUTOCOMPLETE_GENERATION),
|
||||
"task_body": form_data,
|
||||
"chat_id": form_data.get("chat_id", None),
|
||||
},
|
||||
}
|
||||
|
||||
try:
|
||||
return await generate_chat_completion(request, form_data=payload, user=user)
|
||||
except Exception as e:
|
||||
return JSONResponse(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
content={"detail": str(e)},
|
||||
)
|
||||
|
||||
|
||||
@router.post("/emoji/completions")
|
||||
async def generate_emoji(
|
||||
request: Request, form_data: dict, user=Depends(get_verified_user)
|
||||
):
|
||||
|
||||
models = request.app.state.MODELS
|
||||
|
||||
model_id = form_data["model"]
|
||||
if model_id not in models:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Model not found",
|
||||
)
|
||||
|
||||
# Check if the user has a custom task model
|
||||
# If the user has a custom task model, use that model
|
||||
task_model_id = get_task_model_id(
|
||||
model_id,
|
||||
request.app.state.config.TASK_MODEL,
|
||||
request.app.state.config.TASK_MODEL_EXTERNAL,
|
||||
models,
|
||||
)
|
||||
|
||||
log.debug(f"generating emoji using model {task_model_id} for user {user.email} ")
|
||||
|
||||
template = DEFAULT_EMOJI_GENERATION_PROMPT_TEMPLATE
|
||||
|
||||
content = emoji_generation_template(
|
||||
template,
|
||||
form_data["prompt"],
|
||||
{
|
||||
"name": user.name,
|
||||
"location": user.info.get("location") if user.info else None,
|
||||
},
|
||||
)
|
||||
|
||||
payload = {
|
||||
"model": task_model_id,
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
"stream": False,
|
||||
**(
|
||||
{"max_tokens": 4}
|
||||
if models[task_model_id]["owned_by"] == "ollama"
|
||||
else {
|
||||
"max_completion_tokens": 4,
|
||||
}
|
||||
),
|
||||
"chat_id": form_data.get("chat_id", None),
|
||||
"metadata": {"task": str(TASKS.EMOJI_GENERATION), "task_body": form_data},
|
||||
}
|
||||
|
||||
try:
|
||||
return await generate_chat_completion(request, form_data=payload, user=user)
|
||||
except Exception as e:
|
||||
return JSONResponse(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
content={"detail": str(e)},
|
||||
)
|
||||
|
||||
|
||||
@router.post("/moa/completions")
|
||||
async def generate_moa_response(
|
||||
request: Request, form_data: dict, user=Depends(get_verified_user)
|
||||
):
|
||||
|
||||
models = request.app.state.MODELS
|
||||
model_id = form_data["model"]
|
||||
|
||||
if model_id not in models:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Model not found",
|
||||
)
|
||||
|
||||
# Check if the user has a custom task model
|
||||
# If the user has a custom task model, use that model
|
||||
task_model_id = get_task_model_id(
|
||||
model_id,
|
||||
request.app.state.config.TASK_MODEL,
|
||||
request.app.state.config.TASK_MODEL_EXTERNAL,
|
||||
models,
|
||||
)
|
||||
|
||||
log.debug(f"generating MOA model {task_model_id} for user {user.email} ")
|
||||
|
||||
template = DEFAULT_MOA_GENERATION_PROMPT_TEMPLATE
|
||||
|
||||
content = moa_response_generation_template(
|
||||
template,
|
||||
form_data["prompt"],
|
||||
form_data["responses"],
|
||||
)
|
||||
|
||||
payload = {
|
||||
"model": task_model_id,
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
"stream": form_data.get("stream", False),
|
||||
"chat_id": form_data.get("chat_id", None),
|
||||
"metadata": {
|
||||
"task": str(TASKS.MOA_RESPONSE_GENERATION),
|
||||
"task_body": form_data,
|
||||
},
|
||||
}
|
||||
|
||||
try:
|
||||
return await generate_chat_completion(request, form_data=payload, user=user)
|
||||
except Exception as e:
|
||||
return JSONResponse(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
content={"detail": str(e)},
|
||||
)
|
||||
@@ -0,0 +1,418 @@
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from open_webui.models.tools import (
|
||||
ToolForm,
|
||||
ToolModel,
|
||||
ToolResponse,
|
||||
ToolUserResponse,
|
||||
Tools,
|
||||
)
|
||||
from open_webui.utils.plugin import load_tools_module_by_id, replace_imports
|
||||
from open_webui.config import CACHE_DIR
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, status
|
||||
from open_webui.utils.tools import get_tools_specs
|
||||
from open_webui.utils.auth import get_admin_user, get_verified_user
|
||||
from open_webui.utils.access_control import has_access, has_permission
|
||||
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
############################
|
||||
# GetTools
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/", response_model=list[ToolUserResponse])
|
||||
async def get_tools(user=Depends(get_verified_user)):
|
||||
if user.role == "admin":
|
||||
tools = Tools.get_tools()
|
||||
else:
|
||||
tools = Tools.get_tools_by_user_id(user.id, "read")
|
||||
return tools
|
||||
|
||||
|
||||
############################
|
||||
# GetToolList
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/list", response_model=list[ToolUserResponse])
|
||||
async def get_tool_list(user=Depends(get_verified_user)):
|
||||
if user.role == "admin":
|
||||
tools = Tools.get_tools()
|
||||
else:
|
||||
tools = Tools.get_tools_by_user_id(user.id, "write")
|
||||
return tools
|
||||
|
||||
|
||||
############################
|
||||
# ExportTools
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/export", response_model=list[ToolModel])
|
||||
async def export_tools(user=Depends(get_admin_user)):
|
||||
tools = Tools.get_tools()
|
||||
return tools
|
||||
|
||||
|
||||
############################
|
||||
# CreateNewTools
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/create", response_model=Optional[ToolResponse])
|
||||
async def create_new_tools(
|
||||
request: Request,
|
||||
form_data: ToolForm,
|
||||
user=Depends(get_verified_user),
|
||||
):
|
||||
if user.role != "admin" and not has_permission(
|
||||
user.id, "workspace.knowledge", request.app.state.config.USER_PERMISSIONS
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.UNAUTHORIZED,
|
||||
)
|
||||
|
||||
if not form_data.id.isidentifier():
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Only alphanumeric characters and underscores are allowed in the id",
|
||||
)
|
||||
|
||||
form_data.id = form_data.id.lower()
|
||||
|
||||
tools = Tools.get_tool_by_id(form_data.id)
|
||||
if tools is None:
|
||||
try:
|
||||
form_data.content = replace_imports(form_data.content)
|
||||
tools_module, frontmatter = load_tools_module_by_id(
|
||||
form_data.id, content=form_data.content
|
||||
)
|
||||
form_data.meta.manifest = frontmatter
|
||||
|
||||
TOOLS = request.app.state.TOOLS
|
||||
TOOLS[form_data.id] = tools_module
|
||||
|
||||
specs = get_tools_specs(TOOLS[form_data.id])
|
||||
tools = Tools.insert_new_tool(user.id, form_data, specs)
|
||||
|
||||
tool_cache_dir = Path(CACHE_DIR) / "tools" / form_data.id
|
||||
tool_cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if tools:
|
||||
return tools
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error creating tools"),
|
||||
)
|
||||
except Exception as e:
|
||||
print(e)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(str(e)),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.ID_TAKEN,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# GetToolsById
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/id/{id}", response_model=Optional[ToolModel])
|
||||
async def get_tools_by_id(id: str, user=Depends(get_verified_user)):
|
||||
tools = Tools.get_tool_by_id(id)
|
||||
|
||||
if tools:
|
||||
if (
|
||||
user.role == "admin"
|
||||
or tools.user_id == user.id
|
||||
or has_access(user.id, "read", tools.access_control)
|
||||
):
|
||||
return tools
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# UpdateToolsById
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/id/{id}/update", response_model=Optional[ToolModel])
|
||||
async def update_tools_by_id(
|
||||
request: Request,
|
||||
id: str,
|
||||
form_data: ToolForm,
|
||||
user=Depends(get_verified_user),
|
||||
):
|
||||
tools = Tools.get_tool_by_id(id)
|
||||
if not tools:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
if tools.user_id != user.id and user.role != "admin":
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.UNAUTHORIZED,
|
||||
)
|
||||
|
||||
try:
|
||||
form_data.content = replace_imports(form_data.content)
|
||||
tools_module, frontmatter = load_tools_module_by_id(
|
||||
id, content=form_data.content
|
||||
)
|
||||
form_data.meta.manifest = frontmatter
|
||||
|
||||
TOOLS = request.app.state.TOOLS
|
||||
TOOLS[id] = tools_module
|
||||
|
||||
specs = get_tools_specs(TOOLS[id])
|
||||
|
||||
updated = {
|
||||
**form_data.model_dump(exclude={"id"}),
|
||||
"specs": specs,
|
||||
}
|
||||
|
||||
print(updated)
|
||||
tools = Tools.update_tool_by_id(id, updated)
|
||||
|
||||
if tools:
|
||||
return tools
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error updating tools"),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(str(e)),
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# DeleteToolsById
|
||||
############################
|
||||
|
||||
|
||||
@router.delete("/id/{id}/delete", response_model=bool)
|
||||
async def delete_tools_by_id(
|
||||
request: Request, id: str, user=Depends(get_verified_user)
|
||||
):
|
||||
tools = Tools.get_tool_by_id(id)
|
||||
if not tools:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
if tools.user_id != user.id and user.role != "admin":
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.UNAUTHORIZED,
|
||||
)
|
||||
|
||||
result = Tools.delete_tool_by_id(id)
|
||||
if result:
|
||||
TOOLS = request.app.state.TOOLS
|
||||
if id in TOOLS:
|
||||
del TOOLS[id]
|
||||
|
||||
return result
|
||||
|
||||
|
||||
############################
|
||||
# GetToolValves
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/id/{id}/valves", response_model=Optional[dict])
|
||||
async def get_tools_valves_by_id(id: str, user=Depends(get_verified_user)):
|
||||
tools = Tools.get_tool_by_id(id)
|
||||
if tools:
|
||||
try:
|
||||
valves = Tools.get_tool_valves_by_id(id)
|
||||
return valves
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(str(e)),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# GetToolValvesSpec
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/id/{id}/valves/spec", response_model=Optional[dict])
|
||||
async def get_tools_valves_spec_by_id(
|
||||
request: Request, id: str, user=Depends(get_verified_user)
|
||||
):
|
||||
tools = Tools.get_tool_by_id(id)
|
||||
if tools:
|
||||
if id in request.app.state.TOOLS:
|
||||
tools_module = request.app.state.TOOLS[id]
|
||||
else:
|
||||
tools_module, _ = load_tools_module_by_id(id)
|
||||
request.app.state.TOOLS[id] = tools_module
|
||||
|
||||
if hasattr(tools_module, "Valves"):
|
||||
Valves = tools_module.Valves
|
||||
return Valves.schema()
|
||||
return None
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# UpdateToolValves
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/id/{id}/valves/update", response_model=Optional[dict])
|
||||
async def update_tools_valves_by_id(
|
||||
request: Request, id: str, form_data: dict, user=Depends(get_verified_user)
|
||||
):
|
||||
tools = Tools.get_tool_by_id(id)
|
||||
if not tools:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
if id in request.app.state.TOOLS:
|
||||
tools_module = request.app.state.TOOLS[id]
|
||||
else:
|
||||
tools_module, _ = load_tools_module_by_id(id)
|
||||
request.app.state.TOOLS[id] = tools_module
|
||||
|
||||
if not hasattr(tools_module, "Valves"):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
Valves = tools_module.Valves
|
||||
|
||||
try:
|
||||
form_data = {k: v for k, v in form_data.items() if v is not None}
|
||||
valves = Valves(**form_data)
|
||||
Tools.update_tool_valves_by_id(id, valves.model_dump())
|
||||
return valves.model_dump()
|
||||
except Exception as e:
|
||||
print(e)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(str(e)),
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# ToolUserValves
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/id/{id}/valves/user", response_model=Optional[dict])
|
||||
async def get_tools_user_valves_by_id(id: str, user=Depends(get_verified_user)):
|
||||
tools = Tools.get_tool_by_id(id)
|
||||
if tools:
|
||||
try:
|
||||
user_valves = Tools.get_user_valves_by_id_and_user_id(id, user.id)
|
||||
return user_valves
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(str(e)),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
@router.get("/id/{id}/valves/user/spec", response_model=Optional[dict])
|
||||
async def get_tools_user_valves_spec_by_id(
|
||||
request: Request, id: str, user=Depends(get_verified_user)
|
||||
):
|
||||
tools = Tools.get_tool_by_id(id)
|
||||
if tools:
|
||||
if id in request.app.state.TOOLS:
|
||||
tools_module = request.app.state.TOOLS[id]
|
||||
else:
|
||||
tools_module, _ = load_tools_module_by_id(id)
|
||||
request.app.state.TOOLS[id] = tools_module
|
||||
|
||||
if hasattr(tools_module, "UserValves"):
|
||||
UserValves = tools_module.UserValves
|
||||
return UserValves.schema()
|
||||
return None
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
@router.post("/id/{id}/valves/user/update", response_model=Optional[dict])
|
||||
async def update_tools_user_valves_by_id(
|
||||
request: Request, id: str, form_data: dict, user=Depends(get_verified_user)
|
||||
):
|
||||
tools = Tools.get_tool_by_id(id)
|
||||
|
||||
if tools:
|
||||
if id in request.app.state.TOOLS:
|
||||
tools_module = request.app.state.TOOLS[id]
|
||||
else:
|
||||
tools_module, _ = load_tools_module_by_id(id)
|
||||
request.app.state.TOOLS[id] = tools_module
|
||||
|
||||
if hasattr(tools_module, "UserValves"):
|
||||
UserValves = tools_module.UserValves
|
||||
|
||||
try:
|
||||
form_data = {k: v for k, v in form_data.items() if v is not None}
|
||||
user_valves = UserValves(**form_data)
|
||||
Tools.update_user_valves_by_id_and_user_id(
|
||||
id, user.id, user_valves.model_dump()
|
||||
)
|
||||
return user_valves.model_dump()
|
||||
except Exception as e:
|
||||
print(e)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(str(e)),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
@@ -0,0 +1,295 @@
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from open_webui.models.auths import Auths
|
||||
from open_webui.models.chats import Chats
|
||||
from open_webui.models.users import (
|
||||
UserModel,
|
||||
UserRoleUpdateForm,
|
||||
Users,
|
||||
UserSettings,
|
||||
UserUpdateForm,
|
||||
)
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
from open_webui.env import SRC_LOG_LEVELS
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, status
|
||||
from pydantic import BaseModel
|
||||
from open_webui.utils.auth import get_admin_user, get_password_hash, get_verified_user
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["MODELS"])
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
############################
|
||||
# GetUsers
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/", response_model=list[UserModel])
|
||||
async def get_users(skip: int = 0, limit: int = 50, user=Depends(get_admin_user)):
|
||||
return Users.get_users(skip, limit)
|
||||
|
||||
|
||||
############################
|
||||
# User Groups
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/groups")
|
||||
async def get_user_groups(user=Depends(get_verified_user)):
|
||||
return Users.get_user_groups(user.id)
|
||||
|
||||
|
||||
############################
|
||||
# User Permissions
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/permissions")
|
||||
async def get_user_permissisions(user=Depends(get_verified_user)):
|
||||
return Users.get_user_groups(user.id)
|
||||
|
||||
|
||||
############################
|
||||
# User Default Permissions
|
||||
############################
|
||||
class WorkspacePermissions(BaseModel):
|
||||
models: bool
|
||||
knowledge: bool
|
||||
prompts: bool
|
||||
tools: bool
|
||||
|
||||
|
||||
class ChatPermissions(BaseModel):
|
||||
file_upload: bool
|
||||
delete: bool
|
||||
edit: bool
|
||||
temporary: bool
|
||||
|
||||
|
||||
class UserPermissions(BaseModel):
|
||||
workspace: WorkspacePermissions
|
||||
chat: ChatPermissions
|
||||
|
||||
|
||||
@router.get("/default/permissions")
|
||||
async def get_user_permissions(request: Request, user=Depends(get_admin_user)):
|
||||
return request.app.state.config.USER_PERMISSIONS
|
||||
|
||||
|
||||
@router.post("/default/permissions")
|
||||
async def update_user_permissions(
|
||||
request: Request, form_data: UserPermissions, user=Depends(get_admin_user)
|
||||
):
|
||||
request.app.state.config.USER_PERMISSIONS = form_data.model_dump()
|
||||
return request.app.state.config.USER_PERMISSIONS
|
||||
|
||||
|
||||
############################
|
||||
# UpdateUserRole
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/update/role", response_model=Optional[UserModel])
|
||||
async def update_user_role(form_data: UserRoleUpdateForm, user=Depends(get_admin_user)):
|
||||
if user.id != form_data.id and form_data.id != Users.get_first_user().id:
|
||||
return Users.update_user_role_by_id(form_data.id, form_data.role)
|
||||
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail=ERROR_MESSAGES.ACTION_PROHIBITED,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# GetUserSettingsBySessionUser
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/user/settings", response_model=Optional[UserSettings])
|
||||
async def get_user_settings_by_session_user(user=Depends(get_verified_user)):
|
||||
user = Users.get_user_by_id(user.id)
|
||||
if user:
|
||||
return user.settings
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.USER_NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# UpdateUserSettingsBySessionUser
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/user/settings/update", response_model=UserSettings)
|
||||
async def update_user_settings_by_session_user(
|
||||
form_data: UserSettings, user=Depends(get_verified_user)
|
||||
):
|
||||
user = Users.update_user_by_id(user.id, {"settings": form_data.model_dump()})
|
||||
if user:
|
||||
return user.settings
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.USER_NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# GetUserInfoBySessionUser
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/user/info", response_model=Optional[dict])
|
||||
async def get_user_info_by_session_user(user=Depends(get_verified_user)):
|
||||
user = Users.get_user_by_id(user.id)
|
||||
if user:
|
||||
return user.info
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.USER_NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# UpdateUserInfoBySessionUser
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/user/info/update", response_model=Optional[dict])
|
||||
async def update_user_info_by_session_user(
|
||||
form_data: dict, user=Depends(get_verified_user)
|
||||
):
|
||||
user = Users.get_user_by_id(user.id)
|
||||
if user:
|
||||
if user.info is None:
|
||||
user.info = {}
|
||||
|
||||
user = Users.update_user_by_id(user.id, {"info": {**user.info, **form_data}})
|
||||
if user:
|
||||
return user.info
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.USER_NOT_FOUND,
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.USER_NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# GetUserById
|
||||
############################
|
||||
|
||||
|
||||
class UserResponse(BaseModel):
|
||||
name: str
|
||||
profile_image_url: str
|
||||
|
||||
|
||||
@router.get("/{user_id}", response_model=UserResponse)
|
||||
async def get_user_by_id(user_id: str, user=Depends(get_verified_user)):
|
||||
# Check if user_id is a shared chat
|
||||
# If it is, get the user_id from the chat
|
||||
if user_id.startswith("shared-"):
|
||||
chat_id = user_id.replace("shared-", "")
|
||||
chat = Chats.get_chat_by_id(chat_id)
|
||||
if chat:
|
||||
user_id = chat.user_id
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.USER_NOT_FOUND,
|
||||
)
|
||||
|
||||
user = Users.get_user_by_id(user_id)
|
||||
|
||||
if user:
|
||||
return UserResponse(name=user.name, profile_image_url=user.profile_image_url)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.USER_NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# UpdateUserById
|
||||
############################
|
||||
|
||||
|
||||
@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 = Users.get_user_by_id(user_id)
|
||||
|
||||
if user:
|
||||
if form_data.email.lower() != user.email:
|
||||
email_user = Users.get_user_by_email(form_data.email.lower())
|
||||
if email_user:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.EMAIL_TAKEN,
|
||||
)
|
||||
|
||||
if form_data.password:
|
||||
hashed = get_password_hash(form_data.password)
|
||||
log.debug(f"hashed: {hashed}")
|
||||
Auths.update_user_password_by_id(user_id, hashed)
|
||||
|
||||
Auths.update_email_by_id(user_id, form_data.email.lower())
|
||||
updated_user = Users.update_user_by_id(
|
||||
user_id,
|
||||
{
|
||||
"name": form_data.name,
|
||||
"email": form_data.email.lower(),
|
||||
"profile_image_url": form_data.profile_image_url,
|
||||
},
|
||||
)
|
||||
|
||||
if updated_user:
|
||||
return updated_user
|
||||
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(),
|
||||
)
|
||||
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.USER_NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# DeleteUserById
|
||||
############################
|
||||
|
||||
|
||||
@router.delete("/{user_id}", response_model=bool)
|
||||
async def delete_user_by_id(user_id: str, user=Depends(get_admin_user)):
|
||||
if user.id != user_id:
|
||||
result = Auths.delete_auth_by_id(user_id)
|
||||
|
||||
if result:
|
||||
return True
|
||||
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail=ERROR_MESSAGES.DELETE_USER_ERROR,
|
||||
)
|
||||
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail=ERROR_MESSAGES.ACTION_PROHIBITED,
|
||||
)
|
||||
@@ -0,0 +1,99 @@
|
||||
import black
|
||||
import markdown
|
||||
|
||||
from open_webui.models.chats import ChatTitleMessagesForm
|
||||
from open_webui.config import DATA_DIR, ENABLE_ADMIN_EXPORT
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
from fastapi import APIRouter, Depends, HTTPException, Response, status
|
||||
from pydantic import BaseModel
|
||||
from starlette.responses import FileResponse
|
||||
from open_webui.utils.misc import get_gravatar_url
|
||||
from open_webui.utils.pdf_generator import PDFGenerator
|
||||
from open_webui.utils.auth import get_admin_user
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("/gravatar")
|
||||
async def get_gravatar(
|
||||
email: str,
|
||||
):
|
||||
return get_gravatar_url(email)
|
||||
|
||||
|
||||
class CodeFormatRequest(BaseModel):
|
||||
code: str
|
||||
|
||||
|
||||
@router.post("/code/format")
|
||||
async def format_code(request: CodeFormatRequest):
|
||||
try:
|
||||
formatted_code = black.format_str(request.code, mode=black.Mode())
|
||||
return {"code": formatted_code}
|
||||
except black.NothingChanged:
|
||||
return {"code": request.code}
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
|
||||
|
||||
class MarkdownForm(BaseModel):
|
||||
md: str
|
||||
|
||||
|
||||
@router.post("/markdown")
|
||||
async def get_html_from_markdown(
|
||||
form_data: MarkdownForm,
|
||||
):
|
||||
return {"html": markdown.markdown(form_data.md)}
|
||||
|
||||
|
||||
class ChatForm(BaseModel):
|
||||
title: str
|
||||
messages: list[dict]
|
||||
|
||||
|
||||
@router.post("/pdf")
|
||||
async def download_chat_as_pdf(
|
||||
form_data: ChatTitleMessagesForm,
|
||||
):
|
||||
try:
|
||||
pdf_bytes = PDFGenerator(form_data).generate_chat_pdf()
|
||||
|
||||
return Response(
|
||||
content=pdf_bytes,
|
||||
media_type="application/pdf",
|
||||
headers={"Content-Disposition": "attachment;filename=chat.pdf"},
|
||||
)
|
||||
except Exception as e:
|
||||
print(e)
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/db/download")
|
||||
async def download_db(user=Depends(get_admin_user)):
|
||||
if not ENABLE_ADMIN_EXPORT:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
from open_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(
|
||||
engine.url.database,
|
||||
media_type="application/octet-stream",
|
||||
filename="webui.db",
|
||||
)
|
||||
|
||||
|
||||
@router.get("/litellm/config")
|
||||
async def download_litellm_config_yaml(user=Depends(get_admin_user)):
|
||||
return FileResponse(
|
||||
f"{DATA_DIR}/litellm/config.yaml",
|
||||
media_type="application/octet-stream",
|
||||
filename="config.yaml",
|
||||
)
|
||||
Reference in New Issue
Block a user