2024-06-21 13:35:54 +00:00
import base64
2024-05-27 17:07:38 +00:00
import uuid
2024-05-09 04:00:03 +00:00
from contextlib import asynccontextmanager
2024-05-27 17:07:38 +00:00
from authlib . integrations . starlette_client import OAuth
from authlib . oidc . core import UserInfo
2024-02-23 08:30:26 +00:00
import json
2024-01-07 10:48:21 +00:00
import time
2024-02-24 08:21:53 +00:00
import os
import sys
2024-03-20 23:11:36 +00:00
import logging
2024-04-10 06:03:05 +00:00
import aiohttp
2024-02-25 19:26:58 +00:00
import requests
2024-05-21 22:04:00 +00:00
import mimetypes
2024-06-05 20:57:48 +00:00
import shutil
2024-06-11 17:19:59 +00:00
import inspect
2024-02-23 08:30:26 +00:00
2024-06-05 20:57:48 +00:00
from fastapi import FastAPI , Request , Depends , status , UploadFile , File , Form
2023-11-15 00:28:51 +00:00
from fastapi . staticfiles import StaticFiles
2024-05-28 16:50:17 +00:00
from fastapi . responses import JSONResponse
2023-11-15 00:28:51 +00:00
from fastapi import HTTPException
from fastapi . middleware . cors import CORSMiddleware
2024-06-18 13:03:31 +00:00
from sqlalchemy import text
2023-11-19 00:47:12 +00:00
from starlette . exceptions import HTTPException as StarletteHTTPException
2024-03-09 06:34:47 +00:00
from starlette . middleware . base import BaseHTTPMiddleware
2024-05-27 17:07:38 +00:00
from starlette . middleware . sessions import SessionMiddleware
from starlette . responses import StreamingResponse , Response , RedirectResponse
2024-01-07 06:07:20 +00:00
2024-06-04 06:39:52 +00:00
2024-07-31 12:35:02 +00:00
from apps . socket . main import app as socket_app , get_event_emitter , get_event_call
2024-06-09 20:17:44 +00:00
from apps . ollama . main import (
app as ollama_app ,
get_all_models as get_ollama_models ,
generate_openai_chat_completion as generate_ollama_chat_completion ,
)
from apps . openai . main import (
app as openai_app ,
get_all_models as get_openai_models ,
generate_chat_completion as generate_openai_chat_completion ,
)
2024-04-10 06:03:05 +00:00
2024-02-11 08:17:50 +00:00
from apps . audio . main import app as audio_app
2024-02-22 02:12:01 +00:00
from apps . images . main import app as images_app
from apps . rag . main import app as rag_app
2024-06-24 18:17:18 +00:00
from apps . webui . main import (
app as webui_app ,
get_pipe_models ,
generate_function_chat_completion ,
)
2024-07-09 10:51:43 +00:00
from apps . webui . internal . db import Session
2024-01-07 06:07:20 +00:00
2024-06-05 20:57:48 +00:00
2024-03-10 05:19:20 +00:00
from pydantic import BaseModel
2024-08-14 12:46:31 +00:00
from typing import Optional
2024-02-24 06:44:56 +00:00
2024-05-27 17:07:38 +00:00
from apps . webui . models . auths import Auths
2024-07-09 10:51:43 +00:00
from apps . webui . models . models import Models
2024-06-11 06:40:27 +00:00
from apps . webui . models . tools import Tools
2024-06-20 09:30:00 +00:00
from apps . webui . models . functions import Functions
2024-05-27 17:07:38 +00:00
from apps . webui . models . users import Users
2024-06-20 09:30:00 +00:00
from apps . webui . utils import load_toolkit_module_by_id , load_function_module_by_id
2024-06-11 06:40:27 +00:00
2024-05-28 02:16:07 +00:00
from utils . utils import (
get_admin_user ,
get_verified_user ,
get_current_user ,
get_http_authorization_cred ,
2024-05-27 17:07:38 +00:00
get_password_hash ,
create_token ,
2024-05-28 02:16:07 +00:00
)
2024-06-11 06:40:27 +00:00
from utils . task import (
title_generation_template ,
search_query_generation_template ,
2024-08-11 13:56:16 +00:00
tool_calling_generation_template ,
2024-06-11 06:40:27 +00:00
)
2024-06-20 11:38:59 +00:00
from utils . misc import (
get_last_user_message ,
add_or_update_system_message ,
2024-07-26 11:22:13 +00:00
prepend_to_first_user_message_content ,
2024-06-24 10:46:48 +00:00
parse_duration ,
2024-06-20 11:38:59 +00:00
)
2024-06-09 21:25:31 +00:00
2024-06-11 08:10:24 +00:00
from apps . rag . utils import get_rag_context , rag_template
2024-03-09 06:34:47 +00:00
2024-03-10 05:47:01 +00:00
from config import (
WEBUI_NAME ,
2024-05-07 00:29:16 +00:00
WEBUI_URL ,
2024-05-08 15:40:18 +00:00
WEBUI_AUTH ,
2024-03-10 05:47:01 +00:00
ENV ,
VERSION ,
CHANGELOG ,
FRONTEND_BUILD_DIR ,
2024-04-09 10:32:28 +00:00
CACHE_DIR ,
STATIC_DIR ,
2024-06-30 21:48:05 +00:00
DEFAULT_LOCALE ,
2024-05-24 08:40:48 +00:00
ENABLE_OPENAI_API ,
ENABLE_OLLAMA_API ,
2024-04-26 21:17:18 +00:00
ENABLE_MODEL_FILTER ,
2024-03-10 05:47:01 +00:00
MODEL_FILTER_LIST ,
2024-03-20 23:11:36 +00:00
GLOBAL_LOG_LEVEL ,
SRC_LOG_LEVELS ,
2024-03-21 01:35:02 +00:00
WEBHOOK_URL ,
2024-04-22 18:55:46 +00:00
ENABLE_ADMIN_EXPORT ,
2024-05-26 07:49:30 +00:00
WEBUI_BUILD_HASH ,
2024-06-09 21:53:10 +00:00
TASK_MODEL ,
TASK_MODEL_EXTERNAL ,
2024-06-09 21:25:31 +00:00
TITLE_GENERATION_PROMPT_TEMPLATE ,
2024-06-09 21:53:10 +00:00
SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE ,
2024-06-09 22:19:36 +00:00
SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD ,
2024-06-11 06:40:27 +00:00
TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE ,
2024-06-24 02:28:33 +00:00
SAFE_MODE ,
2024-05-26 07:37:09 +00:00
OAUTH_PROVIDERS ,
2024-05-27 17:07:38 +00:00
ENABLE_OAUTH_SIGNUP ,
OAUTH_MERGE_ACCOUNTS_BY_EMAIL ,
WEBUI_SECRET_KEY ,
2024-06-05 18:21:42 +00:00
WEBUI_SESSION_COOKIE_SAME_SITE ,
2024-06-07 08:13:42 +00:00
WEBUI_SESSION_COOKIE_SECURE ,
2024-08-04 14:16:14 +00:00
ENABLE_ADMIN_CHAT_ACCESS ,
2024-06-09 21:25:31 +00:00
AppConfig ,
2024-03-10 05:47:01 +00:00
)
2024-07-09 10:51:43 +00:00
2024-07-03 22:46:56 +00:00
from constants import ERROR_MESSAGES , WEBHOOK_MESSAGES , TASKS
2024-05-27 17:07:38 +00:00
from utils . webhook import post_webhook
2024-02-25 19:26:58 +00:00
2024-06-24 02:28:33 +00:00
if SAFE_MODE :
print ( " SAFE MODE ENABLED " )
Functions . deactivate_all_functions ( )
2024-03-20 23:11:36 +00:00
logging . basicConfig ( stream = sys . stdout , level = GLOBAL_LOG_LEVEL )
log = logging . getLogger ( __name__ )
log . setLevel ( SRC_LOG_LEVELS [ " MAIN " ] )
2023-11-15 00:28:51 +00:00
2024-03-28 09:45:56 +00:00
2023-11-15 00:28:51 +00:00
class SPAStaticFiles ( StaticFiles ) :
async def get_response ( self , path : str , scope ) :
try :
return await super ( ) . get_response ( path , scope )
except ( HTTPException , StarletteHTTPException ) as ex :
if ex . status_code == 404 :
return await super ( ) . get_response ( " index.html " , scope )
else :
raise ex
2024-04-02 10:03:55 +00:00
print (
2024-05-03 21:23:38 +00:00
rf """
2024-04-02 10:03:55 +00:00
___ __ __ _ _ _ ___
/ _ \ _ __ ___ _ __ \ \ / / __ | | __ | | | | _ _ |
| | | | ' _ \ / _ \ ' _ \ \ \ / \ / / _ \ ' _ \ | | | || |
| | _ | | | _ ) | __ / | | | \ V V / __ / | _ ) | | _ | | | |
\___ / | . __ / \___ | _ | | _ | \_ / \_ / \___ | _ . __ / \___ / | ___ |
| _ |
2024-05-22 19:22:38 +00:00
v { VERSION } - building the best open - source AI user interface .
2024-05-26 07:49:30 +00:00
{ f " Commit: { WEBUI_BUILD_HASH } " if WEBUI_BUILD_HASH != " dev-build " else " " }
2024-04-02 10:03:55 +00:00
https : / / github . com / open - webui / open - webui
"""
)
2023-11-15 00:28:51 +00:00
2024-06-18 13:03:31 +00:00
def run_migrations ( ) :
2024-07-08 19:55:27 +00:00
try :
from alembic . config import Config
from alembic import command
alembic_cfg = Config ( " alembic.ini " )
command . upgrade ( alembic_cfg , " head " )
except Exception as e :
print ( f " Error: { e } " )
2024-06-18 13:03:31 +00:00
2024-05-09 04:00:03 +00:00
@asynccontextmanager
async def lifespan ( app : FastAPI ) :
2024-06-18 13:03:31 +00:00
run_migrations ( )
2024-05-09 04:00:03 +00:00
yield
app = FastAPI (
docs_url = " /docs " if ENV == " dev " else None , redoc_url = None , lifespan = lifespan
)
2023-11-15 00:28:51 +00:00
2024-05-10 07:03:24 +00:00
app . state . config = AppConfig ( )
2024-05-24 08:40:48 +00:00
app . state . config . ENABLE_OPENAI_API = ENABLE_OPENAI_API
app . state . config . ENABLE_OLLAMA_API = ENABLE_OLLAMA_API
2024-05-10 07:03:24 +00:00
app . state . config . ENABLE_MODEL_FILTER = ENABLE_MODEL_FILTER
app . state . config . MODEL_FILTER_LIST = MODEL_FILTER_LIST
2024-03-10 05:19:20 +00:00
2024-05-10 07:03:24 +00:00
app . state . config . WEBHOOK_URL = WEBHOOK_URL
2024-06-09 21:53:10 +00:00
app . state . config . TASK_MODEL = TASK_MODEL
app . state . config . TASK_MODEL_EXTERNAL = TASK_MODEL_EXTERNAL
2024-06-09 21:25:31 +00:00
app . state . config . TITLE_GENERATION_PROMPT_TEMPLATE = TITLE_GENERATION_PROMPT_TEMPLATE
2024-06-09 21:53:10 +00:00
app . state . config . SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE = (
SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE
)
2024-06-09 22:19:36 +00:00
app . state . config . SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD = (
SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD
)
2024-06-11 06:40:27 +00:00
app . state . config . TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE = (
TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE
)
2024-05-25 01:26:36 +00:00
app . state . MODELS = { }
2023-11-15 00:28:51 +00:00
origins = [ " * " ]
2024-05-17 03:49:28 +00:00
2024-06-20 08:51:39 +00:00
##################################
#
# ChatCompletion Middleware
#
##################################
2024-07-02 02:37:54 +00:00
async def get_body_and_model_and_user ( request ) :
# Read the original request body
body = await request . body ( )
body_str = body . decode ( " utf-8 " )
body = json . loads ( body_str ) if body_str else { }
model_id = body [ " model " ]
if model_id not in app . state . MODELS :
2024-07-09 14:57:24 +00:00
raise Exception ( " Model not found " )
2024-07-02 02:37:54 +00:00
model = app . state . MODELS [ model_id ]
user = get_current_user (
request ,
get_http_authorization_cred ( request . headers . get ( " Authorization " ) ) ,
)
return body , model , user
def get_task_model_id ( default_model_id ) :
# Set the task model
task_model_id = default_model_id
# Check if the user has a custom task model and use that model
if app . state . MODELS [ task_model_id ] [ " owned_by " ] == " ollama " :
if (
app . state . config . TASK_MODEL
and app . state . config . TASK_MODEL in app . state . MODELS
) :
task_model_id = app . state . config . TASK_MODEL
else :
if (
app . state . config . TASK_MODEL_EXTERNAL
and app . state . config . TASK_MODEL_EXTERNAL in app . state . MODELS
) :
task_model_id = app . state . config . TASK_MODEL_EXTERNAL
return task_model_id
def get_filter_function_ids ( model ) :
def get_priority ( function_id ) :
function = Functions . get_function_by_id ( function_id )
if function is not None and hasattr ( function , " valves " ) :
2024-08-11 07:31:40 +00:00
# TODO: Fix FunctionModel
2024-07-02 02:37:54 +00:00
return ( function . valves if function . valves else { } ) . get ( " priority " , 0 )
return 0
filter_ids = [ function . id for function in Functions . get_global_filter_functions ( ) ]
if " info " in model and " meta " in model [ " info " ] :
filter_ids . extend ( model [ " info " ] [ " meta " ] . get ( " filterIds " , [ ] ) )
filter_ids = list ( set ( filter_ids ) )
enabled_filter_ids = [
function . id
for function in Functions . get_functions_by_type ( " filter " , active_only = True )
]
filter_ids = [
filter_id for filter_id in filter_ids if filter_id in enabled_filter_ids
]
filter_ids . sort ( key = get_priority )
return filter_ids
2024-08-10 12:40:04 +00:00
async def get_content_from_response ( response ) - > Optional [ str ] :
content = None
if hasattr ( response , " body_iterator " ) :
async for chunk in response . body_iterator :
data = json . loads ( chunk . decode ( " utf-8 " ) )
content = data [ " choices " ] [ 0 ] [ " message " ] [ " content " ]
# Cleanup any remaining background tasks if necessary
if response . background is not None :
await response . background ( )
else :
content = response [ " choices " ] [ 0 ] [ " message " ] [ " content " ]
return content
2024-08-10 13:25:20 +00:00
async def call_tool_from_completion (
result : dict , extra_params : dict , toolkit_module
) - > Optional [ str ] :
if " name " not in result :
return None
tool = getattr ( toolkit_module , result [ " name " ] )
try :
# Get the signature of the function
sig = inspect . signature ( tool )
params = result [ " parameters " ]
for key , value in extra_params . items ( ) :
if key in sig . parameters :
params [ key ] = value
if inspect . iscoroutinefunction ( tool ) :
return await tool ( * * params )
else :
return tool ( * * params )
except Exception as e :
print ( f " Error: { e } " )
return None
2024-08-11 13:56:16 +00:00
def get_tool_calling_payload ( messages , task_model_id , content ) :
2024-06-11 18:15:43 +00:00
user_message = get_last_user_message ( messages )
2024-08-11 08:05:22 +00:00
history = " \n " . join (
f " { message [ ' role ' ] . upper ( ) } : \" \" \" { message [ ' content ' ] } \" \" \" "
for message in messages [ : : - 1 ] [ : 4 ]
2024-06-11 18:15:43 +00:00
)
2024-08-11 08:05:22 +00:00
prompt = f " History: \n { history } \n Query: { user_message } "
2024-06-11 18:15:43 +00:00
2024-08-11 08:05:22 +00:00
return {
2024-06-11 06:40:27 +00:00
" model " : task_model_id ,
" messages " : [
{ " role " : " system " , " content " : content } ,
{ " role " : " user " , " content " : f " Query: { prompt } " } ,
] ,
" stream " : False ,
2024-08-10 12:04:01 +00:00
" metadata " : { " task " : str ( TASKS . FUNCTION_CALLING ) } ,
2024-06-11 06:40:27 +00:00
}
2024-08-11 08:05:22 +00:00
2024-08-11 13:56:16 +00:00
async def get_tool_call_response (
2024-08-11 08:05:22 +00:00
messages , files , tool_id , template , task_model_id , user , extra_params
) - > tuple [ Optional [ str ] , Optional [ dict ] , bool ] :
tool = Tools . get_tool_by_id ( tool_id )
if tool is None :
return None , None , False
tools_specs = json . dumps ( tool . specs , indent = 2 )
2024-08-11 13:56:16 +00:00
content = tool_calling_generation_template ( template , tools_specs )
payload = get_tool_calling_payload ( messages , task_model_id , content )
2024-08-11 08:05:22 +00:00
2024-06-12 20:34:34 +00:00
try :
payload = filter_pipeline ( payload , user )
except Exception as e :
raise e
2024-08-10 13:25:20 +00:00
if tool_id in webui_app . state . TOOLS :
toolkit_module = webui_app . state . TOOLS [ tool_id ]
else :
toolkit_module , _ = load_toolkit_module_by_id ( tool_id )
webui_app . state . TOOLS [ tool_id ] = toolkit_module
2024-08-11 07:31:40 +00:00
custom_params = {
* * extra_params ,
" __model__ " : app . state . MODELS [ task_model_id ] ,
" __id__ " : tool_id ,
" __messages__ " : messages ,
" __files__ " : files ,
2024-08-10 13:25:20 +00:00
}
try :
if hasattr ( toolkit_module , " UserValves " ) :
2024-08-11 07:31:40 +00:00
custom_params [ " __user__ " ] [ " valves " ] = toolkit_module . UserValves (
2024-08-10 13:25:20 +00:00
* * Tools . get_user_valves_by_id_and_user_id ( tool_id , user . id )
)
except Exception as e :
print ( e )
file_handler = hasattr ( toolkit_module , " file_handler " )
if hasattr ( toolkit_module , " valves " ) and hasattr ( toolkit_module , " Valves " ) :
valves = Tools . get_tool_valves_by_id ( tool_id )
toolkit_module . valves = toolkit_module . Valves ( * * ( valves if valves else { } ) )
2024-05-17 03:49:28 +00:00
2024-06-11 07:18:45 +00:00
try :
2024-06-24 18:17:18 +00:00
response = await generate_chat_completions ( form_data = payload , user = user )
2024-08-10 12:40:04 +00:00
content = await get_content_from_response ( response )
2024-06-11 07:18:45 +00:00
2024-07-09 10:51:43 +00:00
if content is None :
return None , None , False
2024-06-11 07:18:45 +00:00
# Parse the function response
2024-08-10 13:25:20 +00:00
log . debug ( f " content: { content } " )
2024-07-09 10:51:43 +00:00
result = json . loads ( content )
2024-06-11 07:18:45 +00:00
2024-08-10 13:25:20 +00:00
function_result = await call_tool_from_completion (
2024-08-11 07:31:40 +00:00
result , custom_params , toolkit_module
2024-08-10 13:25:20 +00:00
)
2024-06-24 02:18:13 +00:00
2024-08-10 13:25:20 +00:00
if hasattr ( toolkit_module , " citation " ) and toolkit_module . citation :
citation = {
" source " : { " name " : f " TOOL: { tool . name } / { result [ ' name ' ] } " } ,
" document " : [ function_result ] ,
" metadata " : [ { " source " : result [ " name " ] } ] ,
2024-07-09 10:51:43 +00:00
}
2024-08-10 13:25:20 +00:00
else :
citation = None
2024-07-09 10:51:43 +00:00
# Add the function result to the system prompt
if function_result is not None :
return function_result , citation , file_handler
2024-06-11 07:18:45 +00:00
except Exception as e :
print ( f " Error: { e } " )
2024-06-11 06:40:27 +00:00
2024-06-20 21:14:12 +00:00
return None , None , False
2024-06-11 06:40:27 +00:00
2024-08-11 07:31:40 +00:00
async def chat_completion_inlets_handler ( body , model , extra_params ) :
2024-07-02 02:33:58 +00:00
skip_files = None
filter_ids = get_filter_function_ids ( model )
for filter_id in filter_ids :
filter = Functions . get_function_by_id ( filter_id )
2024-07-09 10:51:43 +00:00
if not filter :
continue
2024-06-27 20:04:12 +00:00
2024-07-09 10:51:43 +00:00
if filter_id in webui_app . state . FUNCTIONS :
function_module = webui_app . state . FUNCTIONS [ filter_id ]
else :
function_module , _ , _ = load_function_module_by_id ( filter_id )
webui_app . state . FUNCTIONS [ filter_id ] = function_module
2024-07-02 02:33:58 +00:00
2024-07-09 10:51:43 +00:00
# Check if the function has a file_handler variable
if hasattr ( function_module , " file_handler " ) :
skip_files = function_module . file_handler
2024-07-02 02:33:58 +00:00
2024-07-09 10:51:43 +00:00
if hasattr ( function_module , " valves " ) and hasattr ( function_module , " Valves " ) :
valves = Functions . get_function_valves_by_id ( filter_id )
function_module . valves = function_module . Valves (
* * ( valves if valves else { } )
)
2024-07-09 11:15:09 +00:00
if not hasattr ( function_module , " inlet " ) :
continue
2024-07-09 10:51:43 +00:00
try :
2024-07-09 11:15:09 +00:00
inlet = function_module . inlet
# Get the signature of the function
sig = inspect . signature ( inlet )
params = { " body " : body }
# Extra parameters to be passed to the function
2024-08-11 07:31:40 +00:00
custom_params = { * * extra_params , " __model__ " : model , " __id__ " : filter_id }
if hasattr ( function_module , " UserValves " ) and " __user__ " in sig . parameters :
2024-08-11 08:07:12 +00:00
try :
uid = custom_params [ " __user__ " ] [ " id " ]
custom_params [ " __user__ " ] [ " valves " ] = function_module . UserValves (
* * Functions . get_user_valves_by_id_and_user_id ( filter_id , uid )
)
except Exception as e :
print ( e )
2024-07-09 11:15:09 +00:00
# Add extra params in contained in function signature
2024-08-11 07:31:40 +00:00
for key , value in custom_params . items ( ) :
2024-07-09 11:15:09 +00:00
if key in sig . parameters :
params [ key ] = value
if inspect . iscoroutinefunction ( inlet ) :
body = await inlet ( * * params )
else :
body = inlet ( * * params )
2024-07-09 10:51:43 +00:00
except Exception as e :
print ( f " Error: { e } " )
raise e
2024-06-20 09:30:00 +00:00
2024-07-02 02:33:58 +00:00
if skip_files :
if " files " in body :
del body [ " files " ]
2024-06-11 06:40:27 +00:00
2024-07-02 02:33:58 +00:00
return body , { }
2024-06-11 08:10:24 +00:00
2024-06-20 09:06:10 +00:00
2024-08-11 07:31:40 +00:00
async def chat_completion_tools_handler ( body , user , extra_params ) :
2024-07-02 02:33:58 +00:00
skip_files = None
2024-06-11 08:10:24 +00:00
2024-07-02 02:33:58 +00:00
contexts = [ ]
citations = None
2024-06-11 08:10:24 +00:00
2024-07-02 02:33:58 +00:00
task_model_id = get_task_model_id ( body [ " model " ] )
# If tool_ids field is present, call the functions
2024-08-10 10:58:18 +00:00
if " tool_ids " not in body :
return body , { }
2024-07-02 02:33:58 +00:00
2024-08-10 10:58:18 +00:00
print ( body [ " tool_ids " ] )
for tool_id in body [ " tool_ids " ] :
print ( tool_id )
try :
2024-08-11 13:56:16 +00:00
response , citation , file_handler = await get_tool_call_response (
2024-08-10 10:58:18 +00:00
messages = body [ " messages " ] ,
files = body . get ( " files " , [ ] ) ,
tool_id = tool_id ,
template = app . state . config . TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE ,
task_model_id = task_model_id ,
user = user ,
2024-08-11 07:31:40 +00:00
extra_params = extra_params ,
2024-08-10 10:58:18 +00:00
)
2024-07-02 02:33:58 +00:00
2024-08-10 10:58:18 +00:00
print ( file_handler )
if isinstance ( response , str ) :
contexts . append ( response )
2024-07-02 02:33:58 +00:00
2024-08-10 10:58:18 +00:00
if citation :
if citations is None :
citations = [ citation ]
else :
citations . append ( citation )
2024-07-02 02:33:58 +00:00
2024-08-10 10:58:18 +00:00
if file_handler :
skip_files = True
except Exception as e :
print ( f " Error: { e } " )
del body [ " tool_ids " ]
print ( f " tool_contexts: { contexts } " )
2024-07-02 02:33:58 +00:00
if skip_files :
if " files " in body :
del body [ " files " ]
return body , {
* * ( { " contexts " : contexts } if contexts is not None else { } ) ,
* * ( { " citations " : citations } if citations is not None else { } ) ,
}
async def chat_completion_files_handler ( body ) :
contexts = [ ]
citations = None
2024-08-11 07:31:40 +00:00
if files := body . pop ( " files " , None ) :
2024-07-02 02:33:58 +00:00
contexts , citations = get_rag_context (
files = files ,
messages = body [ " messages " ] ,
embedding_function = rag_app . state . EMBEDDING_FUNCTION ,
k = rag_app . state . config . TOP_K ,
reranking_function = rag_app . state . sentence_transformer_rf ,
r = rag_app . state . config . RELEVANCE_THRESHOLD ,
hybrid_search = rag_app . state . config . ENABLE_RAG_HYBRID_SEARCH ,
)
log . debug ( f " rag_contexts: { contexts } , citations: { citations } " )
return body , {
* * ( { " contexts " : contexts } if contexts is not None else { } ) ,
* * ( { " citations " : citations } if citations is not None else { } ) ,
}
2024-08-10 11:03:47 +00:00
def is_chat_completion_request ( request ) :
return request . method == " POST " and any (
endpoint in request . url . path
for endpoint in [ " /ollama/api/chat " , " /chat/completions " ]
)
2024-07-02 02:33:58 +00:00
class ChatCompletionMiddleware ( BaseHTTPMiddleware ) :
async def dispatch ( self , request : Request , call_next ) :
2024-08-10 11:03:47 +00:00
if not is_chat_completion_request ( request ) :
return await call_next ( request )
log . debug ( f " request.url.path: { request . url . path } " )
2024-07-02 02:33:58 +00:00
2024-08-10 11:03:47 +00:00
try :
body , model , user = await get_body_and_model_and_user ( request )
except Exception as e :
return JSONResponse (
status_code = status . HTTP_400_BAD_REQUEST ,
content = { " detail " : str ( e ) } ,
)
2024-07-31 13:01:40 +00:00
2024-08-10 11:03:47 +00:00
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 ) ,
}
2024-07-09 04:39:06 +00:00
2024-08-11 07:31:40 +00:00
__user__ = {
" id " : user . id ,
" email " : user . email ,
" name " : user . name ,
" role " : user . role ,
}
extra_params = {
" __user__ " : __user__ ,
" __event_emitter__ " : get_event_emitter ( metadata ) ,
" __event_call__ " : get_event_call ( metadata ) ,
}
2024-07-02 02:33:58 +00:00
2024-08-10 11:03:47 +00:00
# Initialize data_items to store additional data to be sent to the client
2024-08-10 11:33:27 +00:00
# Initalize contexts and citation
2024-08-10 11:03:47 +00:00
data_items = [ ]
contexts = [ ]
citations = [ ]
2024-07-02 02:33:58 +00:00
2024-08-10 11:03:47 +00:00
try :
2024-08-10 11:33:27 +00:00
body , flags = await chat_completion_inlets_handler (
2024-08-11 07:31:40 +00:00
body , model , extra_params
2024-08-10 11:03:47 +00:00
)
except Exception as e :
return JSONResponse (
status_code = status . HTTP_400_BAD_REQUEST ,
content = { " detail " : str ( e ) } ,
)
2024-07-02 02:33:58 +00:00
2024-08-10 11:03:47 +00:00
try :
2024-08-11 07:31:40 +00:00
body , flags = await chat_completion_tools_handler ( body , user , extra_params )
2024-08-10 11:03:47 +00:00
contexts . extend ( flags . get ( " contexts " , [ ] ) )
citations . extend ( flags . get ( " citations " , [ ] ) )
except Exception as e :
print ( e )
pass
2024-03-09 06:34:47 +00:00
2024-08-10 11:03:47 +00:00
try :
body , flags = await chat_completion_files_handler ( body )
contexts . extend ( flags . get ( " contexts " , [ ] ) )
citations . extend ( flags . get ( " citations " , [ ] ) )
except Exception as e :
print ( e )
pass
# If context is not empty, insert it into the messages
if len ( contexts ) > 0 :
context_string = " /n " . join ( contexts ) . strip ( )
prompt = get_last_user_message ( body [ " messages " ] )
2024-08-10 11:11:41 +00:00
if prompt is None :
raise Exception ( " No user message found " )
2024-08-10 11:03:47 +00:00
# Workaround for Ollama 2.0+ system prompt issue
# TODO: replace with add_or_update_system_message
if model [ " owned_by " ] == " ollama " :
body [ " messages " ] = prepend_to_first_user_message_content (
rag_template (
rag_app . state . config . RAG_TEMPLATE , context_string , prompt
) ,
body [ " messages " ] ,
)
2024-06-20 10:23:50 +00:00
else :
2024-08-10 11:03:47 +00:00
body [ " messages " ] = add_or_update_system_message (
rag_template (
rag_app . state . config . RAG_TEMPLATE , context_string , prompt
) ,
body [ " messages " ] ,
)
# If there are citations, add them to the data_items
if len ( citations ) > 0 :
data_items . append ( { " citations " : citations } )
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
# Set custom header to ensure content-length matches new body length
request . headers . __dict__ [ " _list " ] = [
( b " content-length " , str ( len ( modified_body_bytes ) ) . encode ( " utf-8 " ) ) ,
* [ ( k , v ) for k , v in request . headers . raw if k . lower ( ) != b " content-length " ] ,
]
2024-05-06 13:14:51 +00:00
2024-06-20 10:23:50 +00:00
response = await call_next ( request )
2024-08-10 11:03:47 +00:00
if isinstance ( response , StreamingResponse ) :
# If it's a streaming response, inject it as SSE event or NDJSON line
2024-08-11 07:31:40 +00:00
content_type = response . headers [ " Content-Type " ]
2024-08-10 11:03:47 +00:00
if " text/event-stream " in content_type :
return StreamingResponse (
self . openai_stream_wrapper ( response . body_iterator , data_items ) ,
)
if " application/x-ndjson " in content_type :
return StreamingResponse (
self . ollama_stream_wrapper ( response . body_iterator , data_items ) ,
)
2024-08-10 11:11:41 +00:00
return response
2024-03-09 06:34:47 +00:00
async def _receive ( self , body : bytes ) :
return { " type " : " http.request " , " body " : body , " more_body " : False }
2024-06-20 09:06:10 +00:00
async def openai_stream_wrapper ( self , original_generator , data_items ) :
for item in data_items :
yield f " data: { json . dumps ( item ) } \n \n "
2024-05-06 13:14:51 +00:00
async for data in original_generator :
yield data
2024-06-20 09:06:10 +00:00
async def ollama_stream_wrapper ( self , original_generator , data_items ) :
for item in data_items :
yield f " { json . dumps ( item ) } \n "
2024-05-06 13:14:51 +00:00
async for data in original_generator :
yield data
2024-03-09 06:34:47 +00:00
2024-06-11 06:40:27 +00:00
app . add_middleware ( ChatCompletionMiddleware )
2024-03-09 06:34:47 +00:00
2024-06-20 08:51:39 +00:00
##################################
#
# Pipeline Middleware
#
##################################
2024-03-09 06:34:47 +00:00
2024-07-09 11:51:13 +00:00
def get_sorted_filters ( model_id ) :
2024-06-09 21:25:31 +00:00
filters = [
model
for model in app . state . MODELS . values ( )
if " pipeline " in model
and " type " in model [ " pipeline " ]
and model [ " pipeline " ] [ " type " ] == " filter "
and (
model [ " pipeline " ] [ " pipelines " ] == [ " * " ]
or any (
model_id == target_model_id
for target_model_id in model [ " pipeline " ] [ " pipelines " ]
)
)
]
sorted_filters = sorted ( filters , key = lambda x : x [ " pipeline " ] [ " priority " ] )
2024-07-09 11:51:13 +00:00
return sorted_filters
def filter_pipeline ( payload , user ) :
user = { " id " : user . id , " email " : user . email , " name " : user . name , " role " : user . role }
model_id = payload [ " model " ]
sorted_filters = get_sorted_filters ( model_id )
2024-06-09 21:25:31 +00:00
model = app . state . MODELS [ model_id ]
if " pipeline " in model :
sorted_filters . append ( model )
for filter in sorted_filters :
r = None
try :
urlIdx = filter [ " urlIdx " ]
url = openai_app . state . config . OPENAI_API_BASE_URLS [ urlIdx ]
key = openai_app . state . config . OPENAI_API_KEYS [ urlIdx ]
2024-08-10 12:04:01 +00:00
if key == " " :
continue
headers = { " Authorization " : f " Bearer { key } " }
r = requests . post (
f " { url } / { filter [ ' id ' ] } /filter/inlet " ,
headers = headers ,
json = {
" user " : user ,
" body " : payload ,
} ,
)
2024-06-09 21:25:31 +00:00
2024-08-10 12:04:01 +00:00
r . raise_for_status ( )
payload = r . json ( )
2024-06-09 21:25:31 +00:00
except Exception as e :
# Handle connection error here
print ( f " Connection error: { e } " )
if r is not None :
2024-07-09 11:51:13 +00:00
res = r . json ( )
2024-06-12 20:31:05 +00:00
if " detail " in res :
raise Exception ( r . status_code , res [ " detail " ] )
2024-06-09 21:25:31 +00:00
return payload
2024-05-28 02:03:26 +00:00
class PipelineMiddleware ( BaseHTTPMiddleware ) :
async def dispatch ( self , request : Request , call_next ) :
2024-08-10 11:03:47 +00:00
if not is_chat_completion_request ( request ) :
return await call_next ( request )
2024-06-12 20:31:05 +00:00
2024-08-10 11:03:47 +00:00
log . debug ( f " request.url.path: { request . url . path } " )
2024-05-30 09:04:29 +00:00
2024-08-10 11:03:47 +00:00
# Read the original request body
body = await request . body ( )
# Decode body to string
body_str = body . decode ( " utf-8 " )
# Parse string to JSON
data = json . loads ( body_str ) if body_str else { }
user = get_current_user (
request ,
2024-08-11 07:31:40 +00:00
get_http_authorization_cred ( request . headers [ " Authorization " ] ) ,
2024-08-10 11:03:47 +00:00
)
try :
data = filter_pipeline ( data , user )
except Exception as e :
return JSONResponse (
status_code = e . args [ 0 ] ,
content = { " detail " : e . args [ 1 ] } ,
)
modified_body_bytes = json . dumps ( data ) . encode ( " utf-8 " )
# Replace the request body with the modified one
request . _body = modified_body_bytes
# Set custom header to ensure content-length matches new body length
request . headers . __dict__ [ " _list " ] = [
( b " content-length " , str ( len ( modified_body_bytes ) ) . encode ( " utf-8 " ) ) ,
* [ ( k , v ) for k , v in request . headers . raw if k . lower ( ) != b " content-length " ] ,
]
2024-05-28 02:03:26 +00:00
response = await call_next ( request )
return response
async def _receive ( self , body : bytes ) :
return { " type " : " http.request " , " body " : body , " more_body " : False }
app . add_middleware ( PipelineMiddleware )
2024-05-28 16:50:17 +00:00
app . add_middleware (
CORSMiddleware ,
allow_origins = origins ,
allow_credentials = True ,
allow_methods = [ " * " ] ,
allow_headers = [ " * " ] ,
)
2024-06-24 11:06:15 +00:00
@app.middleware ( " http " )
2024-06-24 11:45:33 +00:00
async def commit_session_after_request ( request : Request , call_next ) :
2024-06-24 11:06:15 +00:00
response = await call_next ( request )
2024-06-24 11:45:33 +00:00
log . debug ( " Commit session after request " )
2024-06-24 11:06:15 +00:00
Session . commit ( )
return response
2024-05-28 16:50:17 +00:00
2023-11-15 00:28:51 +00:00
@app.middleware ( " http " )
async def check_url ( request : Request , call_next ) :
2024-05-25 01:26:36 +00:00
if len ( app . state . MODELS ) == 0 :
await get_all_models ( )
else :
pass
2023-11-15 00:28:51 +00:00
start_time = int ( time . time ( ) )
response = await call_next ( request )
process_time = int ( time . time ( ) ) - start_time
response . headers [ " X-Process-Time " ] = str ( process_time )
return response
2024-05-19 15:00:07 +00:00
@app.middleware ( " http " )
async def update_embedding_function ( request : Request , call_next ) :
response = await call_next ( request )
if " /embedding/update " in request . url . path :
webui_app . state . EMBEDDING_FUNCTION = rag_app . state . EMBEDDING_FUNCTION
return response
2024-02-22 11:22:23 +00:00
2024-05-19 15:00:07 +00:00
2024-06-04 06:39:52 +00:00
app . mount ( " /ws " , socket_app )
2024-03-04 03:55:32 +00:00
app . mount ( " /ollama " , ollama_app )
2024-05-24 08:40:48 +00:00
app . mount ( " /openai " , openai_app )
2024-02-11 08:17:50 +00:00
2024-02-22 02:12:01 +00:00
app . mount ( " /images/api/v1 " , images_app )
2024-02-11 08:17:50 +00:00
app . mount ( " /audio/api/v1 " , audio_app )
2024-01-07 06:07:20 +00:00
app . mount ( " /rag/api/v1 " , rag_app )
2024-05-19 15:00:07 +00:00
app . mount ( " /api/v1 " , webui_app )
webui_app . state . EMBEDDING_FUNCTION = rag_app . state . EMBEDDING_FUNCTION
2024-03-31 20:59:39 +00:00
2024-05-25 01:26:36 +00:00
async def get_all_models ( ) :
2024-07-12 01:41:00 +00:00
# TODO: Optimize this function
2024-06-20 11:21:55 +00:00
pipe_models = [ ]
2024-05-24 08:40:48 +00:00
openai_models = [ ]
ollama_models = [ ]
2024-06-20 11:21:55 +00:00
pipe_models = await get_pipe_models ( )
2024-05-24 08:40:48 +00:00
if app . state . config . ENABLE_OPENAI_API :
openai_models = await get_openai_models ( )
openai_models = openai_models [ " data " ]
if app . state . config . ENABLE_OLLAMA_API :
ollama_models = await get_ollama_models ( )
ollama_models = [
{
" id " : model [ " model " ] ,
" name " : model [ " name " ] ,
" object " : " model " ,
" created " : int ( time . time ( ) ) ,
" owned_by " : " ollama " ,
" ollama " : model ,
}
for model in ollama_models [ " models " ]
]
2024-06-20 11:21:55 +00:00
models = pipe_models + openai_models + ollama_models
2024-05-24 08:40:48 +00:00
2024-07-12 01:41:00 +00:00
global_action_ids = [
function . id for function in Functions . get_global_action_functions ( )
]
enabled_action_ids = [
function . id
for function in Functions . get_functions_by_type ( " action " , active_only = True )
]
2024-06-20 11:21:55 +00:00
custom_models = Models . get_all_models ( )
2024-05-24 08:40:48 +00:00
for custom_model in custom_models :
2024-08-03 13:24:26 +00:00
if custom_model . base_model_id is None :
2024-05-24 08:40:48 +00:00
for model in models :
2024-05-24 09:11:17 +00:00
if (
custom_model . id == model [ " id " ]
or custom_model . id == model [ " id " ] . split ( " : " ) [ 0 ]
) :
2024-05-24 08:40:48 +00:00
model [ " name " ] = custom_model . name
model [ " info " ] = custom_model . model_dump ( )
2024-07-12 01:41:00 +00:00
2024-07-28 12:00:58 +00:00
action_ids = [ ]
2024-07-12 01:41:00 +00:00
if " info " in model and " meta " in model [ " info " ] :
action_ids . extend ( model [ " info " ] [ " meta " ] . get ( " actionIds " , [ ] ) )
2024-07-28 12:00:58 +00:00
model [ " action_ids " ] = action_ids
2024-05-24 08:40:48 +00:00
else :
2024-05-24 10:06:57 +00:00
owned_by = " openai "
2024-07-04 20:41:18 +00:00
pipe = None
2024-07-28 12:00:58 +00:00
action_ids = [ ]
2024-07-04 20:41:18 +00:00
2024-05-24 10:06:57 +00:00
for model in models :
2024-05-25 03:29:13 +00:00
if (
custom_model . base_model_id == model [ " id " ]
or custom_model . base_model_id == model [ " id " ] . split ( " : " ) [ 0 ]
) :
2024-05-24 10:06:57 +00:00
owned_by = model [ " owned_by " ]
2024-07-04 20:41:18 +00:00
if " pipe " in model :
pipe = model [ " pipe " ]
2024-07-12 01:41:00 +00:00
if " info " in model and " meta " in model [ " info " ] :
action_ids . extend ( model [ " info " ] [ " meta " ] . get ( " actionIds " , [ ] ) )
2024-05-24 10:06:57 +00:00
break
2024-05-24 08:40:48 +00:00
models . append (
{
" id " : custom_model . id ,
" name " : custom_model . name ,
" object " : " model " ,
" created " : custom_model . created_at ,
2024-05-24 10:06:57 +00:00
" owned_by " : owned_by ,
2024-05-24 08:40:48 +00:00
" info " : custom_model . model_dump ( ) ,
2024-05-25 03:29:13 +00:00
" preset " : True ,
2024-07-04 20:41:18 +00:00
* * ( { " pipe " : pipe } if pipe is not None else { } ) ,
2024-07-28 12:00:58 +00:00
" action_ids " : action_ids ,
}
)
for model in models :
action_ids = [ ]
if " action_ids " in model :
action_ids = model [ " action_ids " ]
del model [ " action_ids " ]
action_ids = action_ids + global_action_ids
action_ids = list ( set ( action_ids ) )
action_ids = [
action_id for action_id in action_ids if action_id in enabled_action_ids
]
model [ " actions " ] = [ ]
for action_id in action_ids :
action = Functions . get_function_by_id ( action_id )
2024-08-11 07:31:40 +00:00
if action is None :
raise Exception ( f " Action not found: { action_id } " )
2024-07-28 21:02:23 +00:00
if action_id in webui_app . state . FUNCTIONS :
function_module = webui_app . state . FUNCTIONS [ action_id ]
else :
function_module , _ , _ = load_function_module_by_id ( action_id )
webui_app . state . FUNCTIONS [ action_id ] = function_module
2024-08-11 07:31:40 +00:00
icon_url = None
if action . meta . manifest is not None :
icon_url = action . meta . manifest . get ( " icon_url " , None )
2024-07-28 21:02:23 +00:00
if hasattr ( function_module , " actions " ) :
actions = function_module . actions
model [ " actions " ] . extend (
[
{
" id " : f " { action_id } . { _action [ ' id ' ] } " ,
" name " : _action . get (
" name " , f " { action . name } ( { _action [ ' id ' ] } ) "
) ,
" description " : action . meta . description ,
2024-08-11 07:31:40 +00:00
" icon_url " : _action . get ( " icon_url " , icon_url ) ,
2024-07-28 21:02:23 +00:00
}
for _action in actions
]
)
else :
model [ " actions " ] . append (
{
" id " : action_id ,
" name " : action . name ,
" description " : action . meta . description ,
2024-08-11 07:31:40 +00:00
" icon_url " : icon_url ,
2024-07-28 21:02:23 +00:00
}
)
2024-05-24 08:40:48 +00:00
2024-05-25 01:26:36 +00:00
app . state . MODELS = { model [ " id " ] : model for model in models }
webui_app . state . MODELS = app . state . MODELS
return models
@app.get ( " /api/models " )
async def get_models ( user = Depends ( get_verified_user ) ) :
models = await get_all_models ( )
2024-05-28 02:03:26 +00:00
2024-05-28 02:34:05 +00:00
# Filter out filter pipelines
2024-05-28 02:03:26 +00:00
models = [
model
for model in models
2024-05-28 18:43:48 +00:00
if " pipeline " not in model or model [ " pipeline " ] . get ( " type " , None ) != " filter "
2024-05-28 02:03:26 +00:00
]
2024-05-24 08:40:48 +00:00
if app . state . config . ENABLE_MODEL_FILTER :
if user . role == " user " :
models = list (
filter (
lambda model : model [ " id " ] in app . state . config . MODEL_FILTER_LIST ,
models ,
)
)
return { " data " : models }
return { " data " : models }
2024-06-20 08:51:39 +00:00
@app.post ( " /api/chat/completions " )
async def generate_chat_completions ( form_data : dict , user = Depends ( get_verified_user ) ) :
model_id = form_data [ " model " ]
if model_id not in app . state . MODELS :
raise HTTPException (
status_code = status . HTTP_404_NOT_FOUND ,
detail = " Model not found " ,
)
model = app . state . MODELS [ model_id ]
2024-07-11 20:43:44 +00:00
if model . get ( " pipe " ) :
2024-06-24 18:17:18 +00:00
return await generate_function_chat_completion ( form_data , user = user )
2024-06-20 08:51:39 +00:00
if model [ " owned_by " ] == " ollama " :
return await generate_ollama_chat_completion ( form_data , user = user )
else :
return await generate_openai_chat_completion ( form_data , user = user )
@app.post ( " /api/chat/completed " )
async def chat_completed ( form_data : dict , user = Depends ( get_verified_user ) ) :
data = form_data
model_id = data [ " model " ]
2024-06-20 10:23:50 +00:00
if model_id not in app . state . MODELS :
raise HTTPException (
status_code = status . HTTP_404_NOT_FOUND ,
detail = " Model not found " ,
)
model = app . state . MODELS [ model_id ]
2024-06-20 08:51:39 +00:00
2024-07-09 11:51:13 +00:00
sorted_filters = get_sorted_filters ( model_id )
2024-06-20 10:23:50 +00:00
if " pipeline " in model :
sorted_filters = [ model ] + sorted_filters
2024-06-20 08:51:39 +00:00
for filter in sorted_filters :
r = None
try :
urlIdx = filter [ " urlIdx " ]
url = openai_app . state . config . OPENAI_API_BASE_URLS [ urlIdx ]
key = openai_app . state . config . OPENAI_API_KEYS [ urlIdx ]
if key != " " :
headers = { " Authorization " : f " Bearer { key } " }
r = requests . post (
f " { url } / { filter [ ' id ' ] } /filter/outlet " ,
headers = headers ,
json = {
2024-06-22 19:14:12 +00:00
" user " : {
" id " : user . id ,
" name " : user . name ,
" email " : user . email ,
" role " : user . role ,
} ,
2024-06-20 08:51:39 +00:00
" 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 JSONResponse (
status_code = r . status_code ,
content = res ,
)
2024-07-31 12:35:02 +00:00
except Exception :
2024-06-20 08:51:39 +00:00
pass
else :
pass
2024-07-31 12:35:02 +00:00
__event_emitter__ = get_event_emitter (
2024-07-11 17:40:10 +00:00
{
" chat_id " : data [ " chat_id " ] ,
" message_id " : data [ " id " ] ,
" session_id " : data [ " session_id " ] ,
}
)
2024-07-02 04:41:44 +00:00
2024-07-31 12:35:02 +00:00
__event_call__ = get_event_call (
2024-07-11 17:40:10 +00:00
{
" chat_id " : data [ " chat_id " ] ,
" message_id " : data [ " id " ] ,
" session_id " : data [ " session_id " ] ,
}
)
2024-07-09 04:40:22 +00:00
2024-06-24 03:11:08 +00:00
def get_priority ( function_id ) :
function = Functions . get_function_by_id ( function_id )
if function is not None and hasattr ( function , " valves " ) :
2024-08-11 07:31:40 +00:00
# TODO: Fix FunctionModel to include vavles
2024-06-24 03:11:08 +00:00
return ( function . valves if function . valves else { } ) . get ( " priority " , 0 )
return 0
2024-06-27 20:04:12 +00:00
filter_ids = [ function . id for function in Functions . get_global_filter_functions ( ) ]
2024-06-20 11:38:59 +00:00
if " info " in model and " meta " in model [ " info " ] :
2024-06-27 20:04:12 +00:00
filter_ids . extend ( model [ " info " ] [ " meta " ] . get ( " filterIds " , [ ] ) )
2024-06-24 03:11:08 +00:00
filter_ids = list ( set ( filter_ids ) )
2024-06-20 10:23:50 +00:00
2024-06-27 20:04:12 +00:00
enabled_filter_ids = [
function . id
for function in Functions . get_functions_by_type ( " filter " , active_only = True )
]
filter_ids = [
filter_id for filter_id in filter_ids if filter_id in enabled_filter_ids
]
2024-06-24 03:11:08 +00:00
# Sort filter_ids by priority, using the get_priority function
filter_ids . sort ( key = get_priority )
2024-06-24 02:18:13 +00:00
2024-06-24 03:11:08 +00:00
for filter_id in filter_ids :
filter = Functions . get_function_by_id ( filter_id )
2024-07-09 10:51:43 +00:00
if not filter :
continue
2024-06-24 03:11:08 +00:00
2024-07-09 10:51:43 +00:00
if filter_id in webui_app . state . FUNCTIONS :
function_module = webui_app . state . FUNCTIONS [ filter_id ]
else :
function_module , _ , _ = load_function_module_by_id ( filter_id )
webui_app . state . FUNCTIONS [ filter_id ] = function_module
2024-06-24 03:11:08 +00:00
2024-07-09 10:51:43 +00:00
if hasattr ( function_module , " valves " ) and hasattr ( function_module , " Valves " ) :
valves = Functions . get_function_valves_by_id ( filter_id )
function_module . valves = function_module . Valves (
* * ( valves if valves else { } )
)
2024-06-24 03:11:08 +00:00
2024-07-09 11:15:09 +00:00
if not hasattr ( function_module , " outlet " ) :
continue
2024-07-09 10:51:43 +00:00
try :
2024-07-09 11:15:09 +00:00
outlet = function_module . outlet
# Get the signature of the function
sig = inspect . signature ( outlet )
params = { " body " : data }
# Extra parameters to be passed to the function
extra_params = {
" __model__ " : model ,
" __id__ " : filter_id ,
" __event_emitter__ " : __event_emitter__ ,
" __event_call__ " : __event_call__ ,
}
# Add extra params in contained in function signature
for key , value in extra_params . items ( ) :
if key in sig . parameters :
params [ key ] = value
if " __user__ " in sig . parameters :
__user__ = {
" id " : user . id ,
" email " : user . email ,
" name " : user . name ,
" role " : user . role ,
2024-07-09 10:51:43 +00:00
}
2024-07-09 11:15:09 +00:00
try :
if hasattr ( function_module , " UserValves " ) :
__user__ [ " valves " ] = function_module . UserValves (
* * Functions . get_user_valves_by_id_and_user_id (
filter_id , user . id
2024-07-09 10:51:43 +00:00
)
2024-07-09 11:15:09 +00:00
)
except Exception as e :
print ( e )
2024-06-21 03:26:28 +00:00
2024-07-09 11:15:09 +00:00
params = { * * params , " __user__ " : __user__ }
2024-07-09 10:51:43 +00:00
2024-07-09 11:15:09 +00:00
if inspect . iscoroutinefunction ( outlet ) :
data = await outlet ( * * params )
else :
data = outlet ( * * params )
2024-07-09 10:51:43 +00:00
except Exception as e :
print ( f " Error: { e } " )
return JSONResponse (
status_code = status . HTTP_400_BAD_REQUEST ,
content = { " detail " : str ( e ) } ,
)
2024-06-20 10:23:50 +00:00
2024-06-20 08:51:39 +00:00
return data
2024-07-12 01:41:00 +00:00
@app.post ( " /api/chat/actions/ {action_id} " )
2024-08-02 15:45:30 +00:00
async def chat_action ( action_id : str , form_data : dict , user = Depends ( get_verified_user ) ) :
2024-07-28 21:02:23 +00:00
if " . " in action_id :
action_id , sub_action_id = action_id . split ( " . " )
else :
sub_action_id = None
2024-07-12 01:41:00 +00:00
action = Functions . get_function_by_id ( action_id )
if not action :
raise HTTPException (
status_code = status . HTTP_404_NOT_FOUND ,
detail = " Action not found " ,
)
2024-07-12 02:05:59 +00:00
2024-07-12 01:41:00 +00:00
data = form_data
model_id = data [ " model " ]
if model_id not in app . state . MODELS :
raise HTTPException (
status_code = status . HTTP_404_NOT_FOUND ,
detail = " Model not found " ,
)
model = app . state . MODELS [ model_id ]
2024-07-31 12:35:02 +00:00
__event_emitter__ = get_event_emitter (
2024-07-12 01:41:00 +00:00
{
" chat_id " : data [ " chat_id " ] ,
" message_id " : data [ " id " ] ,
" session_id " : data [ " session_id " ] ,
}
)
2024-07-31 12:35:02 +00:00
__event_call__ = get_event_call (
2024-07-12 01:41:00 +00:00
{
" chat_id " : data [ " chat_id " ] ,
" message_id " : data [ " id " ] ,
" session_id " : data [ " session_id " ] ,
}
)
if action_id in webui_app . state . FUNCTIONS :
function_module = webui_app . state . FUNCTIONS [ action_id ]
else :
function_module , _ , _ = load_function_module_by_id ( action_id )
webui_app . state . FUNCTIONS [ action_id ] = function_module
if hasattr ( function_module , " valves " ) and hasattr ( function_module , " Valves " ) :
valves = Functions . get_function_valves_by_id ( action_id )
function_module . valves = function_module . Valves ( * * ( valves if valves else { } ) )
if hasattr ( function_module , " action " ) :
try :
action = function_module . action
# Get the signature of the function
sig = inspect . signature ( action )
params = { " body " : data }
# Extra parameters to be passed to the function
extra_params = {
" __model__ " : model ,
2024-07-28 21:02:23 +00:00
" __id__ " : sub_action_id if sub_action_id is not None else action_id ,
2024-07-12 01:41:00 +00:00
" __event_emitter__ " : __event_emitter__ ,
" __event_call__ " : __event_call__ ,
}
# Add extra params in contained in function signature
for key , value in extra_params . items ( ) :
if key in sig . parameters :
params [ key ] = value
if " __user__ " in sig . parameters :
__user__ = {
" id " : user . id ,
" email " : user . email ,
" name " : user . name ,
" role " : user . role ,
}
try :
if hasattr ( function_module , " UserValves " ) :
__user__ [ " valves " ] = function_module . UserValves (
* * Functions . get_user_valves_by_id_and_user_id (
action_id , user . id
)
)
except Exception as e :
print ( e )
params = { * * params , " __user__ " : __user__ }
if inspect . iscoroutinefunction ( action ) :
data = await action ( * * params )
else :
data = action ( * * params )
except Exception as e :
print ( f " Error: { e } " )
return JSONResponse (
status_code = status . HTTP_400_BAD_REQUEST ,
content = { " detail " : str ( e ) } ,
)
return data
2024-06-20 08:51:39 +00:00
##################################
#
# Task Endpoints
#
##################################
# TODO: Refactor task API endpoints below into a separate file
2024-06-09 21:53:10 +00:00
@app.get ( " /api/task/config " )
async def get_task_config ( user = Depends ( get_verified_user ) ) :
return {
" TASK_MODEL " : app . state . config . TASK_MODEL ,
" TASK_MODEL_EXTERNAL " : app . state . config . TASK_MODEL_EXTERNAL ,
" TITLE_GENERATION_PROMPT_TEMPLATE " : app . state . config . TITLE_GENERATION_PROMPT_TEMPLATE ,
" SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE " : app . state . config . SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE ,
2024-06-09 22:29:55 +00:00
" SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD " : app . state . config . SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD ,
2024-06-11 06:40:27 +00:00
" TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE " : app . state . config . TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE ,
2024-06-09 21:53:10 +00:00
}
class TaskConfigForm ( BaseModel ) :
TASK_MODEL : Optional [ str ]
TASK_MODEL_EXTERNAL : Optional [ str ]
TITLE_GENERATION_PROMPT_TEMPLATE : str
SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE : str
2024-06-09 22:29:55 +00:00
SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD : int
2024-06-11 06:40:27 +00:00
TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE : str
2024-06-09 21:53:10 +00:00
@app.post ( " /api/task/config/update " )
async def update_task_config ( form_data : TaskConfigForm , user = Depends ( get_admin_user ) ) :
app . state . config . TASK_MODEL = form_data . TASK_MODEL
app . state . config . TASK_MODEL_EXTERNAL = form_data . TASK_MODEL_EXTERNAL
app . state . config . TITLE_GENERATION_PROMPT_TEMPLATE = (
form_data . TITLE_GENERATION_PROMPT_TEMPLATE
)
app . state . config . SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE = (
form_data . SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE
)
2024-06-09 22:29:55 +00:00
app . state . config . SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD = (
form_data . SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD
)
2024-06-11 06:40:27 +00:00
app . state . config . TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE = (
form_data . TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE
)
2024-06-09 21:53:10 +00:00
return {
" TASK_MODEL " : app . state . config . TASK_MODEL ,
" TASK_MODEL_EXTERNAL " : app . state . config . TASK_MODEL_EXTERNAL ,
" TITLE_GENERATION_PROMPT_TEMPLATE " : app . state . config . TITLE_GENERATION_PROMPT_TEMPLATE ,
" SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE " : app . state . config . SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE ,
2024-06-09 22:29:55 +00:00
" SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD " : app . state . config . SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD ,
2024-06-11 06:40:27 +00:00
" TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE " : app . state . config . TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE ,
2024-06-09 21:53:10 +00:00
}
2024-06-09 21:26:49 +00:00
@app.post ( " /api/task/title/completions " )
2024-06-09 21:25:31 +00:00
async def generate_title ( form_data : dict , user = Depends ( get_verified_user ) ) :
print ( " generate_title " )
2024-06-09 21:53:10 +00:00
2024-06-09 21:25:31 +00:00
model_id = form_data [ " model " ]
if model_id not in app . state . MODELS :
raise HTTPException (
status_code = status . HTTP_404_NOT_FOUND ,
detail = " Model not found " ,
)
2024-06-09 21:53:10 +00:00
# Check if the user has a custom task model
# If the user has a custom task model, use that model
2024-07-09 15:08:54 +00:00
model_id = get_task_model_id ( model_id )
2024-06-09 21:53:10 +00:00
print ( model_id )
2024-06-09 21:25:31 +00:00
template = app . state . config . TITLE_GENERATION_PROMPT_TEMPLATE
content = title_generation_template (
2024-06-16 22:32:26 +00:00
template ,
form_data [ " prompt " ] ,
{
" name " : user . name ,
" location " : user . info . get ( " location " ) if user . info else None ,
} ,
2024-06-09 21:25:31 +00:00
)
payload = {
" model " : model_id ,
" messages " : [ { " role " : " user " , " content " : content } ] ,
" stream " : False ,
" max_tokens " : 50 ,
" chat_id " : form_data . get ( " chat_id " , None ) ,
2024-08-10 12:04:01 +00:00
" metadata " : { " task " : str ( TASKS . TITLE_GENERATION ) } ,
2024-06-09 21:25:31 +00:00
}
2024-06-16 15:40:16 +00:00
log . debug ( payload )
2024-06-12 20:34:34 +00:00
try :
payload = filter_pipeline ( payload , user )
except Exception as e :
return JSONResponse (
status_code = e . args [ 0 ] ,
content = { " detail " : e . args [ 1 ] } ,
)
2024-06-09 21:25:31 +00:00
2024-07-02 02:33:58 +00:00
if " chat_id " in payload :
del payload [ " chat_id " ]
2024-06-24 18:17:18 +00:00
return await generate_chat_completions ( form_data = payload , user = user )
2024-06-09 21:25:31 +00:00
2024-06-09 21:53:10 +00:00
@app.post ( " /api/task/query/completions " )
async def generate_search_query ( form_data : dict , user = Depends ( get_verified_user ) ) :
print ( " generate_search_query " )
2024-06-09 22:19:36 +00:00
if len ( form_data [ " prompt " ] ) < app . state . config . SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD :
raise HTTPException (
status_code = status . HTTP_400_BAD_REQUEST ,
detail = f " Skip search query generation for short prompts (< { app . state . config . SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD } characters) " ,
)
2024-06-09 21:53:10 +00:00
model_id = form_data [ " model " ]
if model_id not in app . state . 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
2024-07-09 15:08:54 +00:00
model_id = get_task_model_id ( model_id )
2024-06-09 21:53:10 +00:00
print ( model_id )
template = app . state . config . SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE
content = search_query_generation_template (
2024-06-16 22:32:26 +00:00
template , form_data [ " prompt " ] , { " name " : user . name }
2024-06-09 21:53:10 +00:00
)
payload = {
" model " : model_id ,
" messages " : [ { " role " : " user " , " content " : content } ] ,
" stream " : False ,
" max_tokens " : 30 ,
2024-08-10 12:04:01 +00:00
" metadata " : { " task " : str ( TASKS . QUERY_GENERATION ) } ,
2024-06-13 04:18:53 +00:00
}
print ( payload )
try :
payload = filter_pipeline ( payload , user )
except Exception as e :
return JSONResponse (
status_code = e . args [ 0 ] ,
content = { " detail " : e . args [ 1 ] } ,
)
2024-07-02 02:33:58 +00:00
if " chat_id " in payload :
del payload [ " chat_id " ]
2024-06-24 18:17:18 +00:00
return await generate_chat_completions ( form_data = payload , user = user )
2024-06-13 04:18:53 +00:00
@app.post ( " /api/task/emoji/completions " )
async def generate_emoji ( form_data : dict , user = Depends ( get_verified_user ) ) :
print ( " generate_emoji " )
model_id = form_data [ " model " ]
if model_id not in app . state . 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
2024-07-09 15:08:54 +00:00
model_id = get_task_model_id ( model_id )
2024-06-13 04:18:53 +00:00
print ( model_id )
template = '''
2024-06-13 09:29:56 +00:00
Your task is to reflect the speaker ' s likely facial expression through a fitting emoji. Interpret emotions from the message and reflect their facial expression using fitting, diverse emojis (e.g., 😊, 😢, 😡, 😱).
2024-06-13 04:18:53 +00:00
Message : """ {{ prompt}} """
'''
content = title_generation_template (
2024-06-16 22:32:26 +00:00
template ,
form_data [ " prompt " ] ,
{
" name " : user . name ,
" location " : user . info . get ( " location " ) if user . info else None ,
} ,
2024-06-13 04:18:53 +00:00
)
payload = {
" model " : model_id ,
" messages " : [ { " role " : " user " , " content " : content } ] ,
" stream " : False ,
" max_tokens " : 4 ,
" chat_id " : form_data . get ( " chat_id " , None ) ,
2024-08-10 12:04:01 +00:00
" metadata " : { " task " : str ( TASKS . EMOJI_GENERATION ) } ,
2024-06-09 21:53:10 +00:00
}
2024-06-16 15:40:16 +00:00
log . debug ( payload )
2024-06-12 20:34:34 +00:00
try :
payload = filter_pipeline ( payload , user )
except Exception as e :
return JSONResponse (
status_code = e . args [ 0 ] ,
content = { " detail " : e . args [ 1 ] } ,
)
2024-06-09 21:53:10 +00:00
2024-07-02 02:33:58 +00:00
if " chat_id " in payload :
del payload [ " chat_id " ]
2024-06-24 18:17:18 +00:00
return await generate_chat_completions ( form_data = payload , user = user )
2024-06-09 21:53:10 +00:00
2024-06-20 08:51:39 +00:00
##################################
#
# Pipelines Endpoints
#
##################################
2024-05-30 09:04:29 +00:00
2024-06-20 08:51:39 +00:00
# TODO: Refactor pipelines API endpoints below into a separate file
2024-05-30 09:04:29 +00:00
2024-05-30 04:26:57 +00:00
@app.get ( " /api/pipelines/list " )
async def get_pipelines_list ( user = Depends ( get_admin_user ) ) :
2024-08-04 15:42:16 +00:00
responses = await get_openai_models ( raw = True )
2024-05-30 05:41:51 +00:00
print ( responses )
2024-06-02 23:46:33 +00:00
urlIdxs = [
idx
for idx , response in enumerate ( responses )
2024-08-03 13:24:26 +00:00
if response is not None and " pipelines " in response
2024-06-02 23:46:33 +00:00
]
2024-05-30 04:26:57 +00:00
return {
" data " : [
{
" url " : openai_app . state . config . OPENAI_API_BASE_URLS [ urlIdx ] ,
" idx " : urlIdx ,
}
for urlIdx in urlIdxs
]
}
2024-06-05 20:57:48 +00:00
@app.post ( " /api/pipelines/upload " )
async def upload_pipeline (
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
2024-08-11 07:31:40 +00:00
if not ( file . filename and file . filename . endswith ( " .py " ) ) :
2024-06-05 20:57:48 +00:00
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 )
2024-07-09 10:51:43 +00:00
r = None
2024-06-05 20:57:48 +00:00
try :
# Save the uploaded file
with open ( file_path , " wb " ) as buffer :
shutil . copyfileobj ( file . file , buffer )
url = openai_app . state . config . OPENAI_API_BASE_URLS [ urlIdx ]
key = openai_app . state . config . OPENAI_API_KEYS [ urlIdx ]
headers = { " Authorization " : f " Bearer { key } " }
with open ( file_path , " rb " ) as f :
files = { " file " : f }
r = requests . post ( f " { url } /pipelines/upload " , headers = headers , 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 = " Pipeline not found "
2024-07-09 10:51:43 +00:00
status_code = status . HTTP_404_NOT_FOUND
2024-06-05 20:57:48 +00:00
if r is not None :
2024-07-09 10:51:43 +00:00
status_code = r . status_code
2024-06-05 20:57:48 +00:00
try :
res = r . json ( )
if " detail " in res :
detail = res [ " detail " ]
2024-08-03 13:24:26 +00:00
except Exception :
2024-06-05 20:57:48 +00:00
pass
raise HTTPException (
2024-07-09 10:51:43 +00:00
status_code = status_code ,
2024-06-05 20:57:48 +00:00
detail = detail ,
)
finally :
# Ensure the file is deleted after the upload is completed or on failure
if os . path . exists ( file_path ) :
os . remove ( file_path )
2024-05-30 05:03:22 +00:00
class AddPipelineForm ( BaseModel ) :
url : str
urlIdx : int
@app.post ( " /api/pipelines/add " )
async def add_pipeline ( form_data : AddPipelineForm , user = Depends ( get_admin_user ) ) :
r = None
try :
urlIdx = form_data . urlIdx
url = openai_app . state . config . OPENAI_API_BASE_URLS [ urlIdx ]
key = openai_app . state . config . OPENAI_API_KEYS [ urlIdx ]
headers = { " Authorization " : f " Bearer { key } " }
r = requests . post (
f " { url } /pipelines/add " , headers = headers , 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 = " Pipeline not found "
if r is not None :
try :
res = r . json ( )
if " detail " in res :
detail = res [ " detail " ]
2024-08-03 13:24:26 +00:00
except Exception :
2024-05-30 05:03:22 +00:00
pass
raise HTTPException (
status_code = ( r . status_code if r is not None else status . HTTP_404_NOT_FOUND ) ,
detail = detail ,
)
class DeletePipelineForm ( BaseModel ) :
id : str
urlIdx : int
@app.delete ( " /api/pipelines/delete " )
async def delete_pipeline ( form_data : DeletePipelineForm , user = Depends ( get_admin_user ) ) :
r = None
try :
urlIdx = form_data . urlIdx
url = openai_app . state . config . OPENAI_API_BASE_URLS [ urlIdx ]
key = openai_app . state . config . OPENAI_API_KEYS [ urlIdx ]
headers = { " Authorization " : f " Bearer { key } " }
r = requests . delete (
f " { url } /pipelines/delete " , headers = headers , 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 = " Pipeline not found "
if r is not None :
try :
res = r . json ( )
if " detail " in res :
detail = res [ " detail " ]
2024-08-03 13:24:26 +00:00
except Exception :
2024-05-30 05:03:22 +00:00
pass
raise HTTPException (
status_code = ( r . status_code if r is not None else status . HTTP_404_NOT_FOUND ) ,
detail = detail ,
)
2024-05-28 19:04:19 +00:00
@app.get ( " /api/pipelines " )
2024-05-30 04:26:57 +00:00
async def get_pipelines ( urlIdx : Optional [ int ] = None , user = Depends ( get_admin_user ) ) :
2024-05-30 05:18:27 +00:00
r = None
try :
url = openai_app . state . config . OPENAI_API_BASE_URLS [ urlIdx ]
key = openai_app . state . config . OPENAI_API_KEYS [ urlIdx ]
2024-05-30 04:26:57 +00:00
2024-05-30 05:18:27 +00:00
headers = { " Authorization " : f " Bearer { key } " }
r = requests . get ( f " { url } /pipelines " , headers = headers )
2024-05-30 04:26:57 +00:00
2024-05-30 05:18:27 +00:00
r . raise_for_status ( )
data = r . json ( )
2024-05-28 19:04:19 +00:00
2024-05-30 05:18:27 +00:00
return { * * data }
except Exception as e :
# Handle connection error here
print ( f " Connection error: { e } " )
2024-05-28 19:04:19 +00:00
2024-05-30 05:18:27 +00:00
detail = " Pipeline not found "
if r is not None :
try :
res = r . json ( )
if " detail " in res :
detail = res [ " detail " ]
2024-08-03 13:24:26 +00:00
except Exception :
2024-05-30 05:18:27 +00:00
pass
raise HTTPException (
status_code = ( r . status_code if r is not None else status . HTTP_404_NOT_FOUND ) ,
detail = detail ,
)
2024-05-28 19:32:49 +00:00
2024-05-30 05:18:27 +00:00
@app.get ( " /api/pipelines/ {pipeline_id} /valves " )
async def get_pipeline_valves (
2024-06-18 13:03:31 +00:00
urlIdx : Optional [ int ] ,
pipeline_id : str ,
user = Depends ( get_admin_user ) ,
2024-05-30 05:18:27 +00:00
) :
r = None
try :
url = openai_app . state . config . OPENAI_API_BASE_URLS [ urlIdx ]
key = openai_app . state . config . OPENAI_API_KEYS [ urlIdx ]
2024-05-28 19:32:49 +00:00
2024-05-30 05:18:27 +00:00
headers = { " Authorization " : f " Bearer { key } " }
r = requests . get ( f " { url } / { pipeline_id } /valves " , headers = headers )
2024-05-28 19:32:49 +00:00
2024-05-30 05:18:27 +00:00
r . raise_for_status ( )
data = r . json ( )
2024-05-28 19:32:49 +00:00
2024-05-30 05:18:27 +00:00
return { * * data }
except Exception as e :
# Handle connection error here
print ( f " Connection error: { e } " )
2024-05-28 20:05:31 +00:00
2024-05-30 05:18:27 +00:00
detail = " Pipeline not found "
2024-05-28 20:05:31 +00:00
2024-05-30 05:18:27 +00:00
if r is not None :
try :
res = r . json ( )
if " detail " in res :
detail = res [ " detail " ]
2024-08-14 12:38:19 +00:00
except Exception :
2024-05-30 05:18:27 +00:00
pass
2024-05-28 19:32:49 +00:00
raise HTTPException (
2024-05-30 05:18:27 +00:00
status_code = ( r . status_code if r is not None else status . HTTP_404_NOT_FOUND ) ,
detail = detail ,
2024-05-28 19:32:49 +00:00
)
@app.get ( " /api/pipelines/ {pipeline_id} /valves/spec " )
2024-05-30 05:18:27 +00:00
async def get_pipeline_valves_spec (
2024-06-18 13:03:31 +00:00
urlIdx : Optional [ int ] ,
pipeline_id : str ,
user = Depends ( get_admin_user ) ,
2024-05-30 05:18:27 +00:00
) :
r = None
try :
url = openai_app . state . config . OPENAI_API_BASE_URLS [ urlIdx ]
key = openai_app . state . config . OPENAI_API_KEYS [ urlIdx ]
2024-05-28 19:32:49 +00:00
2024-05-30 05:18:27 +00:00
headers = { " Authorization " : f " Bearer { key } " }
r = requests . get ( f " { url } / { pipeline_id } /valves/spec " , headers = headers )
2024-05-28 19:32:49 +00:00
2024-05-30 05:18:27 +00:00
r . raise_for_status ( )
data = r . json ( )
2024-05-28 19:32:49 +00:00
2024-05-30 05:18:27 +00:00
return { * * data }
except Exception as e :
# Handle connection error here
print ( f " Connection error: { e } " )
2024-05-28 19:32:49 +00:00
2024-05-30 05:18:27 +00:00
detail = " Pipeline not found "
if r is not None :
try :
res = r . json ( )
if " detail " in res :
detail = res [ " detail " ]
2024-08-03 13:24:26 +00:00
except Exception :
2024-05-30 05:18:27 +00:00
pass
2024-05-28 20:05:31 +00:00
2024-05-28 19:32:49 +00:00
raise HTTPException (
2024-05-30 05:18:27 +00:00
status_code = ( r . status_code if r is not None else status . HTTP_404_NOT_FOUND ) ,
detail = detail ,
2024-05-28 19:32:49 +00:00
)
@app.post ( " /api/pipelines/ {pipeline_id} /valves/update " )
async def update_pipeline_valves (
2024-05-30 05:18:27 +00:00
urlIdx : Optional [ int ] ,
pipeline_id : str ,
form_data : dict ,
user = Depends ( get_admin_user ) ,
2024-05-28 19:32:49 +00:00
) :
2024-05-30 05:18:27 +00:00
r = None
try :
url = openai_app . state . config . OPENAI_API_BASE_URLS [ urlIdx ]
key = openai_app . state . config . OPENAI_API_KEYS [ urlIdx ]
2024-05-28 19:32:49 +00:00
2024-05-30 05:18:27 +00:00
headers = { " Authorization " : f " Bearer { key } " }
r = requests . post (
f " { url } / { pipeline_id } /valves/update " ,
headers = headers ,
json = { * * form_data } ,
)
2024-05-28 19:32:49 +00:00
2024-05-30 05:18:27 +00:00
r . raise_for_status ( )
data = r . json ( )
2024-05-28 19:32:49 +00:00
2024-05-30 05:18:27 +00:00
return { * * data }
except Exception as e :
# Handle connection error here
print ( f " Connection error: { e } " )
2024-05-28 19:32:49 +00:00
2024-05-30 05:18:27 +00:00
detail = " Pipeline not found "
2024-05-28 20:05:31 +00:00
2024-05-30 05:18:27 +00:00
if r is not None :
try :
res = r . json ( )
if " detail " in res :
detail = res [ " detail " ]
2024-08-03 13:24:26 +00:00
except Exception :
2024-05-30 05:18:27 +00:00
pass
2024-05-28 20:05:31 +00:00
2024-05-28 19:32:49 +00:00
raise HTTPException (
2024-05-30 05:41:51 +00:00
status_code = ( r . status_code if r is not None else status . HTTP_404_NOT_FOUND ) ,
2024-05-30 05:18:27 +00:00
detail = detail ,
2024-05-28 19:32:49 +00:00
)
2024-06-20 08:51:39 +00:00
##################################
#
# Config Endpoints
#
##################################
2024-02-22 02:12:01 +00:00
@app.get ( " /api/config " )
async def get_app_config ( ) :
return {
" status " : True ,
2024-02-24 01:12:19 +00:00
" name " : WEBUI_NAME ,
2024-02-23 08:30:26 +00:00
" version " : VERSION ,
2024-06-30 21:48:05 +00:00
" default_locale " : str ( DEFAULT_LOCALE ) ,
2024-05-10 07:03:24 +00:00
" default_models " : webui_app . state . config . DEFAULT_MODELS ,
" default_prompt_suggestions " : webui_app . state . config . DEFAULT_PROMPT_SUGGESTIONS ,
2024-05-26 20:02:40 +00:00
" features " : {
2024-05-26 16:05:26 +00:00
" auth " : WEBUI_AUTH ,
" auth_trusted_header " : bool ( webui_app . state . AUTH_TRUSTED_EMAIL_HEADER ) ,
2024-05-26 20:02:40 +00:00
" enable_signup " : webui_app . state . config . ENABLE_SIGNUP ,
2024-07-25 01:44:40 +00:00
" enable_login_form " : webui_app . state . config . ENABLE_LOGIN_FORM ,
2024-06-02 02:03:56 +00:00
" enable_web_search " : rag_app . state . config . ENABLE_RAG_WEB_SEARCH ,
2024-05-26 16:05:26 +00:00
" enable_image_generation " : images_app . state . config . ENABLED ,
2024-05-26 16:10:25 +00:00
" enable_community_sharing " : webui_app . state . config . ENABLE_COMMUNITY_SHARING ,
2024-05-27 06:40:05 +00:00
" enable_admin_export " : ENABLE_ADMIN_EXPORT ,
2024-08-04 14:16:14 +00:00
" enable_admin_chat_access " : ENABLE_ADMIN_CHAT_ACCESS ,
2024-05-26 16:05:26 +00:00
} ,
2024-06-08 03:18:48 +00:00
" audio " : {
" tts " : {
" engine " : audio_app . state . config . TTS_ENGINE ,
" voice " : audio_app . state . config . TTS_VOICE ,
} ,
" stt " : {
" engine " : audio_app . state . config . STT_ENGINE ,
} ,
} ,
2024-05-26 07:37:09 +00:00
" oauth " : {
" providers " : {
name : config . get ( " name " , name )
for name , config in OAUTH_PROVIDERS . items ( )
}
} ,
2024-02-22 02:12:01 +00:00
}
2024-03-10 05:19:20 +00:00
@app.get ( " /api/config/model/filter " )
async def get_model_filter_config ( user = Depends ( get_admin_user ) ) :
2024-03-10 05:47:01 +00:00
return {
2024-05-10 07:03:24 +00:00
" enabled " : app . state . config . ENABLE_MODEL_FILTER ,
" models " : app . state . config . MODEL_FILTER_LIST ,
2024-03-10 05:47:01 +00:00
}
2024-03-10 05:19:20 +00:00
class ModelFilterConfigForm ( BaseModel ) :
enabled : bool
2024-08-14 12:46:31 +00:00
models : list [ str ]
2024-03-10 05:19:20 +00:00
@app.post ( " /api/config/model/filter " )
2024-03-21 01:35:02 +00:00
async def update_model_filter_config (
2024-03-10 05:19:20 +00:00
form_data : ModelFilterConfigForm , user = Depends ( get_admin_user )
) :
2024-05-17 17:35:33 +00:00
app . state . config . ENABLE_MODEL_FILTER = form_data . enabled
app . state . config . MODEL_FILTER_LIST = form_data . models
2024-03-10 05:19:20 +00:00
2024-03-10 05:47:01 +00:00
return {
2024-05-10 07:03:24 +00:00
" enabled " : app . state . config . ENABLE_MODEL_FILTER ,
" models " : app . state . config . MODEL_FILTER_LIST ,
2024-03-10 05:47:01 +00:00
}
2024-03-10 05:19:20 +00:00
2024-06-20 08:51:39 +00:00
# TODO: webhook endpoint should be under config endpoints
2024-03-21 01:35:02 +00:00
@app.get ( " /api/webhook " )
async def get_webhook_url ( user = Depends ( get_admin_user ) ) :
return {
2024-05-10 07:03:24 +00:00
" url " : app . state . config . WEBHOOK_URL ,
2024-03-21 01:35:02 +00:00
}
class UrlForm ( BaseModel ) :
url : str
@app.post ( " /api/webhook " )
async def update_webhook_url ( form_data : UrlForm , user = Depends ( get_admin_user ) ) :
2024-05-10 07:03:24 +00:00
app . state . config . WEBHOOK_URL = form_data . url
webui_app . state . WEBHOOK_URL = app . state . config . WEBHOOK_URL
2024-06-04 04:17:43 +00:00
return { " url " : app . state . config . WEBHOOK_URL }
2024-05-26 16:23:24 +00:00
2024-03-05 08:59:35 +00:00
@app.get ( " /api/version " )
2024-08-03 13:24:26 +00:00
async def get_app_version ( ) :
2024-03-05 08:59:35 +00:00
return {
" version " : VERSION ,
}
2024-02-23 08:30:26 +00:00
@app.get ( " /api/changelog " )
async def get_app_changelog ( ) :
2024-03-31 08:10:57 +00:00
return { key : CHANGELOG [ key ] for idx , key in enumerate ( CHANGELOG ) if idx < 5 }
2024-02-23 08:30:26 +00:00
2024-02-25 19:26:58 +00:00
@app.get ( " /api/version/updates " )
async def get_app_latest_release_version ( ) :
try :
2024-06-06 00:24:59 +00:00
async with aiohttp . ClientSession ( trust_env = True ) as session :
2024-04-10 06:03:05 +00:00
async with session . get (
" https://api.github.com/repos/open-webui/open-webui/releases/latest "
) as response :
response . raise_for_status ( )
data = await response . json ( )
latest_version = data [ " tag_name " ]
return { " current " : VERSION , " latest " : latest_version [ 1 : ] }
2024-08-03 13:24:26 +00:00
except aiohttp . ClientError :
2024-02-25 19:26:58 +00:00
raise HTTPException (
status_code = status . HTTP_503_SERVICE_UNAVAILABLE ,
2024-02-25 19:55:15 +00:00
detail = ERROR_MESSAGES . RATE_LIMIT_EXCEEDED ,
2024-02-25 19:26:58 +00:00
)
2024-04-10 08:27:19 +00:00
2024-05-27 17:07:38 +00:00
############################
# OAuth Login & Callback
############################
oauth = OAuth ( )
for provider_name , provider_config in OAUTH_PROVIDERS . items ( ) :
oauth . register (
name = provider_name ,
client_id = provider_config [ " client_id " ] ,
client_secret = provider_config [ " client_secret " ] ,
server_metadata_url = provider_config [ " server_metadata_url " ] ,
client_kwargs = {
" scope " : provider_config [ " scope " ] ,
} ,
2024-07-19 07:03:41 +00:00
redirect_uri = provider_config [ " redirect_uri " ] ,
2024-05-27 17:07:38 +00:00
)
# SessionMiddleware is used by authlib for oauth
if len ( OAUTH_PROVIDERS ) > 0 :
app . add_middleware (
2024-06-05 18:21:42 +00:00
SessionMiddleware ,
secret_key = WEBUI_SECRET_KEY ,
session_cookie = " oui-session " ,
same_site = WEBUI_SESSION_COOKIE_SAME_SITE ,
2024-06-07 08:13:42 +00:00
https_only = WEBUI_SESSION_COOKIE_SECURE ,
2024-05-27 17:07:38 +00:00
)
@app.get ( " /oauth/ {provider} /login " )
async def oauth_login ( provider : str , request : Request ) :
if provider not in OAUTH_PROVIDERS :
raise HTTPException ( 404 )
2024-07-19 07:03:41 +00:00
# If the provider has a custom redirect URL, use that, otherwise automatically generate one
2024-08-01 11:54:06 +00:00
redirect_uri = OAUTH_PROVIDERS [ provider ] . get ( " redirect_uri " ) or request . url_for (
2024-07-19 07:03:41 +00:00
" oauth_callback " , provider = provider
)
2024-08-11 07:31:40 +00:00
client = oauth . create_client ( provider )
if client is None :
raise HTTPException ( 404 )
return await client . authorize_redirect ( request , redirect_uri )
2024-05-27 17:07:38 +00:00
2024-06-21 17:25:19 +00:00
# OAuth login logic is as follows:
# 1. Attempt to find a user with matching subject ID, tied to the provider
# 2. If OAUTH_MERGE_ACCOUNTS_BY_EMAIL is true, find a user with the email address provided via OAuth
# - This is considered insecure in general, as OAuth providers do not always verify email addresses
# 3. If there is no user, and ENABLE_OAUTH_SIGNUP is true, create a user
# - Email addresses are considered unique, so we fail registration if the email address is alreayd taken
2024-05-27 17:07:38 +00:00
@app.get ( " /oauth/ {provider} /callback " )
2024-06-21 13:35:11 +00:00
async def oauth_callback ( provider : str , request : Request , response : Response ) :
2024-05-27 17:07:38 +00:00
if provider not in OAUTH_PROVIDERS :
raise HTTPException ( 404 )
client = oauth . create_client ( provider )
2024-06-05 18:21:42 +00:00
try :
token = await client . authorize_access_token ( request )
except Exception as e :
2024-06-24 02:43:53 +00:00
log . warning ( f " OAuth callback error: { e } " )
2024-06-05 18:21:42 +00:00
raise HTTPException ( 400 , detail = ERROR_MESSAGES . INVALID_CRED )
2024-05-27 17:07:38 +00:00
user_data : UserInfo = token [ " userinfo " ]
sub = user_data . get ( " sub " )
if not sub :
2024-06-24 02:43:53 +00:00
log . warning ( f " OAuth callback failed, sub is missing: { user_data } " )
2024-05-27 17:07:38 +00:00
raise HTTPException ( 400 , detail = ERROR_MESSAGES . INVALID_CRED )
provider_sub = f " { provider } @ { sub } "
2024-08-07 18:39:51 +00:00
email_claim = webui_app . state . config . OAUTH_EMAIL_CLAIM
email = user_data . get ( email_claim , " " ) . lower ( )
2024-06-21 17:25:19 +00:00
# We currently mandate that email addresses are provided
if not email :
2024-06-24 02:43:53 +00:00
log . warning ( f " OAuth callback failed, email is missing: { user_data } " )
2024-06-21 17:25:19 +00:00
raise HTTPException ( 400 , detail = ERROR_MESSAGES . INVALID_CRED )
2024-05-27 17:07:38 +00:00
# Check if the user exists
user = Users . get_user_by_oauth_sub ( provider_sub )
if not user :
# If the user does not exist, check if merging is enabled
2024-05-27 17:16:36 +00:00
if OAUTH_MERGE_ACCOUNTS_BY_EMAIL . value :
2024-05-27 17:07:38 +00:00
# Check if the user exists by email
2024-06-21 17:25:19 +00:00
user = Users . get_user_by_email ( email )
2024-05-27 17:07:38 +00:00
if user :
# Update the user with the new oauth sub
Users . update_user_oauth_sub_by_id ( user . id , provider_sub )
if not user :
# If the user does not exist, check if signups are enabled
if ENABLE_OAUTH_SIGNUP . value :
2024-06-21 17:25:19 +00:00
# Check if an existing user with the same email already exists
existing_user = Users . get_user_by_email ( user_data . get ( " email " , " " ) . lower ( ) )
if existing_user :
raise HTTPException ( 400 , detail = ERROR_MESSAGES . EMAIL_TAKEN )
2024-07-01 07:25:25 +00:00
picture_claim = webui_app . state . config . OAUTH_PICTURE_CLAIM
picture_url = user_data . get ( picture_claim , " " )
2024-06-21 13:35:54 +00:00
if picture_url :
# Download the profile image into a base64 string
try :
async with aiohttp . ClientSession ( ) as session :
async with session . get ( picture_url ) as resp :
picture = await resp . read ( )
base64_encoded_picture = base64 . b64encode ( picture ) . decode (
" utf-8 "
)
guessed_mime_type = mimetypes . guess_type ( picture_url ) [ 0 ]
if guessed_mime_type is None :
# assume JPG, browsers are tolerant enough of image formats
guessed_mime_type = " image/jpeg "
picture_url = f " data: { guessed_mime_type } ;base64, { base64_encoded_picture } "
except Exception as e :
2024-06-21 17:25:19 +00:00
log . error ( f " Error downloading profile image ' { picture_url } ' : { e } " )
2024-06-21 13:35:54 +00:00
picture_url = " "
if not picture_url :
picture_url = " /user.png "
2024-06-28 14:08:32 +00:00
username_claim = webui_app . state . config . OAUTH_USERNAME_CLAIM
2024-06-28 13:20:34 +00:00
role = (
" admin "
if Users . get_num_users ( ) == 0
else webui_app . state . config . DEFAULT_USER_ROLE
)
2024-05-27 17:07:38 +00:00
user = Auths . insert_new_auth (
2024-06-21 17:25:19 +00:00
email = email ,
2024-05-27 17:07:38 +00:00
password = get_password_hash (
str ( uuid . uuid4 ( ) )
) , # Random password, not used
2024-06-28 14:08:32 +00:00
name = user_data . get ( username_claim , " User " ) ,
2024-06-21 13:35:54 +00:00
profile_image_url = picture_url ,
2024-06-28 13:20:34 +00:00
role = role ,
2024-05-27 17:07:38 +00:00
oauth_sub = provider_sub ,
)
if webui_app . state . config . WEBHOOK_URL :
post_webhook (
webui_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 ) ,
} ,
)
else :
2024-06-24 02:43:53 +00:00
raise HTTPException (
status . HTTP_403_FORBIDDEN , detail = ERROR_MESSAGES . ACCESS_PROHIBITED
)
2024-05-27 17:07:38 +00:00
jwt_token = create_token (
data = { " id " : user . id } ,
expires_delta = parse_duration ( webui_app . state . config . JWT_EXPIRES_IN ) ,
)
2024-06-21 13:35:11 +00:00
# Set the cookie token
response . set_cookie (
key = " token " ,
2024-06-28 13:20:57 +00:00
value = jwt_token ,
2024-06-21 13:35:11 +00:00
httponly = True , # Ensures the cookie is not accessible via JavaScript
)
2024-05-27 17:07:38 +00:00
# Redirect back to the frontend with the JWT token
redirect_url = f " { request . base_url } auth#token= { jwt_token } "
return RedirectResponse ( url = redirect_url )
2024-04-02 18:55:00 +00:00
@app.get ( " /manifest.json " )
async def get_manifest_json ( ) :
return {
2024-04-04 03:43:55 +00:00
" name " : WEBUI_NAME ,
" short_name " : WEBUI_NAME ,
2024-04-02 18:55:00 +00:00
" start_url " : " / " ,
" display " : " standalone " ,
" background_color " : " #343541 " ,
" orientation " : " portrait-primary " ,
2024-05-02 02:32:36 +00:00
" icons " : [ { " src " : " /static/logo.png " , " type " : " image/png " , " sizes " : " 500x500 " } ] ,
2024-04-02 18:55:00 +00:00
}
2024-04-10 08:27:19 +00:00
2024-05-07 00:29:16 +00:00
@app.get ( " /opensearch.xml " )
async def get_opensearch_xml ( ) :
xml_content = rf """
< OpenSearchDescription xmlns = " http://a9.com/-/spec/opensearch/1.1/ " xmlns : moz = " http://www.mozilla.org/2006/browser/search/ " >
< ShortName > { WEBUI_NAME } < / ShortName >
< Description > Search { WEBUI_NAME } < / Description >
< InputEncoding > UTF - 8 < / InputEncoding >
2024-07-09 06:07:23 +00:00
< Image width = " 16 " height = " 16 " type = " image/x-icon " > { WEBUI_URL } / static / favicon . png < / Image >
2024-05-07 00:29:16 +00:00
< Url type = " text/html " method = " get " template = " {WEBUI_URL} /?q= { " { searchTerms } " } " / >
< moz : SearchForm > { WEBUI_URL } < / moz : SearchForm >
< / OpenSearchDescription >
"""
return Response ( content = xml_content , media_type = " application/xml " )
2024-05-15 18:17:18 +00:00
@app.get ( " /health " )
async def healthcheck ( ) :
return { " status " : True }
2024-06-18 13:03:31 +00:00
@app.get ( " /health/db " )
2024-06-21 12:58:57 +00:00
async def healthcheck_with_db ( ) :
2024-06-24 11:06:15 +00:00
Session . execute ( text ( " SELECT 1; " ) ) . all ( )
2024-06-18 13:03:31 +00:00
return { " status " : True }
2024-04-09 10:32:28 +00:00
app . mount ( " /static " , StaticFiles ( directory = STATIC_DIR ) , name = " static " )
app . mount ( " /cache " , StaticFiles ( directory = CACHE_DIR ) , name = " cache " )
2024-02-24 01:12:19 +00:00
2024-04-28 15:03:30 +00:00
if os . path . exists ( FRONTEND_BUILD_DIR ) :
2024-05-22 04:38:58 +00:00
mimetypes . add_type ( " text/javascript " , " .js " )
2024-04-28 15:03:30 +00:00
app . mount (
" / " ,
SPAStaticFiles ( directory = FRONTEND_BUILD_DIR , html = True ) ,
name = " spa-static-files " ,
)
else :
log . warning (
f " Frontend build directory not found at ' { FRONTEND_BUILD_DIR } ' . Serving API only. "
)