Merge branch 'origin/dev' into sidebar-pagination [skip ci]

This commit is contained in:
Aryan Kothari
2024-08-03 09:56:23 -04:00
80 changed files with 3065 additions and 1802 deletions

View File

@@ -10,12 +10,12 @@ from fastapi import (
File,
Form,
)
from fastapi.responses import StreamingResponse, JSONResponse, FileResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List
import uuid
import requests
import hashlib
@@ -31,6 +31,7 @@ from utils.utils import (
)
from utils.misc import calculate_sha256
from config import (
SRC_LOG_LEVELS,
CACHE_DIR,
@@ -252,15 +253,15 @@ async def speech(request: Request, user=Depends(get_verified_user)):
)
elif app.state.config.TTS_ENGINE == "elevenlabs":
payload = None
try:
payload = json.loads(body.decode("utf-8"))
except Exception as e:
log.exception(e)
pass
raise HTTPException(status_code=400, detail="Invalid JSON payload")
url = f"https://api.elevenlabs.io/v1/text-to-speech/{payload['voice']}"
voice_id = payload.get("voice", "")
url = f"https://api.elevenlabs.io/v1/text-to-speech/{voice_id}"
headers = {
"Accept": "audio/mpeg",
@@ -435,3 +436,69 @@ def transcribe(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.DEFAULT(e),
)
def get_available_models() -> List[dict]:
if app.state.config.TTS_ENGINE == "openai":
return [{"id": "tts-1"}, {"id": "tts-1-hd"}]
elif app.state.config.TTS_ENGINE == "elevenlabs":
headers = {
"xi-api-key": app.state.config.TTS_API_KEY,
"Content-Type": "application/json",
}
try:
response = requests.get(
"https://api.elevenlabs.io/v1/models", headers=headers
)
response.raise_for_status()
models = response.json()
return [
{"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 []
@app.get("/models")
async def get_models(user=Depends(get_verified_user)):
return {"models": get_available_models()}
def get_available_voices() -> List[dict]:
if app.state.config.TTS_ENGINE == "openai":
return [
{"name": "alloy", "id": "alloy"},
{"name": "echo", "id": "echo"},
{"name": "fable", "id": "fable"},
{"name": "onyx", "id": "onyx"},
{"name": "nova", "id": "nova"},
{"name": "shimmer", "id": "shimmer"},
]
elif app.state.config.TTS_ENGINE == "elevenlabs":
headers = {
"xi-api-key": app.state.config.TTS_API_KEY,
"Content-Type": "application/json",
}
try:
response = requests.get(
"https://api.elevenlabs.io/v1/voices", headers=headers
)
response.raise_for_status()
voices_data = response.json()
voices = []
for voice in voices_data.get("voices", []):
voices.append({"name": voice["name"], "id": voice["voice_id"]})
return voices
except requests.RequestException as e:
log.error(f"Error fetching voices: {str(e)}")
return []
@app.get("/voices")
async def get_voices(user=Depends(get_verified_user)):
return {"voices": get_available_voices()}

View File

@@ -42,6 +42,9 @@ from config import (
COMFYUI_SAMPLER,
COMFYUI_SCHEDULER,
COMFYUI_SD3,
COMFYUI_FLUX,
COMFYUI_FLUX_WEIGHT_DTYPE,
COMFYUI_FLUX_FP8_CLIP,
IMAGES_OPENAI_API_BASE_URL,
IMAGES_OPENAI_API_KEY,
IMAGE_GENERATION_MODEL,
@@ -85,6 +88,9 @@ app.state.config.COMFYUI_CFG_SCALE = COMFYUI_CFG_SCALE
app.state.config.COMFYUI_SAMPLER = COMFYUI_SAMPLER
app.state.config.COMFYUI_SCHEDULER = COMFYUI_SCHEDULER
app.state.config.COMFYUI_SD3 = COMFYUI_SD3
app.state.config.COMFYUI_FLUX = COMFYUI_FLUX
app.state.config.COMFYUI_FLUX_WEIGHT_DTYPE = COMFYUI_FLUX_WEIGHT_DTYPE
app.state.config.COMFYUI_FLUX_FP8_CLIP = COMFYUI_FLUX_FP8_CLIP
def get_automatic1111_api_auth():
@@ -497,6 +503,15 @@ async def image_generations(
if app.state.config.COMFYUI_SD3 is not None:
data["sd3"] = app.state.config.COMFYUI_SD3
if app.state.config.COMFYUI_FLUX is not None:
data["flux"] = app.state.config.COMFYUI_FLUX
if app.state.config.COMFYUI_FLUX_WEIGHT_DTYPE is not None:
data["flux_weight_dtype"] = app.state.config.COMFYUI_FLUX_WEIGHT_DTYPE
if app.state.config.COMFYUI_FLUX_FP8_CLIP is not None:
data["flux_fp8_clip"] = app.state.config.COMFYUI_FLUX_FP8_CLIP
data = ImageGenerationPayload(**data)
res = comfyui_generate_image(

View File

@@ -125,6 +125,135 @@ COMFYUI_DEFAULT_PROMPT = """
}
"""
FLUX_DEFAULT_PROMPT = """
{
"5": {
"inputs": {
"width": 1024,
"height": 1024,
"batch_size": 1
},
"class_type": "EmptyLatentImage"
},
"6": {
"inputs": {
"text": "Input Text Here",
"clip": [
"11",
0
]
},
"class_type": "CLIPTextEncode"
},
"8": {
"inputs": {
"samples": [
"13",
0
],
"vae": [
"10",
0
]
},
"class_type": "VAEDecode"
},
"9": {
"inputs": {
"filename_prefix": "ComfyUI",
"images": [
"8",
0
]
},
"class_type": "SaveImage"
},
"10": {
"inputs": {
"vae_name": "ae.sft"
},
"class_type": "VAELoader"
},
"11": {
"inputs": {
"clip_name1": "clip_l.safetensors",
"clip_name2": "t5xxl_fp16.safetensors",
"type": "flux"
},
"class_type": "DualCLIPLoader"
},
"12": {
"inputs": {
"unet_name": "flux1-dev.sft",
"weight_dtype": "default"
},
"class_type": "UNETLoader"
},
"13": {
"inputs": {
"noise": [
"25",
0
],
"guider": [
"22",
0
],
"sampler": [
"16",
0
],
"sigmas": [
"17",
0
],
"latent_image": [
"5",
0
]
},
"class_type": "SamplerCustomAdvanced"
},
"16": {
"inputs": {
"sampler_name": "euler"
},
"class_type": "KSamplerSelect"
},
"17": {
"inputs": {
"scheduler": "simple",
"steps": 20,
"denoise": 1,
"model": [
"12",
0
]
},
"class_type": "BasicScheduler"
},
"22": {
"inputs": {
"model": [
"12",
0
],
"conditioning": [
"6",
0
]
},
"class_type": "BasicGuider"
},
"25": {
"inputs": {
"noise_seed": 778937779713005
},
"class_type": "RandomNoise"
}
}
"""
def queue_prompt(prompt, client_id, base_url):
log.info("queue_prompt")
@@ -194,6 +323,9 @@ class ImageGenerationPayload(BaseModel):
sampler: Optional[str] = None
scheduler: Optional[str] = None
sd3: Optional[bool] = None
flux: Optional[bool] = None
flux_weight_dtype: Optional[str] = None
flux_fp8_clip: Optional[bool] = None
def comfyui_generate_image(
@@ -215,21 +347,46 @@ def comfyui_generate_image(
if payload.sd3:
comfyui_prompt["5"]["class_type"] = "EmptySD3LatentImage"
if payload.steps:
comfyui_prompt["3"]["inputs"]["steps"] = payload.steps
comfyui_prompt["4"]["inputs"]["ckpt_name"] = model
comfyui_prompt["7"]["inputs"]["text"] = payload.negative_prompt
comfyui_prompt["3"]["inputs"]["seed"] = (
payload.seed if payload.seed else random.randint(0, 18446744073709551614)
)
# as Flux uses a completely different workflow, we must treat it specially
if payload.flux:
comfyui_prompt = json.loads(FLUX_DEFAULT_PROMPT)
comfyui_prompt["12"]["inputs"]["unet_name"] = model
comfyui_prompt["25"]["inputs"]["noise_seed"] = (
payload.seed if payload.seed else random.randint(0, 18446744073709551614)
)
if payload.sampler:
comfyui_prompt["16"]["inputs"]["sampler_name"] = payload.sampler
if payload.steps:
comfyui_prompt["17"]["inputs"]["steps"] = payload.steps
if payload.scheduler:
comfyui_prompt["17"]["inputs"]["scheduler"] = payload.scheduler
if payload.flux_weight_dtype:
comfyui_prompt["12"]["inputs"]["weight_dtype"] = payload.flux_weight_dtype
if payload.flux_fp8_clip:
comfyui_prompt["11"]["inputs"][
"clip_name2"
] = "t5xxl_fp8_e4m3fn.safetensors"
comfyui_prompt["5"]["inputs"]["batch_size"] = payload.n
comfyui_prompt["5"]["inputs"]["width"] = payload.width
comfyui_prompt["5"]["inputs"]["height"] = payload.height
# set the text prompt for our positive CLIPTextEncode
comfyui_prompt["6"]["inputs"]["text"] = payload.prompt
comfyui_prompt["7"]["inputs"]["text"] = payload.negative_prompt
if payload.steps:
comfyui_prompt["3"]["inputs"]["steps"] = payload.steps
comfyui_prompt["3"]["inputs"]["seed"] = (
payload.seed if payload.seed else random.randint(0, 18446744073709551614)
)
try:
ws = websocket.WebSocket()

View File

@@ -857,6 +857,12 @@ async def generate_chat_completion(
):
payload["options"]["top_p"] = model_info.params.get("top_p", None)
if (
model_info.params.get("min_p", None)
and payload["options"].get("min_p") is None
):
payload["options"]["min_p"] = model_info.params.get("min_p", None)
if (
model_info.params.get("use_mmap", None)
and payload["options"].get("use_mmap") is None

View File

@@ -52,7 +52,6 @@ async def user_join(sid, data):
user = Users.get_user_by_id(data["id"])
if user:
SESSION_POOL[sid] = user.id
if user.id in USER_POOL:
USER_POOL[user.id].append(sid)
@@ -80,7 +79,6 @@ def get_models_in_use():
@sio.on("usage")
async def usage(sid, data):
model_id = data["model"]
# Cancel previous callback if there is one
@@ -139,7 +137,7 @@ async def disconnect(sid):
print(f"Unknown session ID {sid} disconnected")
async def get_event_emitter(request_info):
def get_event_emitter(request_info):
async def __event_emitter__(event_data):
await sio.emit(
"chat-events",
@@ -154,7 +152,7 @@ async def get_event_emitter(request_info):
return __event_emitter__
async def get_event_call(request_info):
def get_event_call(request_info):
async def __event_call__(event_data):
response = await sio.call(
"chat-events",

View File

@@ -1,9 +1,6 @@
from fastapi import FastAPI, Depends
from fastapi.routing import APIRoute
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from starlette.middleware.sessions import SessionMiddleware
from sqlalchemy.orm import Session
from apps.webui.routers import (
auths,
users,
@@ -22,12 +19,15 @@ from apps.webui.models.functions import Functions
from apps.webui.models.models import Models
from apps.webui.utils import load_function_module_by_id
from utils.misc import stream_message_template
from utils.misc import (
openai_chat_chunk_message_template,
openai_chat_completion_message_template,
add_or_update_system_message,
)
from utils.task import prompt_template
from config import (
WEBUI_BUILD_HASH,
SHOW_ADMIN_DETAILS,
ADMIN_EMAIL,
WEBUI_AUTH,
@@ -51,11 +51,9 @@ from config import (
from apps.socket.main import get_event_call, get_event_emitter
import inspect
import uuid
import time
import json
from typing import Iterator, Generator, AsyncGenerator, Optional
from typing import Iterator, Generator, AsyncGenerator
from pydantic import BaseModel
app = FastAPI()
@@ -127,60 +125,58 @@ async def get_status():
}
def get_function_module(pipe_id: str):
# Check if function is already loaded
if pipe_id not in app.state.FUNCTIONS:
function_module, _, _ = load_function_module_by_id(pipe_id)
app.state.FUNCTIONS[pipe_id] = function_module
else:
function_module = app.state.FUNCTIONS[pipe_id]
if hasattr(function_module, "valves") and hasattr(function_module, "Valves"):
valves = Functions.get_function_valves_by_id(pipe_id)
function_module.valves = function_module.Valves(**(valves if valves else {}))
return function_module
async def get_pipe_models():
pipes = Functions.get_functions_by_type("pipe", active_only=True)
pipe_models = []
for pipe in pipes:
# Check if function is already loaded
if pipe.id not in app.state.FUNCTIONS:
function_module, function_type, frontmatter = load_function_module_by_id(
pipe.id
)
app.state.FUNCTIONS[pipe.id] = function_module
else:
function_module = app.state.FUNCTIONS[pipe.id]
if hasattr(function_module, "valves") and hasattr(function_module, "Valves"):
valves = Functions.get_function_valves_by_id(pipe.id)
function_module.valves = function_module.Valves(
**(valves if valves else {})
)
function_module = get_function_module(pipe.id)
# Check if function is a manifold
if hasattr(function_module, "type"):
if function_module.type == "manifold":
manifold_pipes = []
if hasattr(function_module, "pipes"):
manifold_pipes = []
# Check if pipes is a function or a list
if callable(function_module.pipes):
manifold_pipes = function_module.pipes()
else:
manifold_pipes = function_module.pipes
# Check if pipes is a function or a list
if callable(function_module.pipes):
manifold_pipes = function_module.pipes()
else:
manifold_pipes = function_module.pipes
for p in manifold_pipes:
manifold_pipe_id = f'{pipe.id}.{p["id"]}'
manifold_pipe_name = p["name"]
for p in manifold_pipes:
manifold_pipe_id = f'{pipe.id}.{p["id"]}'
manifold_pipe_name = p["name"]
if hasattr(function_module, "name"):
manifold_pipe_name = (
f"{function_module.name}{manifold_pipe_name}"
)
if hasattr(function_module, "name"):
manifold_pipe_name = f"{function_module.name}{manifold_pipe_name}"
pipe_flag = {"type": pipe.type}
if hasattr(function_module, "ChatValves"):
pipe_flag["valves_spec"] = function_module.ChatValves.schema()
pipe_flag = {"type": pipe.type}
if hasattr(function_module, "ChatValves"):
pipe_flag["valves_spec"] = function_module.ChatValves.schema()
pipe_models.append(
{
"id": manifold_pipe_id,
"name": manifold_pipe_name,
"object": "model",
"created": pipe.created_at,
"owned_by": "openai",
"pipe": pipe_flag,
}
)
pipe_models.append(
{
"id": manifold_pipe_id,
"name": manifold_pipe_name,
"object": "model",
"created": pipe.created_at,
"owned_by": "openai",
"pipe": pipe_flag,
}
)
else:
pipe_flag = {"type": "pipe"}
if hasattr(function_module, "ChatValves"):
@@ -200,284 +196,211 @@ async def get_pipe_models():
return pipe_models
async def execute_pipe(pipe, params):
if inspect.iscoroutinefunction(pipe):
return await pipe(**params)
else:
return pipe(**params)
async def get_message_content(res: str | Generator | AsyncGenerator) -> str:
if isinstance(res, str):
return res
if isinstance(res, Generator):
return "".join(map(str, res))
if isinstance(res, AsyncGenerator):
return "".join([str(stream) async for stream in res])
def process_line(form_data: dict, line):
if isinstance(line, BaseModel):
line = line.model_dump_json()
line = f"data: {line}"
if isinstance(line, dict):
line = f"data: {json.dumps(line)}"
try:
line = line.decode("utf-8")
except Exception:
pass
if line.startswith("data:"):
return f"{line}\n\n"
else:
line = openai_chat_chunk_message_template(form_data["model"], line)
return f"data: {json.dumps(line)}\n\n"
def get_pipe_id(form_data: dict) -> str:
pipe_id = form_data["model"]
if "." in pipe_id:
pipe_id, _ = pipe_id.split(".", 1)
print(pipe_id)
return pipe_id
def get_function_params(function_module, form_data, user, extra_params={}):
pipe_id = get_pipe_id(form_data)
# Get the signature of the function
sig = inspect.signature(function_module.pipe)
params = {"body": form_data}
for key, value in extra_params.items():
if key in sig.parameters:
params[key] = value
if "__user__" in sig.parameters:
__user__ = {
"id": user.id,
"email": user.email,
"name": user.name,
"role": user.role,
}
try:
if hasattr(function_module, "UserValves"):
__user__["valves"] = function_module.UserValves(
**Functions.get_user_valves_by_id_and_user_id(pipe_id, user.id)
)
except Exception as e:
print(e)
params["__user__"] = __user__
return params
# inplace function: form_data is modified
def apply_model_params_to_body(params: dict, form_data: dict) -> dict:
if not params:
return form_data
mappings = {
"temperature": float,
"top_p": int,
"max_tokens": int,
"frequency_penalty": int,
"seed": lambda x: x,
"stop": lambda x: [bytes(s, "utf-8").decode("unicode_escape") for s in x],
}
for key, cast_func in mappings.items():
if (value := params.get(key)) is not None:
form_data[key] = cast_func(value)
return form_data
# inplace function: form_data is modified
def apply_model_system_prompt_to_body(params: dict, form_data: dict, user) -> dict:
system = params.get("system", None)
if not system:
return form_data
if user:
template_params = {
"user_name": user.name,
"user_location": user.info.get("location") if user.info else None,
}
else:
template_params = {}
system = prompt_template(system, **template_params)
form_data["messages"] = add_or_update_system_message(
system, form_data.get("messages", [])
)
return form_data
async def generate_function_chat_completion(form_data, user):
model_id = form_data.get("model")
model_info = Models.get_model_by_id(model_id)
metadata = None
if "metadata" in form_data:
metadata = form_data["metadata"]
del form_data["metadata"]
metadata = form_data.pop("metadata", None)
__event_emitter__ = None
__event_call__ = None
__task__ = None
if metadata:
if (
metadata.get("session_id")
and metadata.get("chat_id")
and metadata.get("message_id")
):
__event_emitter__ = await get_event_emitter(metadata)
__event_call__ = await get_event_call(metadata)
if metadata.get("task"):
__task__ = metadata.get("task")
if all(k in metadata for k in ("session_id", "chat_id", "message_id")):
__event_emitter__ = get_event_emitter(metadata)
__event_call__ = get_event_call(metadata)
__task__ = metadata.get("task", None)
if model_info:
if model_info.base_model_id:
form_data["model"] = model_info.base_model_id
model_info.params = model_info.params.model_dump()
params = model_info.params.model_dump()
form_data = apply_model_params_to_body(params, form_data)
form_data = apply_model_system_prompt_to_body(params, form_data, user)
if model_info.params:
if model_info.params.get("temperature", None) is not None:
form_data["temperature"] = float(model_info.params.get("temperature"))
pipe_id = get_pipe_id(form_data)
function_module = get_function_module(pipe_id)
if model_info.params.get("top_p", None):
form_data["top_p"] = int(model_info.params.get("top_p", None))
pipe = function_module.pipe
params = get_function_params(
function_module,
form_data,
user,
{
"__event_emitter__": __event_emitter__,
"__event_call__": __event_call__,
"__task__": __task__,
},
)
if model_info.params.get("max_tokens", None):
form_data["max_tokens"] = int(model_info.params.get("max_tokens", None))
if model_info.params.get("frequency_penalty", None):
form_data["frequency_penalty"] = int(
model_info.params.get("frequency_penalty", None)
)
if model_info.params.get("seed", None):
form_data["seed"] = model_info.params.get("seed", None)
if model_info.params.get("stop", None):
form_data["stop"] = (
[
bytes(stop, "utf-8").decode("unicode_escape")
for stop in model_info.params["stop"]
]
if model_info.params.get("stop", None)
else None
)
system = model_info.params.get("system", None)
if system:
system = prompt_template(
system,
**(
{
"user_name": user.name,
"user_location": (
user.info.get("location") if user.info else None
),
}
if user
else {}
),
)
# Check if the payload already has a system message
# If not, add a system message to the payload
if form_data.get("messages"):
for message in form_data["messages"]:
if message.get("role") == "system":
message["content"] = system + message["content"]
break
else:
form_data["messages"].insert(
0,
{
"role": "system",
"content": system,
},
)
else:
pass
async def job():
pipe_id = form_data["model"]
if "." in pipe_id:
pipe_id, sub_pipe_id = pipe_id.split(".", 1)
print(pipe_id)
# Check if function is already loaded
if pipe_id not in app.state.FUNCTIONS:
function_module, function_type, frontmatter = load_function_module_by_id(
pipe_id
)
app.state.FUNCTIONS[pipe_id] = function_module
else:
function_module = app.state.FUNCTIONS[pipe_id]
if hasattr(function_module, "valves") and hasattr(function_module, "Valves"):
valves = Functions.get_function_valves_by_id(pipe_id)
function_module.valves = function_module.Valves(
**(valves if valves else {})
)
pipe = function_module.pipe
# Get the signature of the function
sig = inspect.signature(pipe)
params = {"body": form_data}
if "__user__" in sig.parameters:
__user__ = {
"id": user.id,
"email": user.email,
"name": user.name,
"role": user.role,
}
if form_data["stream"]:
async def stream_content():
try:
if hasattr(function_module, "UserValves"):
__user__["valves"] = function_module.UserValves(
**Functions.get_user_valves_by_id_and_user_id(pipe_id, user.id)
)
except Exception as e:
print(e)
res = await execute_pipe(pipe, params)
params = {**params, "__user__": __user__}
if "__event_emitter__" in sig.parameters:
params = {**params, "__event_emitter__": __event_emitter__}
if "__event_call__" in sig.parameters:
params = {**params, "__event_call__": __event_call__}
if "__task__" in sig.parameters:
params = {**params, "__task__": __task__}
if form_data["stream"]:
async def stream_content():
try:
if inspect.iscoroutinefunction(pipe):
res = await pipe(**params)
else:
res = pipe(**params)
# Directly return if the response is a StreamingResponse
if isinstance(res, StreamingResponse):
async for data in res.body_iterator:
yield data
return
if isinstance(res, dict):
yield f"data: {json.dumps(res)}\n\n"
return
except Exception as e:
print(f"Error: {e}")
yield f"data: {json.dumps({'error': {'detail':str(e)}})}\n\n"
# Directly return if the response is a StreamingResponse
if isinstance(res, StreamingResponse):
async for data in res.body_iterator:
yield data
return
if isinstance(res, dict):
yield f"data: {json.dumps(res)}\n\n"
return
if isinstance(res, str):
message = stream_message_template(form_data["model"], res)
yield f"data: {json.dumps(message)}\n\n"
if isinstance(res, Iterator):
for line in res:
if isinstance(line, BaseModel):
line = line.model_dump_json()
line = f"data: {line}"
if isinstance(line, dict):
line = f"data: {json.dumps(line)}"
try:
line = line.decode("utf-8")
except:
pass
if line.startswith("data:"):
yield f"{line}\n\n"
else:
line = stream_message_template(form_data["model"], line)
yield f"data: {json.dumps(line)}\n\n"
if isinstance(res, str) or isinstance(res, Generator):
finish_message = {
"id": f"{form_data['model']}-{str(uuid.uuid4())}",
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": form_data["model"],
"choices": [
{
"index": 0,
"delta": {},
"logprobs": None,
"finish_reason": "stop",
}
],
}
yield f"data: {json.dumps(finish_message)}\n\n"
yield f"data: [DONE]"
if isinstance(res, AsyncGenerator):
async for line in res:
if isinstance(line, BaseModel):
line = line.model_dump_json()
line = f"data: {line}"
if isinstance(line, dict):
line = f"data: {json.dumps(line)}"
try:
line = line.decode("utf-8")
except:
pass
if line.startswith("data:"):
yield f"{line}\n\n"
else:
line = stream_message_template(form_data["model"], line)
yield f"data: {json.dumps(line)}\n\n"
return StreamingResponse(stream_content(), media_type="text/event-stream")
else:
try:
if inspect.iscoroutinefunction(pipe):
res = await pipe(**params)
else:
res = pipe(**params)
if isinstance(res, StreamingResponse):
return res
except Exception as e:
print(f"Error: {e}")
return {"error": {"detail": str(e)}}
yield f"data: {json.dumps({'error': {'detail':str(e)}})}\n\n"
return
if isinstance(res, dict):
return res
elif isinstance(res, BaseModel):
return res.model_dump()
else:
message = ""
if isinstance(res, str):
message = res
elif isinstance(res, Generator):
for stream in res:
message = f"{message}{stream}"
elif isinstance(res, AsyncGenerator):
async for stream in res:
message = f"{message}{stream}"
if isinstance(res, str):
message = openai_chat_chunk_message_template(form_data["model"], res)
yield f"data: {json.dumps(message)}\n\n"
return {
"id": f"{form_data['model']}-{str(uuid.uuid4())}",
"object": "chat.completion",
"created": int(time.time()),
"model": form_data["model"],
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": message,
},
"logprobs": None,
"finish_reason": "stop",
}
],
}
if isinstance(res, Iterator):
for line in res:
yield process_line(form_data, line)
return await job()
if isinstance(res, AsyncGenerator):
async for line in res:
yield process_line(form_data, line)
if isinstance(res, str) or isinstance(res, Generator):
finish_message = openai_chat_chunk_message_template(
form_data["model"], ""
)
finish_message["choices"][0]["finish_reason"] = "stop"
yield f"data: {json.dumps(finish_message)}\n\n"
yield "data: [DONE]"
return StreamingResponse(stream_content(), media_type="text/event-stream")
else:
try:
res = await execute_pipe(pipe, params)
except Exception as e:
print(f"Error: {e}")
return {"error": {"detail": str(e)}}
if isinstance(res, StreamingResponse) or isinstance(res, dict):
return res
if isinstance(res, BaseModel):
return res.model_dump()
message = await get_message_content(res)
return openai_chat_completion_message_template(form_data["model"], message)

View File

@@ -1,13 +1,11 @@
import json
import logging
from typing import Optional
from typing import Optional, List
from pydantic import BaseModel, ConfigDict
from sqlalchemy import String, Column, BigInteger, Text
from sqlalchemy import Column, BigInteger, Text
from apps.webui.internal.db import Base, JSONField, get_db
from typing import List, Union, Optional
from config import SRC_LOG_LEVELS
import time
@@ -113,7 +111,6 @@ class ModelForm(BaseModel):
class ModelsTable:
def insert_new_model(
self, form_data: ModelForm, user_id: str
) -> Optional[ModelModel]:
@@ -126,9 +123,7 @@ class ModelsTable:
}
)
try:
with get_db() as db:
result = Model(**model.model_dump())
db.add(result)
db.commit()
@@ -144,13 +139,11 @@ class ModelsTable:
def get_all_models(self) -> List[ModelModel]:
with get_db() as db:
return [ModelModel.model_validate(model) for model in db.query(Model).all()]
def get_model_by_id(self, id: str) -> Optional[ModelModel]:
try:
with get_db() as db:
model = db.get(Model, id)
return ModelModel.model_validate(model)
except:
@@ -178,7 +171,6 @@ class ModelsTable:
def delete_model_by_id(self, id: str) -> bool:
try:
with get_db() as db:
db.query(Model).filter_by(id=id).delete()
db.commit()

View File

@@ -1,3 +1,6 @@
from pathlib import Path
import site
from fastapi import APIRouter, UploadFile, File, Response
from fastapi import Depends, HTTPException, status
from starlette.responses import StreamingResponse, FileResponse
@@ -64,8 +67,18 @@ async def download_chat_as_pdf(
pdf = FPDF()
pdf.add_page()
STATIC_DIR = "./static"
FONTS_DIR = f"{STATIC_DIR}/fonts"
# When running in docker, workdir is /app/backend, so fonts is in /app/backend/static/fonts
FONTS_DIR = Path("./static/fonts")
# Non Docker Installation
# When running using `pip install` the static directory is in the site packages.
if not FONTS_DIR.exists():
FONTS_DIR = Path(site.getsitepackages()[0]) / "static/fonts"
# When running using `pip install -e .` the static directory is in the site packages.
# This path only works if `open-webui serve` is run from the root of this project.
if not FONTS_DIR.exists():
FONTS_DIR = Path("./backend/static/fonts")
pdf.add_font("NotoSans", "", f"{FONTS_DIR}/NotoSans-Regular.ttf")
pdf.add_font("NotoSans", "b", f"{FONTS_DIR}/NotoSans-Bold.ttf")

View File

@@ -349,6 +349,12 @@ GOOGLE_OAUTH_SCOPE = PersistentConfig(
os.environ.get("GOOGLE_OAUTH_SCOPE", "openid email profile"),
)
GOOGLE_REDIRECT_URI = PersistentConfig(
"GOOGLE_REDIRECT_URI",
"oauth.google.redirect_uri",
os.environ.get("GOOGLE_REDIRECT_URI", ""),
)
MICROSOFT_CLIENT_ID = PersistentConfig(
"MICROSOFT_CLIENT_ID",
"oauth.microsoft.client_id",
@@ -373,6 +379,12 @@ MICROSOFT_OAUTH_SCOPE = PersistentConfig(
os.environ.get("MICROSOFT_OAUTH_SCOPE", "openid email profile"),
)
MICROSOFT_REDIRECT_URI = PersistentConfig(
"MICROSOFT_REDIRECT_URI",
"oauth.microsoft.redirect_uri",
os.environ.get("MICROSOFT_REDIRECT_URI", ""),
)
OAUTH_CLIENT_ID = PersistentConfig(
"OAUTH_CLIENT_ID",
"oauth.oidc.client_id",
@@ -391,6 +403,12 @@ OPENID_PROVIDER_URL = PersistentConfig(
os.environ.get("OPENID_PROVIDER_URL", ""),
)
OPENID_REDIRECT_URI = PersistentConfig(
"OPENID_REDIRECT_URI",
"oauth.oidc.redirect_uri",
os.environ.get("OPENID_REDIRECT_URI", ""),
)
OAUTH_SCOPES = PersistentConfig(
"OAUTH_SCOPES",
"oauth.oidc.scopes",
@@ -424,6 +442,7 @@ def load_oauth_providers():
"client_secret": GOOGLE_CLIENT_SECRET.value,
"server_metadata_url": "https://accounts.google.com/.well-known/openid-configuration",
"scope": GOOGLE_OAUTH_SCOPE.value,
"redirect_uri": GOOGLE_REDIRECT_URI.value,
}
if (
@@ -436,6 +455,7 @@ def load_oauth_providers():
"client_secret": MICROSOFT_CLIENT_SECRET.value,
"server_metadata_url": f"https://login.microsoftonline.com/{MICROSOFT_CLIENT_TENANT_ID.value}/v2.0/.well-known/openid-configuration",
"scope": MICROSOFT_OAUTH_SCOPE.value,
"redirect_uri": MICROSOFT_REDIRECT_URI.value,
}
if (
@@ -449,6 +469,7 @@ def load_oauth_providers():
"server_metadata_url": OPENID_PROVIDER_URL.value,
"scope": OAUTH_SCOPES.value,
"name": OAUTH_PROVIDER_NAME.value,
"redirect_uri": OPENID_REDIRECT_URI.value,
}
@@ -1281,6 +1302,24 @@ COMFYUI_SD3 = PersistentConfig(
os.environ.get("COMFYUI_SD3", "").lower() == "true",
)
COMFYUI_FLUX = PersistentConfig(
"COMFYUI_FLUX",
"image_generation.comfyui.flux",
os.environ.get("COMFYUI_FLUX", "").lower() == "true",
)
COMFYUI_FLUX_WEIGHT_DTYPE = PersistentConfig(
"COMFYUI_FLUX_WEIGHT_DTYPE",
"image_generation.comfyui.flux_weight_dtype",
os.getenv("COMFYUI_FLUX_WEIGHT_DTYPE", ""),
)
COMFYUI_FLUX_FP8_CLIP = PersistentConfig(
"COMFYUI_FLUX_FP8_CLIP",
"image_generation.comfyui.flux_fp8_clip",
os.getenv("COMFYUI_FLUX_FP8_CLIP", ""),
)
IMAGES_OPENAI_API_BASE_URL = PersistentConfig(
"IMAGES_OPENAI_API_BASE_URL",
"image_generation.openai.api_base_url",

View File

@@ -13,8 +13,6 @@ import aiohttp
import requests
import mimetypes
import shutil
import os
import uuid
import inspect
from fastapi import FastAPI, Request, Depends, status, UploadFile, File, Form
@@ -29,7 +27,7 @@ from starlette.middleware.sessions import SessionMiddleware
from starlette.responses import StreamingResponse, Response, RedirectResponse
from apps.socket.main import sio, app as socket_app, get_event_emitter, get_event_call
from apps.socket.main import app as socket_app, get_event_emitter, get_event_call
from apps.ollama.main import (
app as ollama_app,
get_all_models as get_ollama_models,
@@ -619,32 +617,15 @@ class ChatCompletionMiddleware(BaseHTTPMiddleware):
content={"detail": str(e)},
)
# Extract valves from the request body
valves = None
if "valves" in body:
valves = body["valves"]
del body["valves"]
metadata = {
"chat_id": body.pop("chat_id", None),
"message_id": body.pop("id", None),
"session_id": body.pop("session_id", None),
"valves": body.pop("valves", None),
}
# Extract session_id, chat_id and message_id from the request body
session_id = None
if "session_id" in body:
session_id = body["session_id"]
del body["session_id"]
chat_id = None
if "chat_id" in body:
chat_id = body["chat_id"]
del body["chat_id"]
message_id = None
if "id" in body:
message_id = body["id"]
del body["id"]
__event_emitter__ = await get_event_emitter(
{"chat_id": chat_id, "message_id": message_id, "session_id": session_id}
)
__event_call__ = await get_event_call(
{"chat_id": chat_id, "message_id": message_id, "session_id": session_id}
)
__event_emitter__ = get_event_emitter(metadata)
__event_call__ = get_event_call(metadata)
# Initialize data_items to store additional data to be sent to the client
data_items = []
@@ -709,13 +690,7 @@ class ChatCompletionMiddleware(BaseHTTPMiddleware):
if len(citations) > 0:
data_items.append({"citations": citations})
body["metadata"] = {
"session_id": session_id,
"chat_id": chat_id,
"message_id": message_id,
"valves": valves,
}
body["metadata"] = metadata
modified_body_bytes = json.dumps(body).encode("utf-8")
# Replace the request body with the modified one
request._body = modified_body_bytes
@@ -1191,13 +1166,13 @@ async def chat_completed(form_data: dict, user=Depends(get_verified_user)):
status_code=r.status_code,
content=res,
)
except:
except Exception:
pass
else:
pass
__event_emitter__ = await get_event_emitter(
__event_emitter__ = get_event_emitter(
{
"chat_id": data["chat_id"],
"message_id": data["id"],
@@ -1205,7 +1180,7 @@ async def chat_completed(form_data: dict, user=Depends(get_verified_user)):
}
)
__event_call__ = await get_event_call(
__event_call__ = get_event_call(
{
"chat_id": data["chat_id"],
"message_id": data["id"],
@@ -1310,9 +1285,7 @@ async def chat_completed(form_data: dict, user=Depends(get_verified_user)):
@app.post("/api/chat/actions/{action_id}")
async def chat_completed(
action_id: str, form_data: dict, user=Depends(get_verified_user)
):
async def chat_action(action_id: str, form_data: dict, user=Depends(get_verified_user)):
if "." in action_id:
action_id, sub_action_id = action_id.split(".")
else:
@@ -1334,14 +1307,14 @@ async def chat_completed(
)
model = app.state.MODELS[model_id]
__event_emitter__ = await get_event_emitter(
__event_emitter__ = get_event_emitter(
{
"chat_id": data["chat_id"],
"message_id": data["id"],
"session_id": data["session_id"],
}
)
__event_call__ = await get_event_call(
__event_call__ = get_event_call(
{
"chat_id": data["chat_id"],
"message_id": data["id"],
@@ -1770,7 +1743,6 @@ class AddPipelineForm(BaseModel):
@app.post("/api/pipelines/add")
async def add_pipeline(form_data: AddPipelineForm, user=Depends(get_admin_user)):
r = None
try:
urlIdx = form_data.urlIdx
@@ -1813,7 +1785,6 @@ class DeletePipelineForm(BaseModel):
@app.delete("/api/pipelines/delete")
async def delete_pipeline(form_data: DeletePipelineForm, user=Depends(get_admin_user)):
r = None
try:
urlIdx = form_data.urlIdx
@@ -1891,7 +1862,6 @@ async def get_pipeline_valves(
models = await get_all_models()
r = None
try:
url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
@@ -2143,6 +2113,7 @@ for provider_name, provider_config in OAUTH_PROVIDERS.items():
client_kwargs={
"scope": provider_config["scope"],
},
redirect_uri=provider_config["redirect_uri"],
)
# SessionMiddleware is used by authlib for oauth
@@ -2160,7 +2131,10 @@ if len(OAUTH_PROVIDERS) > 0:
async def oauth_login(provider: str, request: Request):
if provider not in OAUTH_PROVIDERS:
raise HTTPException(404)
redirect_uri = request.url_for("oauth_callback", provider=provider)
# If the provider has a custom redirect URL, use that, otherwise automatically generate one
redirect_uri = OAUTH_PROVIDERS[provider].get("redirect_uri") or request.url_for(
"oauth_callback", provider=provider
)
return await oauth.create_client(provider).authorize_redirect(request, redirect_uri)

View File

@@ -12,6 +12,7 @@ passlib[bcrypt]==1.7.4
requests==2.32.3
aiohttp==3.9.5
sqlalchemy==2.0.31
alembic==1.13.2
peewee==3.17.6
@@ -19,7 +20,7 @@ peewee-migrate==1.12.2
psycopg2-binary==2.9.9
PyMySQL==1.1.1
bcrypt==4.1.3
SQLAlchemy
pymongo
redis
boto3==1.34.110

View File

@@ -1,6 +1,5 @@
from pathlib import Path
import hashlib
import json
import re
from datetime import timedelta
from typing import Optional, List, Tuple
@@ -8,37 +7,39 @@ import uuid
import time
def get_last_user_message_item(messages: List[dict]) -> str:
def get_last_user_message_item(messages: List[dict]) -> Optional[dict]:
for message in reversed(messages):
if message["role"] == "user":
return message
return None
def get_last_user_message(messages: List[dict]) -> str:
message = get_last_user_message_item(messages)
if message is not None:
if isinstance(message["content"], list):
for item in message["content"]:
if item["type"] == "text":
return item["text"]
def get_content_from_message(message: dict) -> Optional[str]:
if isinstance(message["content"], list):
for item in message["content"]:
if item["type"] == "text":
return item["text"]
else:
return message["content"]
return None
def get_last_assistant_message(messages: List[dict]) -> str:
def get_last_user_message(messages: List[dict]) -> Optional[str]:
message = get_last_user_message_item(messages)
if message is None:
return None
return get_content_from_message(message)
def get_last_assistant_message(messages: List[dict]) -> Optional[str]:
for message in reversed(messages):
if message["role"] == "assistant":
if isinstance(message["content"], list):
for item in message["content"]:
if item["type"] == "text":
return item["text"]
return message["content"]
return get_content_from_message(message)
return None
def get_system_message(messages: List[dict]) -> dict:
def get_system_message(messages: List[dict]) -> Optional[dict]:
for message in messages:
if message["role"] == "system":
return message
@@ -49,7 +50,7 @@ def remove_system_message(messages: List[dict]) -> List[dict]:
return [message for message in messages if message["role"] != "system"]
def pop_system_message(messages: List[dict]) -> Tuple[dict, List[dict]]:
def pop_system_message(messages: List[dict]) -> Tuple[Optional[dict], List[dict]]:
return get_system_message(messages), remove_system_message(messages)
@@ -87,23 +88,29 @@ def add_or_update_system_message(content: str, messages: List[dict]):
return messages
def stream_message_template(model: str, message: str):
def openai_chat_message_template(model: str):
return {
"id": f"{model}-{str(uuid.uuid4())}",
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": model,
"choices": [
{
"index": 0,
"delta": {"content": message},
"logprobs": None,
"finish_reason": None,
}
],
"choices": [{"index": 0, "logprobs": None, "finish_reason": None}],
}
def openai_chat_chunk_message_template(model: str, message: str):
template = openai_chat_message_template(model)
template["object"] = "chat.completion.chunk"
template["choices"][0]["delta"] = {"content": message}
return template
def openai_chat_completion_message_template(model: str, message: str):
template = openai_chat_message_template(model)
template["object"] = "chat.completion"
template["choices"][0]["message"] = {"content": message, "role": "assistant"}
template["choices"][0]["finish_reason"] = "stop"
def get_gravatar_url(email):
# Trim leading and trailing whitespace from
# an email address and force all characters
@@ -174,7 +181,7 @@ def extract_folders_after_data_docs(path):
tags = []
folders = parts[index_docs:-1]
for idx, part in enumerate(folders):
for idx, _ in enumerate(folders):
tags.append("/".join(folders[: idx + 1]))
return tags
@@ -270,11 +277,11 @@ def parse_ollama_modelfile(model_text):
value = param_match.group(1)
try:
if param_type == int:
if param_type is int:
value = int(value)
elif param_type == float:
elif param_type is float:
value = float(value)
elif param_type == bool:
elif param_type is bool:
value = value.lower() == "true"
except Exception as e:
print(e)