Merge remote-tracking branch 'upstream/dev' into feat/oauth
This commit is contained in:
Jun Siang Cheah
+207
-59
@@ -17,13 +17,12 @@ from fastapi.middleware.cors import CORSMiddleware
|
||||
from faster_whisper import WhisperModel
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
import uuid
|
||||
import requests
|
||||
import hashlib
|
||||
from pathlib import Path
|
||||
import json
|
||||
|
||||
|
||||
from constants import ERROR_MESSAGES
|
||||
from utils.utils import (
|
||||
decode_token,
|
||||
@@ -41,10 +40,15 @@ from config import (
|
||||
WHISPER_MODEL_DIR,
|
||||
WHISPER_MODEL_AUTO_UPDATE,
|
||||
DEVICE_TYPE,
|
||||
AUDIO_OPENAI_API_BASE_URL,
|
||||
AUDIO_OPENAI_API_KEY,
|
||||
AUDIO_OPENAI_API_MODEL,
|
||||
AUDIO_OPENAI_API_VOICE,
|
||||
AUDIO_STT_OPENAI_API_BASE_URL,
|
||||
AUDIO_STT_OPENAI_API_KEY,
|
||||
AUDIO_TTS_OPENAI_API_BASE_URL,
|
||||
AUDIO_TTS_OPENAI_API_KEY,
|
||||
AUDIO_STT_ENGINE,
|
||||
AUDIO_STT_MODEL,
|
||||
AUDIO_TTS_ENGINE,
|
||||
AUDIO_TTS_MODEL,
|
||||
AUDIO_TTS_VOICE,
|
||||
AppConfig,
|
||||
)
|
||||
|
||||
@@ -61,10 +65,17 @@ app.add_middleware(
|
||||
)
|
||||
|
||||
app.state.config = AppConfig()
|
||||
app.state.config.OPENAI_API_BASE_URL = AUDIO_OPENAI_API_BASE_URL
|
||||
app.state.config.OPENAI_API_KEY = AUDIO_OPENAI_API_KEY
|
||||
app.state.config.OPENAI_API_MODEL = AUDIO_OPENAI_API_MODEL
|
||||
app.state.config.OPENAI_API_VOICE = AUDIO_OPENAI_API_VOICE
|
||||
|
||||
app.state.config.STT_OPENAI_API_BASE_URL = AUDIO_STT_OPENAI_API_BASE_URL
|
||||
app.state.config.STT_OPENAI_API_KEY = AUDIO_STT_OPENAI_API_KEY
|
||||
app.state.config.STT_ENGINE = AUDIO_STT_ENGINE
|
||||
app.state.config.STT_MODEL = AUDIO_STT_MODEL
|
||||
|
||||
app.state.config.TTS_OPENAI_API_BASE_URL = AUDIO_TTS_OPENAI_API_BASE_URL
|
||||
app.state.config.TTS_OPENAI_API_KEY = AUDIO_TTS_OPENAI_API_KEY
|
||||
app.state.config.TTS_ENGINE = AUDIO_TTS_ENGINE
|
||||
app.state.config.TTS_MODEL = AUDIO_TTS_MODEL
|
||||
app.state.config.TTS_VOICE = AUDIO_TTS_VOICE
|
||||
|
||||
# setting device type for whisper model
|
||||
whisper_device_type = DEVICE_TYPE if DEVICE_TYPE and DEVICE_TYPE == "cuda" else "cpu"
|
||||
@@ -74,41 +85,101 @@ SPEECH_CACHE_DIR = Path(CACHE_DIR).joinpath("./audio/speech/")
|
||||
SPEECH_CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
|
||||
class OpenAIConfigUpdateForm(BaseModel):
|
||||
url: str
|
||||
key: str
|
||||
model: str
|
||||
speaker: str
|
||||
class TTSConfigForm(BaseModel):
|
||||
OPENAI_API_BASE_URL: str
|
||||
OPENAI_API_KEY: str
|
||||
ENGINE: str
|
||||
MODEL: str
|
||||
VOICE: str
|
||||
|
||||
|
||||
class STTConfigForm(BaseModel):
|
||||
OPENAI_API_BASE_URL: str
|
||||
OPENAI_API_KEY: str
|
||||
ENGINE: str
|
||||
MODEL: str
|
||||
|
||||
|
||||
class AudioConfigUpdateForm(BaseModel):
|
||||
tts: TTSConfigForm
|
||||
stt: STTConfigForm
|
||||
|
||||
|
||||
from pydub import AudioSegment
|
||||
from pydub.utils import mediainfo
|
||||
|
||||
|
||||
def is_mp4_audio(file_path):
|
||||
"""Check if the given file is an MP4 audio file."""
|
||||
if not os.path.isfile(file_path):
|
||||
print(f"File not found: {file_path}")
|
||||
return False
|
||||
|
||||
info = mediainfo(file_path)
|
||||
if (
|
||||
info.get("codec_name") == "aac"
|
||||
and info.get("codec_type") == "audio"
|
||||
and info.get("codec_tag_string") == "mp4a"
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def convert_mp4_to_wav(file_path, output_path):
|
||||
"""Convert MP4 audio file to WAV format."""
|
||||
audio = AudioSegment.from_file(file_path, format="mp4")
|
||||
audio.export(output_path, format="wav")
|
||||
print(f"Converted {file_path} to {output_path}")
|
||||
|
||||
|
||||
@app.get("/config")
|
||||
async def get_openai_config(user=Depends(get_admin_user)):
|
||||
async def get_audio_config(user=Depends(get_admin_user)):
|
||||
return {
|
||||
"OPENAI_API_BASE_URL": app.state.config.OPENAI_API_BASE_URL,
|
||||
"OPENAI_API_KEY": app.state.config.OPENAI_API_KEY,
|
||||
"OPENAI_API_MODEL": app.state.config.OPENAI_API_MODEL,
|
||||
"OPENAI_API_VOICE": app.state.config.OPENAI_API_VOICE,
|
||||
"tts": {
|
||||
"OPENAI_API_BASE_URL": app.state.config.TTS_OPENAI_API_BASE_URL,
|
||||
"OPENAI_API_KEY": app.state.config.TTS_OPENAI_API_KEY,
|
||||
"ENGINE": app.state.config.TTS_ENGINE,
|
||||
"MODEL": app.state.config.TTS_MODEL,
|
||||
"VOICE": app.state.config.TTS_VOICE,
|
||||
},
|
||||
"stt": {
|
||||
"OPENAI_API_BASE_URL": app.state.config.STT_OPENAI_API_BASE_URL,
|
||||
"OPENAI_API_KEY": app.state.config.STT_OPENAI_API_KEY,
|
||||
"ENGINE": app.state.config.STT_ENGINE,
|
||||
"MODEL": app.state.config.STT_MODEL,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@app.post("/config/update")
|
||||
async def update_openai_config(
|
||||
form_data: OpenAIConfigUpdateForm, user=Depends(get_admin_user)
|
||||
async def update_audio_config(
|
||||
form_data: AudioConfigUpdateForm, user=Depends(get_admin_user)
|
||||
):
|
||||
if form_data.key == "":
|
||||
raise HTTPException(status_code=400, detail=ERROR_MESSAGES.API_KEY_NOT_FOUND)
|
||||
app.state.config.TTS_OPENAI_API_BASE_URL = form_data.tts.OPENAI_API_BASE_URL
|
||||
app.state.config.TTS_OPENAI_API_KEY = form_data.tts.OPENAI_API_KEY
|
||||
app.state.config.TTS_ENGINE = form_data.tts.ENGINE
|
||||
app.state.config.TTS_MODEL = form_data.tts.MODEL
|
||||
app.state.config.TTS_VOICE = form_data.tts.VOICE
|
||||
|
||||
app.state.config.OPENAI_API_BASE_URL = form_data.url
|
||||
app.state.config.OPENAI_API_KEY = form_data.key
|
||||
app.state.config.OPENAI_API_MODEL = form_data.model
|
||||
app.state.config.OPENAI_API_VOICE = form_data.speaker
|
||||
app.state.config.STT_OPENAI_API_BASE_URL = form_data.stt.OPENAI_API_BASE_URL
|
||||
app.state.config.STT_OPENAI_API_KEY = form_data.stt.OPENAI_API_KEY
|
||||
app.state.config.STT_ENGINE = form_data.stt.ENGINE
|
||||
app.state.config.STT_MODEL = form_data.stt.MODEL
|
||||
|
||||
return {
|
||||
"status": True,
|
||||
"OPENAI_API_BASE_URL": app.state.config.OPENAI_API_BASE_URL,
|
||||
"OPENAI_API_KEY": app.state.config.OPENAI_API_KEY,
|
||||
"OPENAI_API_MODEL": app.state.config.OPENAI_API_MODEL,
|
||||
"OPENAI_API_VOICE": app.state.config.OPENAI_API_VOICE,
|
||||
"tts": {
|
||||
"OPENAI_API_BASE_URL": app.state.config.TTS_OPENAI_API_BASE_URL,
|
||||
"OPENAI_API_KEY": app.state.config.TTS_OPENAI_API_KEY,
|
||||
"ENGINE": app.state.config.TTS_ENGINE,
|
||||
"MODEL": app.state.config.TTS_MODEL,
|
||||
"VOICE": app.state.config.TTS_VOICE,
|
||||
},
|
||||
"stt": {
|
||||
"OPENAI_API_BASE_URL": app.state.config.STT_OPENAI_API_BASE_URL,
|
||||
"OPENAI_API_KEY": app.state.config.STT_OPENAI_API_KEY,
|
||||
"ENGINE": app.state.config.STT_ENGINE,
|
||||
"MODEL": app.state.config.STT_MODEL,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@@ -125,13 +196,21 @@ async def speech(request: Request, user=Depends(get_verified_user)):
|
||||
return FileResponse(file_path)
|
||||
|
||||
headers = {}
|
||||
headers["Authorization"] = f"Bearer {app.state.config.OPENAI_API_KEY}"
|
||||
headers["Authorization"] = f"Bearer {app.state.config.TTS_OPENAI_API_KEY}"
|
||||
headers["Content-Type"] = "application/json"
|
||||
|
||||
try:
|
||||
body = body.decode("utf-8")
|
||||
body = json.loads(body)
|
||||
body["model"] = app.state.config.TTS_MODEL
|
||||
body = json.dumps(body).encode("utf-8")
|
||||
except Exception as e:
|
||||
pass
|
||||
|
||||
r = None
|
||||
try:
|
||||
r = requests.post(
|
||||
url=f"{app.state.config.OPENAI_API_BASE_URL}/audio/speech",
|
||||
url=f"{app.state.config.TTS_OPENAI_API_BASE_URL}/audio/speech",
|
||||
data=body,
|
||||
headers=headers,
|
||||
stream=True,
|
||||
@@ -181,41 +260,110 @@ def transcribe(
|
||||
)
|
||||
|
||||
try:
|
||||
filename = file.filename
|
||||
file_path = f"{UPLOAD_DIR}/{filename}"
|
||||
ext = file.filename.split(".")[-1]
|
||||
|
||||
id = uuid.uuid4()
|
||||
filename = f"{id}.{ext}"
|
||||
|
||||
file_dir = f"{CACHE_DIR}/audio/transcriptions"
|
||||
os.makedirs(file_dir, exist_ok=True)
|
||||
file_path = f"{file_dir}/{filename}"
|
||||
|
||||
print(filename)
|
||||
|
||||
contents = file.file.read()
|
||||
with open(file_path, "wb") as f:
|
||||
f.write(contents)
|
||||
f.close()
|
||||
|
||||
whisper_kwargs = {
|
||||
"model_size_or_path": WHISPER_MODEL,
|
||||
"device": whisper_device_type,
|
||||
"compute_type": "int8",
|
||||
"download_root": WHISPER_MODEL_DIR,
|
||||
"local_files_only": not WHISPER_MODEL_AUTO_UPDATE,
|
||||
}
|
||||
if app.state.config.STT_ENGINE == "":
|
||||
whisper_kwargs = {
|
||||
"model_size_or_path": WHISPER_MODEL,
|
||||
"device": whisper_device_type,
|
||||
"compute_type": "int8",
|
||||
"download_root": WHISPER_MODEL_DIR,
|
||||
"local_files_only": not WHISPER_MODEL_AUTO_UPDATE,
|
||||
}
|
||||
|
||||
log.debug(f"whisper_kwargs: {whisper_kwargs}")
|
||||
log.debug(f"whisper_kwargs: {whisper_kwargs}")
|
||||
|
||||
try:
|
||||
model = WhisperModel(**whisper_kwargs)
|
||||
except:
|
||||
log.warning(
|
||||
"WhisperModel initialization failed, attempting download with local_files_only=False"
|
||||
try:
|
||||
model = WhisperModel(**whisper_kwargs)
|
||||
except:
|
||||
log.warning(
|
||||
"WhisperModel initialization failed, attempting download with local_files_only=False"
|
||||
)
|
||||
whisper_kwargs["local_files_only"] = False
|
||||
model = WhisperModel(**whisper_kwargs)
|
||||
|
||||
segments, info = model.transcribe(file_path, beam_size=5)
|
||||
log.info(
|
||||
"Detected language '%s' with probability %f"
|
||||
% (info.language, info.language_probability)
|
||||
)
|
||||
whisper_kwargs["local_files_only"] = False
|
||||
model = WhisperModel(**whisper_kwargs)
|
||||
|
||||
segments, info = model.transcribe(file_path, beam_size=5)
|
||||
log.info(
|
||||
"Detected language '%s' with probability %f"
|
||||
% (info.language, info.language_probability)
|
||||
)
|
||||
transcript = "".join([segment.text for segment in list(segments)])
|
||||
|
||||
transcript = "".join([segment.text for segment in list(segments)])
|
||||
data = {"text": transcript.strip()}
|
||||
|
||||
return {"text": transcript.strip()}
|
||||
# save the transcript to a json file
|
||||
transcript_file = f"{file_dir}/{id}.json"
|
||||
with open(transcript_file, "w") as f:
|
||||
json.dump(data, f)
|
||||
|
||||
print(data)
|
||||
|
||||
return data
|
||||
|
||||
elif app.state.config.STT_ENGINE == "openai":
|
||||
if is_mp4_audio(file_path):
|
||||
print("is_mp4_audio")
|
||||
os.rename(file_path, file_path.replace(".wav", ".mp4"))
|
||||
# Convert MP4 audio file to WAV format
|
||||
convert_mp4_to_wav(file_path.replace(".wav", ".mp4"), file_path)
|
||||
|
||||
headers = {"Authorization": f"Bearer {app.state.config.STT_OPENAI_API_KEY}"}
|
||||
|
||||
files = {"file": (filename, open(file_path, "rb"))}
|
||||
data = {"model": "whisper-1"}
|
||||
|
||||
print(files, data)
|
||||
|
||||
r = None
|
||||
try:
|
||||
r = requests.post(
|
||||
url=f"{app.state.config.STT_OPENAI_API_BASE_URL}/audio/transcriptions",
|
||||
headers=headers,
|
||||
files=files,
|
||||
data=data,
|
||||
)
|
||||
|
||||
r.raise_for_status()
|
||||
|
||||
data = r.json()
|
||||
|
||||
# save the transcript to a json file
|
||||
transcript_file = f"{file_dir}/{id}.json"
|
||||
with open(transcript_file, "w") as f:
|
||||
json.dump(data, f)
|
||||
|
||||
print(data)
|
||||
return data
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
error_detail = "Open WebUI: Server Connection Error"
|
||||
if r is not None:
|
||||
try:
|
||||
res = r.json()
|
||||
if "error" in res:
|
||||
error_detail = f"External: {res['error']['message']}"
|
||||
except:
|
||||
error_detail = f"External: {e}"
|
||||
|
||||
raise HTTPException(
|
||||
status_code=r.status_code if r != None else 500,
|
||||
detail=error_detail,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
|
||||
@@ -41,13 +41,12 @@ from utils.utils import (
|
||||
get_admin_user,
|
||||
)
|
||||
|
||||
from utils.models import get_model_id_from_custom_model_id
|
||||
|
||||
|
||||
from config import (
|
||||
SRC_LOG_LEVELS,
|
||||
OLLAMA_BASE_URLS,
|
||||
ENABLE_OLLAMA_API,
|
||||
AIOHTTP_CLIENT_TIMEOUT,
|
||||
ENABLE_MODEL_FILTER,
|
||||
MODEL_FILTER_LIST,
|
||||
UPLOAD_DIR,
|
||||
@@ -156,7 +155,9 @@ async def cleanup_response(
|
||||
async def post_streaming_url(url: str, payload: str):
|
||||
r = None
|
||||
try:
|
||||
session = aiohttp.ClientSession(trust_env=True)
|
||||
session = aiohttp.ClientSession(
|
||||
trust_env=True, timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT)
|
||||
)
|
||||
r = await session.post(url, data=payload)
|
||||
r.raise_for_status()
|
||||
|
||||
@@ -728,7 +729,6 @@ async def generate_chat_completion(
|
||||
model_info = Models.get_model_by_id(model_id)
|
||||
|
||||
if model_info:
|
||||
print(model_info)
|
||||
if model_info.base_model_id:
|
||||
payload["model"] = model_info.base_model_id
|
||||
|
||||
@@ -754,6 +754,14 @@ async def generate_chat_completion(
|
||||
if model_info.params.get("num_ctx", None):
|
||||
payload["options"]["num_ctx"] = model_info.params.get("num_ctx", None)
|
||||
|
||||
if model_info.params.get("num_batch", None):
|
||||
payload["options"]["num_batch"] = model_info.params.get(
|
||||
"num_batch", None
|
||||
)
|
||||
|
||||
if model_info.params.get("num_keep", None):
|
||||
payload["options"]["num_keep"] = model_info.params.get("num_keep", None)
|
||||
|
||||
if model_info.params.get("repeat_last_n", None):
|
||||
payload["options"]["repeat_last_n"] = model_info.params.get(
|
||||
"repeat_last_n", None
|
||||
@@ -764,7 +772,7 @@ async def generate_chat_completion(
|
||||
"frequency_penalty", None
|
||||
)
|
||||
|
||||
if model_info.params.get("temperature", None):
|
||||
if model_info.params.get("temperature", None) is not None:
|
||||
payload["options"]["temperature"] = model_info.params.get(
|
||||
"temperature", None
|
||||
)
|
||||
@@ -849,9 +857,14 @@ async def generate_chat_completion(
|
||||
|
||||
|
||||
# TODO: we should update this part once Ollama supports other types
|
||||
class OpenAIChatMessageContent(BaseModel):
|
||||
type: str
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
|
||||
class OpenAIChatMessage(BaseModel):
|
||||
role: str
|
||||
content: str
|
||||
content: Union[str, OpenAIChatMessageContent]
|
||||
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
@@ -879,7 +892,6 @@ async def generate_openai_chat_completion(
|
||||
model_info = Models.get_model_by_id(model_id)
|
||||
|
||||
if model_info:
|
||||
print(model_info)
|
||||
if model_info.base_model_id:
|
||||
payload["model"] = model_info.base_model_id
|
||||
|
||||
|
||||
+134
-87
@@ -345,113 +345,97 @@ async def get_models(url_idx: Optional[int] = None, user=Depends(get_current_use
|
||||
)
|
||||
|
||||
|
||||
@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
|
||||
async def proxy(path: str, request: Request, user=Depends(get_verified_user)):
|
||||
@app.post("/chat/completions")
|
||||
@app.post("/chat/completions/{url_idx}")
|
||||
async def generate_chat_completion(
|
||||
form_data: dict,
|
||||
url_idx: Optional[int] = None,
|
||||
user=Depends(get_verified_user),
|
||||
):
|
||||
idx = 0
|
||||
payload = {**form_data}
|
||||
|
||||
body = await request.body()
|
||||
# TODO: Remove below after gpt-4-vision fix from Open AI
|
||||
# Try to decode the body of the request from bytes to a UTF-8 string (Require add max_token to fix gpt-4-vision)
|
||||
model_id = form_data.get("model")
|
||||
model_info = Models.get_model_by_id(model_id)
|
||||
|
||||
payload = None
|
||||
if model_info:
|
||||
if model_info.base_model_id:
|
||||
payload["model"] = model_info.base_model_id
|
||||
|
||||
try:
|
||||
if "chat/completions" in path:
|
||||
body = body.decode("utf-8")
|
||||
body = json.loads(body)
|
||||
model_info.params = model_info.params.model_dump()
|
||||
|
||||
payload = {**body}
|
||||
if model_info.params:
|
||||
if model_info.params.get("temperature", None) is not None:
|
||||
payload["temperature"] = float(model_info.params.get("temperature"))
|
||||
|
||||
model_id = body.get("model")
|
||||
model_info = Models.get_model_by_id(model_id)
|
||||
if model_info.params.get("top_p", None):
|
||||
payload["top_p"] = int(model_info.params.get("top_p", None))
|
||||
|
||||
if model_info:
|
||||
print(model_info)
|
||||
if model_info.base_model_id:
|
||||
payload["model"] = model_info.base_model_id
|
||||
if model_info.params.get("max_tokens", None):
|
||||
payload["max_tokens"] = int(model_info.params.get("max_tokens", None))
|
||||
|
||||
model_info.params = model_info.params.model_dump()
|
||||
if model_info.params.get("frequency_penalty", None):
|
||||
payload["frequency_penalty"] = int(
|
||||
model_info.params.get("frequency_penalty", None)
|
||||
)
|
||||
|
||||
if model_info.params:
|
||||
if model_info.params.get("temperature", None):
|
||||
payload["temperature"] = int(
|
||||
model_info.params.get("temperature")
|
||||
if model_info.params.get("seed", None):
|
||||
payload["seed"] = model_info.params.get("seed", None)
|
||||
|
||||
if model_info.params.get("stop", None):
|
||||
payload["stop"] = (
|
||||
[
|
||||
bytes(stop, "utf-8").decode("unicode_escape")
|
||||
for stop in model_info.params["stop"]
|
||||
]
|
||||
if model_info.params.get("stop", None)
|
||||
else None
|
||||
)
|
||||
|
||||
if model_info.params.get("system", None):
|
||||
# Check if the payload already has a system message
|
||||
# If not, add a system message to the payload
|
||||
if payload.get("messages"):
|
||||
for message in payload["messages"]:
|
||||
if message.get("role") == "system":
|
||||
message["content"] = (
|
||||
model_info.params.get("system", None) + message["content"]
|
||||
)
|
||||
break
|
||||
else:
|
||||
payload["messages"].insert(
|
||||
0,
|
||||
{
|
||||
"role": "system",
|
||||
"content": model_info.params.get("system", None),
|
||||
},
|
||||
)
|
||||
|
||||
if model_info.params.get("top_p", None):
|
||||
payload["top_p"] = int(model_info.params.get("top_p", None))
|
||||
else:
|
||||
pass
|
||||
|
||||
if model_info.params.get("max_tokens", None):
|
||||
payload["max_tokens"] = int(
|
||||
model_info.params.get("max_tokens", None)
|
||||
)
|
||||
model = app.state.MODELS[payload.get("model")]
|
||||
idx = model["urlIdx"]
|
||||
|
||||
if model_info.params.get("frequency_penalty", None):
|
||||
payload["frequency_penalty"] = int(
|
||||
model_info.params.get("frequency_penalty", None)
|
||||
)
|
||||
if "pipeline" in model and model.get("pipeline"):
|
||||
payload["user"] = {"name": user.name, "id": user.id}
|
||||
|
||||
if model_info.params.get("seed", None):
|
||||
payload["seed"] = model_info.params.get("seed", None)
|
||||
# Check if the model is "gpt-4-vision-preview" and set "max_tokens" to 4000
|
||||
# This is a workaround until OpenAI fixes the issue with this model
|
||||
if payload.get("model") == "gpt-4-vision-preview":
|
||||
if "max_tokens" not in payload:
|
||||
payload["max_tokens"] = 4000
|
||||
log.debug("Modified payload:", payload)
|
||||
|
||||
if model_info.params.get("stop", None):
|
||||
payload["stop"] = (
|
||||
[
|
||||
bytes(stop, "utf-8").decode("unicode_escape")
|
||||
for stop in model_info.params["stop"]
|
||||
]
|
||||
if model_info.params.get("stop", None)
|
||||
else None
|
||||
)
|
||||
|
||||
if model_info.params.get("system", None):
|
||||
# Check if the payload already has a system message
|
||||
# If not, add a system message to the payload
|
||||
if payload.get("messages"):
|
||||
for message in payload["messages"]:
|
||||
if message.get("role") == "system":
|
||||
message["content"] = (
|
||||
model_info.params.get("system", None)
|
||||
+ message["content"]
|
||||
)
|
||||
break
|
||||
else:
|
||||
payload["messages"].insert(
|
||||
0,
|
||||
{
|
||||
"role": "system",
|
||||
"content": model_info.params.get("system", None),
|
||||
},
|
||||
)
|
||||
else:
|
||||
pass
|
||||
|
||||
model = app.state.MODELS[payload.get("model")]
|
||||
|
||||
idx = model["urlIdx"]
|
||||
|
||||
if "pipeline" in model and model.get("pipeline"):
|
||||
payload["user"] = {"name": user.name, "id": user.id}
|
||||
|
||||
# Check if the model is "gpt-4-vision-preview" and set "max_tokens" to 4000
|
||||
# This is a workaround until OpenAI fixes the issue with this model
|
||||
if payload.get("model") == "gpt-4-vision-preview":
|
||||
if "max_tokens" not in payload:
|
||||
payload["max_tokens"] = 4000
|
||||
log.debug("Modified payload:", payload)
|
||||
|
||||
# Convert the modified body back to JSON
|
||||
payload = json.dumps(payload)
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
log.error("Error loading request body into a dictionary:", e)
|
||||
# Convert the modified body back to JSON
|
||||
payload = json.dumps(payload)
|
||||
|
||||
print(payload)
|
||||
|
||||
url = app.state.config.OPENAI_API_BASE_URLS[idx]
|
||||
key = app.state.config.OPENAI_API_KEYS[idx]
|
||||
|
||||
target_url = f"{url}/{path}"
|
||||
print(payload)
|
||||
|
||||
headers = {}
|
||||
headers["Authorization"] = f"Bearer {key}"
|
||||
@@ -464,9 +448,72 @@ async def proxy(path: str, request: Request, user=Depends(get_verified_user)):
|
||||
try:
|
||||
session = aiohttp.ClientSession(trust_env=True)
|
||||
r = await session.request(
|
||||
method="POST",
|
||||
url=f"{url}/chat/completions",
|
||||
data=payload,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
r.raise_for_status()
|
||||
|
||||
# Check if response is SSE
|
||||
if "text/event-stream" in r.headers.get("Content-Type", ""):
|
||||
streaming = True
|
||||
return StreamingResponse(
|
||||
r.content,
|
||||
status_code=r.status,
|
||||
headers=dict(r.headers),
|
||||
background=BackgroundTask(
|
||||
cleanup_response, response=r, session=session
|
||||
),
|
||||
)
|
||||
else:
|
||||
response_data = await r.json()
|
||||
return response_data
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
error_detail = "Open WebUI: Server Connection Error"
|
||||
if r is not None:
|
||||
try:
|
||||
res = await r.json()
|
||||
print(res)
|
||||
if "error" in res:
|
||||
error_detail = f"External: {res['error']['message'] if 'message' in res['error'] else res['error']}"
|
||||
except:
|
||||
error_detail = f"External: {e}"
|
||||
raise HTTPException(status_code=r.status if r else 500, detail=error_detail)
|
||||
finally:
|
||||
if not streaming and session:
|
||||
if r:
|
||||
r.close()
|
||||
await session.close()
|
||||
|
||||
|
||||
@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
|
||||
async def proxy(path: str, request: Request, user=Depends(get_verified_user)):
|
||||
idx = 0
|
||||
|
||||
body = await request.body()
|
||||
|
||||
url = app.state.config.OPENAI_API_BASE_URLS[idx]
|
||||
key = app.state.config.OPENAI_API_KEYS[idx]
|
||||
|
||||
target_url = f"{url}/{path}"
|
||||
|
||||
headers = {}
|
||||
headers["Authorization"] = f"Bearer {key}"
|
||||
headers["Content-Type"] = "application/json"
|
||||
|
||||
r = None
|
||||
session = None
|
||||
streaming = False
|
||||
|
||||
try:
|
||||
session = aiohttp.ClientSession(trust_env=True)
|
||||
r = await session.request(
|
||||
method=request.method,
|
||||
url=target_url,
|
||||
data=payload if payload else body,
|
||||
data=body,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
|
||||
@@ -8,12 +8,15 @@ from fastapi import (
|
||||
Form,
|
||||
)
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
import requests
|
||||
import os, shutil, logging, re
|
||||
from datetime import datetime
|
||||
|
||||
from pathlib import Path
|
||||
from typing import List, Union, Sequence
|
||||
from typing import List, Union, Sequence, Iterator, Any
|
||||
|
||||
from chromadb.utils.batch_utils import create_batches
|
||||
from langchain_core.documents import Document
|
||||
|
||||
from langchain_community.document_loaders import (
|
||||
WebBaseLoader,
|
||||
@@ -30,6 +33,7 @@ from langchain_community.document_loaders import (
|
||||
UnstructuredExcelLoader,
|
||||
UnstructuredPowerPointLoader,
|
||||
YoutubeLoader,
|
||||
OutlookMessageLoader,
|
||||
)
|
||||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||||
|
||||
@@ -67,7 +71,9 @@ from apps.rag.search.main import SearchResult
|
||||
from apps.rag.search.searxng import search_searxng
|
||||
from apps.rag.search.serper import search_serper
|
||||
from apps.rag.search.serpstack import search_serpstack
|
||||
|
||||
from apps.rag.search.serply import search_serply
|
||||
from apps.rag.search.duckduckgo import search_duckduckgo
|
||||
from apps.rag.search.tavily import search_tavily
|
||||
|
||||
from utils.misc import (
|
||||
calculate_sha256,
|
||||
@@ -113,6 +119,8 @@ from config import (
|
||||
SERPSTACK_API_KEY,
|
||||
SERPSTACK_HTTPS,
|
||||
SERPER_API_KEY,
|
||||
SERPLY_API_KEY,
|
||||
TAVILY_API_KEY,
|
||||
RAG_WEB_SEARCH_RESULT_COUNT,
|
||||
RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
|
||||
RAG_EMBEDDING_OPENAI_BATCH_SIZE,
|
||||
@@ -165,6 +173,8 @@ app.state.config.BRAVE_SEARCH_API_KEY = BRAVE_SEARCH_API_KEY
|
||||
app.state.config.SERPSTACK_API_KEY = SERPSTACK_API_KEY
|
||||
app.state.config.SERPSTACK_HTTPS = SERPSTACK_HTTPS
|
||||
app.state.config.SERPER_API_KEY = SERPER_API_KEY
|
||||
app.state.config.SERPLY_API_KEY = SERPLY_API_KEY
|
||||
app.state.config.TAVILY_API_KEY = TAVILY_API_KEY
|
||||
app.state.config.RAG_WEB_SEARCH_RESULT_COUNT = RAG_WEB_SEARCH_RESULT_COUNT
|
||||
app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS = RAG_WEB_SEARCH_CONCURRENT_REQUESTS
|
||||
|
||||
@@ -392,6 +402,8 @@ async def get_rag_config(user=Depends(get_admin_user)):
|
||||
"serpstack_api_key": app.state.config.SERPSTACK_API_KEY,
|
||||
"serpstack_https": app.state.config.SERPSTACK_HTTPS,
|
||||
"serper_api_key": app.state.config.SERPER_API_KEY,
|
||||
"serply_api_key": app.state.config.SERPLY_API_KEY,
|
||||
"tavily_api_key": app.state.config.TAVILY_API_KEY,
|
||||
"result_count": app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
|
||||
"concurrent_requests": app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
|
||||
},
|
||||
@@ -419,6 +431,8 @@ class WebSearchConfig(BaseModel):
|
||||
serpstack_api_key: Optional[str] = None
|
||||
serpstack_https: Optional[bool] = None
|
||||
serper_api_key: Optional[str] = None
|
||||
serply_api_key: Optional[str] = None
|
||||
tavily_api_key: Optional[str] = None
|
||||
result_count: Optional[int] = None
|
||||
concurrent_requests: Optional[int] = None
|
||||
|
||||
@@ -469,6 +483,8 @@ async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_
|
||||
app.state.config.SERPSTACK_API_KEY = form_data.web.search.serpstack_api_key
|
||||
app.state.config.SERPSTACK_HTTPS = form_data.web.search.serpstack_https
|
||||
app.state.config.SERPER_API_KEY = form_data.web.search.serper_api_key
|
||||
app.state.config.SERPLY_API_KEY = form_data.web.search.serply_api_key
|
||||
app.state.config.TAVILY_API_KEY = form_data.web.search.tavily_api_key
|
||||
app.state.config.RAG_WEB_SEARCH_RESULT_COUNT = form_data.web.search.result_count
|
||||
app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS = (
|
||||
form_data.web.search.concurrent_requests
|
||||
@@ -497,6 +513,8 @@ async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_
|
||||
"serpstack_api_key": app.state.config.SERPSTACK_API_KEY,
|
||||
"serpstack_https": app.state.config.SERPSTACK_HTTPS,
|
||||
"serper_api_key": app.state.config.SERPER_API_KEY,
|
||||
"serply_api_key": app.state.config.SERPLY_API_KEY,
|
||||
"tavily_api_key": app.state.config.TAVILY_API_KEY,
|
||||
"result_count": app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
|
||||
"concurrent_requests": app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
|
||||
},
|
||||
@@ -693,7 +711,7 @@ def get_web_loader(url: Union[str, Sequence[str]], verify_ssl: bool = True):
|
||||
# Check if the URL is valid
|
||||
if not validate_url(url):
|
||||
raise ValueError(ERROR_MESSAGES.INVALID_URL)
|
||||
return WebBaseLoader(
|
||||
return SafeWebBaseLoader(
|
||||
url,
|
||||
verify_ssl=verify_ssl,
|
||||
requests_per_second=RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
|
||||
@@ -744,7 +762,8 @@ def search_web(engine: str, query: str) -> list[SearchResult]:
|
||||
- BRAVE_SEARCH_API_KEY
|
||||
- SERPSTACK_API_KEY
|
||||
- SERPER_API_KEY
|
||||
|
||||
- SERPLY_API_KEY
|
||||
- TAVILY_API_KEY
|
||||
Args:
|
||||
query (str): The query to search for
|
||||
"""
|
||||
@@ -802,6 +821,26 @@ def search_web(engine: str, query: str) -> list[SearchResult]:
|
||||
)
|
||||
else:
|
||||
raise Exception("No SERPER_API_KEY found in environment variables")
|
||||
elif engine == "serply":
|
||||
if app.state.config.SERPLY_API_KEY:
|
||||
return search_serply(
|
||||
app.state.config.SERPLY_API_KEY,
|
||||
query,
|
||||
app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
|
||||
)
|
||||
else:
|
||||
raise Exception("No SERPLY_API_KEY found in environment variables")
|
||||
elif engine == "duckduckgo":
|
||||
return search_duckduckgo(query, app.state.config.RAG_WEB_SEARCH_RESULT_COUNT)
|
||||
elif engine == "tavily":
|
||||
if app.state.config.TAVILY_API_KEY:
|
||||
return search_tavily(
|
||||
app.state.config.TAVILY_API_KEY,
|
||||
query,
|
||||
app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
|
||||
)
|
||||
else:
|
||||
raise Exception("No TAVILY_API_KEY found in environment variables")
|
||||
else:
|
||||
raise Exception("No search engine API key found in environment variables")
|
||||
|
||||
@@ -809,6 +848,9 @@ def search_web(engine: str, query: str) -> list[SearchResult]:
|
||||
@app.post("/web/search")
|
||||
def store_web_search(form_data: SearchForm, user=Depends(get_current_user)):
|
||||
try:
|
||||
logging.info(
|
||||
f"trying to web search with {app.state.config.RAG_WEB_SEARCH_ENGINE, form_data.query}"
|
||||
)
|
||||
web_results = search_web(
|
||||
app.state.config.RAG_WEB_SEARCH_ENGINE, form_data.query
|
||||
)
|
||||
@@ -879,6 +921,13 @@ def store_docs_in_vector_db(docs, collection_name, overwrite: bool = False) -> b
|
||||
texts = [doc.page_content for doc in docs]
|
||||
metadatas = [doc.metadata for doc in docs]
|
||||
|
||||
# ChromaDB does not like datetime formats
|
||||
# for meta-data so convert them to string.
|
||||
for metadata in metadatas:
|
||||
for key, value in metadata.items():
|
||||
if isinstance(value, datetime):
|
||||
metadata[key] = str(value)
|
||||
|
||||
try:
|
||||
if overwrite:
|
||||
for collection in CHROMA_CLIENT.list_collections():
|
||||
@@ -965,6 +1014,7 @@ def get_loader(filename: str, file_content_type: str, file_path: str):
|
||||
"swift",
|
||||
"vue",
|
||||
"svelte",
|
||||
"msg",
|
||||
]
|
||||
|
||||
if file_ext == "pdf":
|
||||
@@ -999,6 +1049,8 @@ def get_loader(filename: str, file_content_type: str, file_path: str):
|
||||
"application/vnd.openxmlformats-officedocument.presentationml.presentation",
|
||||
] or file_ext in ["ppt", "pptx"]:
|
||||
loader = UnstructuredPowerPointLoader(file_path)
|
||||
elif file_ext == "msg":
|
||||
loader = OutlookMessageLoader(file_path)
|
||||
elif file_ext in known_source_ext or (
|
||||
file_content_type and file_content_type.find("text/") >= 0
|
||||
):
|
||||
@@ -1209,6 +1261,33 @@ def reset(user=Depends(get_admin_user)) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
class SafeWebBaseLoader(WebBaseLoader):
|
||||
"""WebBaseLoader with enhanced error handling for URLs."""
|
||||
|
||||
def lazy_load(self) -> Iterator[Document]:
|
||||
"""Lazy load text from the url(s) in web_path with error handling."""
|
||||
for path in self.web_paths:
|
||||
try:
|
||||
soup = self._scrape(path, bs_kwargs=self.bs_kwargs)
|
||||
text = soup.get_text(**self.bs_get_text_kwargs)
|
||||
|
||||
# Build metadata
|
||||
metadata = {"source": path}
|
||||
if title := soup.find("title"):
|
||||
metadata["title"] = title.get_text()
|
||||
if description := soup.find("meta", attrs={"name": "description"}):
|
||||
metadata["description"] = description.get(
|
||||
"content", "No description found."
|
||||
)
|
||||
if html := soup.find("html"):
|
||||
metadata["language"] = html.get("lang", "No language found.")
|
||||
|
||||
yield Document(page_content=text, metadata=metadata)
|
||||
except Exception as e:
|
||||
# Log the error and continue with the next URL
|
||||
log.error(f"Error loading {path}: {e}")
|
||||
|
||||
|
||||
if ENV == "dev":
|
||||
|
||||
@app.get("/ef")
|
||||
|
||||
@@ -0,0 +1,46 @@
|
||||
import logging
|
||||
|
||||
from apps.rag.search.main import SearchResult
|
||||
from duckduckgo_search import DDGS
|
||||
from config import SRC_LOG_LEVELS
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["RAG"])
|
||||
|
||||
|
||||
def search_duckduckgo(query: str, count: int) -> list[SearchResult]:
|
||||
"""
|
||||
Search using DuckDuckGo's Search API and return the results as a list of SearchResult objects.
|
||||
Args:
|
||||
query (str): The query to search for
|
||||
count (int): The number of results to return
|
||||
|
||||
Returns:
|
||||
List[SearchResult]: A list of search results
|
||||
"""
|
||||
# Use the DDGS context manager to create a DDGS object
|
||||
with DDGS() as ddgs:
|
||||
# Use the ddgs.text() method to perform the search
|
||||
ddgs_gen = ddgs.text(
|
||||
query, safesearch="moderate", max_results=count, backend="api"
|
||||
)
|
||||
# Check if there are search results
|
||||
if ddgs_gen:
|
||||
# Convert the search results into a list
|
||||
search_results = [r for r in ddgs_gen]
|
||||
|
||||
# Create an empty list to store the SearchResult objects
|
||||
results = []
|
||||
# Iterate over each search result
|
||||
for result in search_results:
|
||||
# Create a SearchResult object and append it to the results list
|
||||
results.append(
|
||||
SearchResult(
|
||||
link=result["href"],
|
||||
title=result.get("title"),
|
||||
snippet=result.get("body"),
|
||||
)
|
||||
)
|
||||
print(results)
|
||||
# Return the list of search results
|
||||
return results
|
||||
@@ -25,6 +25,7 @@ def search_searxng(
|
||||
|
||||
Keyword Args:
|
||||
language (str): Language filter for the search results; e.g., "en-US". Defaults to an empty string.
|
||||
safesearch (int): Safe search filter for safer web results; 0 = off, 1 = moderate, 2 = strict. Defaults to 1 (moderate).
|
||||
time_range (str): Time range for filtering results by date; e.g., "2023-04-05..today" or "all-time". Defaults to ''.
|
||||
categories: (Optional[List[str]]): Specific categories within which the search should be performed, defaulting to an empty string if not provided.
|
||||
|
||||
@@ -37,6 +38,7 @@ def search_searxng(
|
||||
|
||||
# Default values for optional parameters are provided as empty strings or None when not specified.
|
||||
language = kwargs.get("language", "en-US")
|
||||
safesearch = kwargs.get("safesearch", "1")
|
||||
time_range = kwargs.get("time_range", "")
|
||||
categories = "".join(kwargs.get("categories", []))
|
||||
|
||||
@@ -44,6 +46,7 @@ def search_searxng(
|
||||
"q": query,
|
||||
"format": "json",
|
||||
"pageno": 1,
|
||||
"safesearch": safesearch,
|
||||
"language": language,
|
||||
"time_range": time_range,
|
||||
"categories": categories,
|
||||
|
||||
@@ -0,0 +1,68 @@
|
||||
import json
|
||||
import logging
|
||||
|
||||
import requests
|
||||
from urllib.parse import urlencode
|
||||
|
||||
from apps.rag.search.main import SearchResult
|
||||
from config import SRC_LOG_LEVELS
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["RAG"])
|
||||
|
||||
|
||||
def search_serply(
|
||||
api_key: str,
|
||||
query: str,
|
||||
count: int,
|
||||
hl: str = "us",
|
||||
limit: int = 10,
|
||||
device_type: str = "desktop",
|
||||
proxy_location: str = "US",
|
||||
) -> list[SearchResult]:
|
||||
"""Search using serper.dev's API and return the results as a list of SearchResult objects.
|
||||
|
||||
Args:
|
||||
api_key (str): A serply.io API key
|
||||
query (str): The query to search for
|
||||
hl (str): Host Language code to display results in (reference https://developers.google.com/custom-search/docs/xml_results?hl=en#wsInterfaceLanguages)
|
||||
limit (int): The maximum number of results to return [10-100, defaults to 10]
|
||||
"""
|
||||
log.info("Searching with Serply")
|
||||
|
||||
url = "https://api.serply.io/v1/search/"
|
||||
|
||||
query_payload = {
|
||||
"q": query,
|
||||
"language": "en",
|
||||
"num": limit,
|
||||
"gl": proxy_location.upper(),
|
||||
"hl": hl.lower(),
|
||||
}
|
||||
|
||||
url = f"{url}{urlencode(query_payload)}"
|
||||
headers = {
|
||||
"X-API-KEY": api_key,
|
||||
"X-User-Agent": device_type,
|
||||
"User-Agent": "open-webui",
|
||||
"X-Proxy-Location": proxy_location,
|
||||
}
|
||||
|
||||
response = requests.request("GET", url, headers=headers)
|
||||
response.raise_for_status()
|
||||
|
||||
json_response = response.json()
|
||||
log.info(f"results from serply search: {json_response}")
|
||||
|
||||
results = sorted(
|
||||
json_response.get("results", []), key=lambda x: x.get("realPosition", 0)
|
||||
)
|
||||
|
||||
return [
|
||||
SearchResult(
|
||||
link=result["link"],
|
||||
title=result.get("title"),
|
||||
snippet=result.get("description"),
|
||||
)
|
||||
for result in results[:count]
|
||||
]
|
||||
@@ -0,0 +1,39 @@
|
||||
import logging
|
||||
|
||||
import requests
|
||||
|
||||
from apps.rag.search.main import SearchResult
|
||||
from config import SRC_LOG_LEVELS
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["RAG"])
|
||||
|
||||
|
||||
def search_tavily(api_key: str, query: str, count: int) -> list[SearchResult]:
|
||||
"""Search using Tavily's Search API and return the results as a list of SearchResult objects.
|
||||
|
||||
Args:
|
||||
api_key (str): A Tavily Search API key
|
||||
query (str): The query to search for
|
||||
|
||||
Returns:
|
||||
List[SearchResult]: A list of search results
|
||||
"""
|
||||
url = "https://api.tavily.com/search"
|
||||
data = {"query": query, "api_key": api_key}
|
||||
|
||||
response = requests.post(url, json=data)
|
||||
response.raise_for_status()
|
||||
|
||||
json_response = response.json()
|
||||
|
||||
raw_search_results = json_response.get("results", [])
|
||||
|
||||
return [
|
||||
SearchResult(
|
||||
link=result["url"],
|
||||
title=result.get("title", ""),
|
||||
snippet=result.get("content"),
|
||||
)
|
||||
for result in raw_search_results[:count]
|
||||
]
|
||||
+206
@@ -0,0 +1,206 @@
|
||||
{
|
||||
"ads": [],
|
||||
"ads_count": 0,
|
||||
"answers": [],
|
||||
"results": [
|
||||
{
|
||||
"title": "Apple",
|
||||
"link": "https://www.apple.com/",
|
||||
"description": "Discover the innovative world of Apple and shop everything iPhone, iPad, Apple Watch, Mac, and Apple TV, plus explore accessories, entertainment, ...",
|
||||
"additional_links": [
|
||||
{
|
||||
"text": "AppleApplehttps://www.apple.com",
|
||||
"href": "https://www.apple.com/"
|
||||
}
|
||||
],
|
||||
"cite": {},
|
||||
"subdomains": [
|
||||
{
|
||||
"title": "Support",
|
||||
"link": "https://support.apple.com/",
|
||||
"description": "SupportContact - iPhone Support - Billing and Subscriptions - Apple Repair"
|
||||
},
|
||||
{
|
||||
"title": "Store",
|
||||
"link": "https://www.apple.com/store",
|
||||
"description": "StoreShop iPhone - Shop iPad - App Store - Shop Mac - ..."
|
||||
},
|
||||
{
|
||||
"title": "Mac",
|
||||
"link": "https://www.apple.com/mac/",
|
||||
"description": "MacMacBook Air - MacBook Pro - iMac - Compare Mac models - Mac mini"
|
||||
},
|
||||
{
|
||||
"title": "iPad",
|
||||
"link": "https://www.apple.com/ipad/",
|
||||
"description": "iPadShop iPad - iPad Pro - iPad Air - Compare iPad models - ..."
|
||||
},
|
||||
{
|
||||
"title": "Watch",
|
||||
"link": "https://www.apple.com/watch/",
|
||||
"description": "WatchShop Apple Watch - Series 9 - SE - Ultra 2 - Nike - Hermès - ..."
|
||||
}
|
||||
],
|
||||
"realPosition": 1
|
||||
},
|
||||
{
|
||||
"title": "Apple",
|
||||
"link": "https://www.apple.com/",
|
||||
"description": "Discover the innovative world of Apple and shop everything iPhone, iPad, Apple Watch, Mac, and Apple TV, plus explore accessories, entertainment, ...",
|
||||
"additional_links": [
|
||||
{
|
||||
"text": "AppleApplehttps://www.apple.com",
|
||||
"href": "https://www.apple.com/"
|
||||
}
|
||||
],
|
||||
"cite": {},
|
||||
"realPosition": 2
|
||||
},
|
||||
{
|
||||
"title": "Apple Inc.",
|
||||
"link": "https://en.wikipedia.org/wiki/Apple_Inc.",
|
||||
"description": "Apple Inc. (formerly Apple Computer, Inc.) is an American multinational corporation and technology company headquartered in Cupertino, California, ...",
|
||||
"additional_links": [
|
||||
{
|
||||
"text": "Apple Inc.Wikipediahttps://en.wikipedia.org › wiki › Apple_Inc",
|
||||
"href": "https://en.wikipedia.org/wiki/Apple_Inc."
|
||||
},
|
||||
{
|
||||
"text": "",
|
||||
"href": "https://en.wikipedia.org/wiki/Apple_Inc."
|
||||
},
|
||||
{
|
||||
"text": "History",
|
||||
"href": "https://en.wikipedia.org/wiki/History_of_Apple_Inc."
|
||||
},
|
||||
{
|
||||
"text": "List of Apple products",
|
||||
"href": "https://en.wikipedia.org/wiki/List_of_Apple_products"
|
||||
},
|
||||
{
|
||||
"text": "Litigation involving Apple Inc.",
|
||||
"href": "https://en.wikipedia.org/wiki/Litigation_involving_Apple_Inc."
|
||||
},
|
||||
{
|
||||
"text": "Apple Park",
|
||||
"href": "https://en.wikipedia.org/wiki/Apple_Park"
|
||||
}
|
||||
],
|
||||
"cite": {
|
||||
"domain": "https://en.wikipedia.org › wiki › Apple_Inc",
|
||||
"span": " › wiki › Apple_Inc"
|
||||
},
|
||||
"realPosition": 3
|
||||
},
|
||||
{
|
||||
"title": "Apple Inc. (AAPL) Company Profile & Facts",
|
||||
"link": "https://finance.yahoo.com/quote/AAPL/profile/",
|
||||
"description": "Apple Inc. designs, manufactures, and markets smartphones, personal computers, tablets, wearables, and accessories worldwide. The company offers iPhone, a line ...",
|
||||
"additional_links": [
|
||||
{
|
||||
"text": "Apple Inc. (AAPL) Company Profile & FactsYahoo Financehttps://finance.yahoo.com › quote › AAPL › profile",
|
||||
"href": "https://finance.yahoo.com/quote/AAPL/profile/"
|
||||
}
|
||||
],
|
||||
"cite": {
|
||||
"domain": "https://finance.yahoo.com › quote › AAPL › profile",
|
||||
"span": " › quote › AAPL › profile"
|
||||
},
|
||||
"realPosition": 4
|
||||
},
|
||||
{
|
||||
"title": "Apple Inc - Company Profile and News",
|
||||
"link": "https://www.bloomberg.com/profile/company/AAPL:US",
|
||||
"description": "Apple Inc. Apple Inc. designs, manufactures, and markets smartphones, personal computers, tablets, wearables and accessories, and sells a variety of related ...",
|
||||
"additional_links": [
|
||||
{
|
||||
"text": "Apple Inc - Company Profile and NewsBloomberghttps://www.bloomberg.com › company › AAPL:US",
|
||||
"href": "https://www.bloomberg.com/profile/company/AAPL:US"
|
||||
},
|
||||
{
|
||||
"text": "",
|
||||
"href": "https://www.bloomberg.com/profile/company/AAPL:US"
|
||||
}
|
||||
],
|
||||
"cite": {
|
||||
"domain": "https://www.bloomberg.com › company › AAPL:US",
|
||||
"span": " › company › AAPL:US"
|
||||
},
|
||||
"realPosition": 5
|
||||
},
|
||||
{
|
||||
"title": "Apple Inc. | History, Products, Headquarters, & Facts",
|
||||
"link": "https://www.britannica.com/money/Apple-Inc",
|
||||
"description": "May 22, 2024 — Apple Inc. is an American multinational technology company that revolutionized the technology sector through its innovation of computer ...",
|
||||
"additional_links": [
|
||||
{
|
||||
"text": "Apple Inc. | History, Products, Headquarters, & FactsBritannicahttps://www.britannica.com › money › Apple-Inc",
|
||||
"href": "https://www.britannica.com/money/Apple-Inc"
|
||||
},
|
||||
{
|
||||
"text": "",
|
||||
"href": "https://www.britannica.com/money/Apple-Inc"
|
||||
}
|
||||
],
|
||||
"cite": {
|
||||
"domain": "https://www.britannica.com › money › Apple-Inc",
|
||||
"span": " › money › Apple-Inc"
|
||||
},
|
||||
"realPosition": 6
|
||||
}
|
||||
],
|
||||
"shopping_ads": [],
|
||||
"places": [
|
||||
{
|
||||
"title": "Apple Inc."
|
||||
},
|
||||
{
|
||||
"title": "Apple Inc"
|
||||
},
|
||||
{
|
||||
"title": "Apple Inc"
|
||||
}
|
||||
],
|
||||
"related_searches": {
|
||||
"images": [],
|
||||
"text": [
|
||||
{
|
||||
"title": "apple inc full form",
|
||||
"link": "https://www.google.com/search?sca_esv=6b6df170a5c9891b&sca_upv=1&q=Apple+Inc+full+form&sa=X&ved=2ahUKEwjLxuSJwM-GAxUHODQIHYuJBhgQ1QJ6BAhPEAE"
|
||||
},
|
||||
{
|
||||
"title": "apple company history",
|
||||
"link": "https://www.google.com/search?sca_esv=6b6df170a5c9891b&sca_upv=1&q=Apple+company+history&sa=X&ved=2ahUKEwjLxuSJwM-GAxUHODQIHYuJBhgQ1QJ6BAhOEAE"
|
||||
},
|
||||
{
|
||||
"title": "apple store",
|
||||
"link": "https://www.google.com/search?sca_esv=6b6df170a5c9891b&sca_upv=1&q=Apple+Store&sa=X&ved=2ahUKEwjLxuSJwM-GAxUHODQIHYuJBhgQ1QJ6BAhQEAE"
|
||||
},
|
||||
{
|
||||
"title": "apple id",
|
||||
"link": "https://www.google.com/search?sca_esv=6b6df170a5c9891b&sca_upv=1&q=Apple+id&sa=X&ved=2ahUKEwjLxuSJwM-GAxUHODQIHYuJBhgQ1QJ6BAhSEAE"
|
||||
},
|
||||
{
|
||||
"title": "apple inc industry",
|
||||
"link": "https://www.google.com/search?sca_esv=6b6df170a5c9891b&sca_upv=1&q=Apple+Inc+industry&sa=X&ved=2ahUKEwjLxuSJwM-GAxUHODQIHYuJBhgQ1QJ6BAhREAE"
|
||||
},
|
||||
{
|
||||
"title": "apple login",
|
||||
"link": "https://www.google.com/search?sca_esv=6b6df170a5c9891b&sca_upv=1&q=Apple+login&sa=X&ved=2ahUKEwjLxuSJwM-GAxUHODQIHYuJBhgQ1QJ6BAhTEAE"
|
||||
}
|
||||
]
|
||||
},
|
||||
"image_results": [],
|
||||
"carousel": [],
|
||||
"total": 2450000000,
|
||||
"knowledge_graph": "",
|
||||
"related_questions": [
|
||||
"What does the Apple Inc do?",
|
||||
"Why did Apple change to Apple Inc?",
|
||||
"Who owns Apple Inc.?",
|
||||
"What is Apple Inc best known for?"
|
||||
],
|
||||
"carousel_count": 0,
|
||||
"ts": 2.491065263748169,
|
||||
"device_type": null
|
||||
}
|
||||
@@ -20,7 +20,7 @@ from langchain.retrievers import (
|
||||
|
||||
from typing import Optional
|
||||
|
||||
|
||||
from utils.misc import get_last_user_message, add_or_update_system_message
|
||||
from config import SRC_LOG_LEVELS, CHROMA_CLIENT
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
@@ -236,10 +236,9 @@ def get_embedding_function(
|
||||
return lambda query: generate_multiple(query, func)
|
||||
|
||||
|
||||
def rag_messages(
|
||||
def get_rag_context(
|
||||
docs,
|
||||
messages,
|
||||
template,
|
||||
embedding_function,
|
||||
k,
|
||||
reranking_function,
|
||||
@@ -247,31 +246,7 @@ def rag_messages(
|
||||
hybrid_search,
|
||||
):
|
||||
log.debug(f"docs: {docs} {messages} {embedding_function} {reranking_function}")
|
||||
|
||||
last_user_message_idx = None
|
||||
for i in range(len(messages) - 1, -1, -1):
|
||||
if messages[i]["role"] == "user":
|
||||
last_user_message_idx = i
|
||||
break
|
||||
|
||||
user_message = messages[last_user_message_idx]
|
||||
|
||||
if isinstance(user_message["content"], list):
|
||||
# Handle list content input
|
||||
content_type = "list"
|
||||
query = ""
|
||||
for content_item in user_message["content"]:
|
||||
if content_item["type"] == "text":
|
||||
query = content_item["text"]
|
||||
break
|
||||
elif isinstance(user_message["content"], str):
|
||||
# Handle text content input
|
||||
content_type = "text"
|
||||
query = user_message["content"]
|
||||
else:
|
||||
# Fallback in case the input does not match expected types
|
||||
content_type = None
|
||||
query = ""
|
||||
query = get_last_user_message(messages)
|
||||
|
||||
extracted_collections = []
|
||||
relevant_contexts = []
|
||||
@@ -342,33 +317,7 @@ def rag_messages(
|
||||
|
||||
context_string = context_string.strip()
|
||||
|
||||
ra_content = rag_template(
|
||||
template=template,
|
||||
context=context_string,
|
||||
query=query,
|
||||
)
|
||||
|
||||
log.debug(f"ra_content: {ra_content}")
|
||||
|
||||
if content_type == "list":
|
||||
new_content = []
|
||||
for content_item in user_message["content"]:
|
||||
if content_item["type"] == "text":
|
||||
# Update the text item's content with ra_content
|
||||
new_content.append({"type": "text", "text": ra_content})
|
||||
else:
|
||||
# Keep other types of content as they are
|
||||
new_content.append(content_item)
|
||||
new_user_message = {**user_message, "content": new_content}
|
||||
else:
|
||||
new_user_message = {
|
||||
**user_message,
|
||||
"content": ra_content,
|
||||
}
|
||||
|
||||
messages[last_user_message_idx] = new_user_message
|
||||
|
||||
return messages, citations
|
||||
return context_string, citations
|
||||
|
||||
|
||||
def get_model_path(model: str, update_model: bool = False):
|
||||
|
||||
+22
-15
@@ -10,7 +10,7 @@ app = socketio.ASGIApp(sio, socketio_path="/ws/socket.io")
|
||||
|
||||
# Dictionary to maintain the user pool
|
||||
|
||||
|
||||
SESSION_POOL = {}
|
||||
USER_POOL = {}
|
||||
USAGE_POOL = {}
|
||||
# Timeout duration in seconds
|
||||
@@ -19,8 +19,6 @@ TIMEOUT_DURATION = 3
|
||||
|
||||
@sio.event
|
||||
async def connect(sid, environ, auth):
|
||||
print("connect ", sid)
|
||||
|
||||
user = None
|
||||
if auth and "token" in auth:
|
||||
data = decode_token(auth["token"])
|
||||
@@ -29,10 +27,14 @@ async def connect(sid, environ, auth):
|
||||
user = Users.get_user_by_id(data["id"])
|
||||
|
||||
if user:
|
||||
USER_POOL[sid] = user.id
|
||||
SESSION_POOL[sid] = user.id
|
||||
if user.id in USER_POOL:
|
||||
USER_POOL[user.id].append(sid)
|
||||
else:
|
||||
USER_POOL[user.id] = [sid]
|
||||
|
||||
print(f"user {user.name}({user.id}) connected with session ID {sid}")
|
||||
|
||||
print(len(set(USER_POOL)))
|
||||
await sio.emit("user-count", {"count": len(set(USER_POOL))})
|
||||
await sio.emit("usage", {"models": get_models_in_use()})
|
||||
|
||||
@@ -50,16 +52,20 @@ async def user_join(sid, data):
|
||||
user = Users.get_user_by_id(data["id"])
|
||||
|
||||
if user:
|
||||
USER_POOL[sid] = user.id
|
||||
|
||||
SESSION_POOL[sid] = user.id
|
||||
if user.id in USER_POOL:
|
||||
USER_POOL[user.id].append(sid)
|
||||
else:
|
||||
USER_POOL[user.id] = [sid]
|
||||
|
||||
print(f"user {user.name}({user.id}) connected with session ID {sid}")
|
||||
|
||||
print(len(set(USER_POOL)))
|
||||
await sio.emit("user-count", {"count": len(set(USER_POOL))})
|
||||
|
||||
|
||||
@sio.on("user-count")
|
||||
async def user_count(sid):
|
||||
print("user-count", sid)
|
||||
await sio.emit("user-count", {"count": len(set(USER_POOL))})
|
||||
|
||||
|
||||
@@ -68,14 +74,12 @@ def get_models_in_use():
|
||||
models_in_use = []
|
||||
for model_id, data in USAGE_POOL.items():
|
||||
models_in_use.append(model_id)
|
||||
print(f"Models in use: {models_in_use}")
|
||||
|
||||
return models_in_use
|
||||
|
||||
|
||||
@sio.on("usage")
|
||||
async def usage(sid, data):
|
||||
print(f'Received "usage" event from {sid}: {data}')
|
||||
|
||||
model_id = data["model"]
|
||||
|
||||
@@ -103,7 +107,6 @@ async def usage(sid, data):
|
||||
|
||||
async def remove_after_timeout(sid, model_id):
|
||||
try:
|
||||
print("remove_after_timeout", sid, model_id)
|
||||
await asyncio.sleep(TIMEOUT_DURATION)
|
||||
if model_id in USAGE_POOL:
|
||||
print(USAGE_POOL[model_id]["sids"])
|
||||
@@ -113,7 +116,6 @@ async def remove_after_timeout(sid, model_id):
|
||||
if len(USAGE_POOL[model_id]["sids"]) == 0:
|
||||
del USAGE_POOL[model_id]
|
||||
|
||||
print(f"Removed usage data for {model_id} due to timeout")
|
||||
# Broadcast the usage data to all clients
|
||||
await sio.emit("usage", {"models": get_models_in_use()})
|
||||
except asyncio.CancelledError:
|
||||
@@ -123,9 +125,14 @@ async def remove_after_timeout(sid, model_id):
|
||||
|
||||
@sio.event
|
||||
async def disconnect(sid):
|
||||
if sid in USER_POOL:
|
||||
disconnected_user = USER_POOL.pop(sid)
|
||||
print(f"user {disconnected_user} disconnected with session ID {sid}")
|
||||
if sid in SESSION_POOL:
|
||||
user_id = SESSION_POOL[sid]
|
||||
del SESSION_POOL[sid]
|
||||
|
||||
USER_POOL[user_id].remove(sid)
|
||||
|
||||
if len(USER_POOL[user_id]) == 0:
|
||||
del USER_POOL[user_id]
|
||||
|
||||
await sio.emit("user-count", {"count": len(USER_POOL)})
|
||||
else:
|
||||
|
||||
@@ -0,0 +1,61 @@
|
||||
"""Peewee migrations -- 009_add_models.py.
|
||||
|
||||
Some examples (model - class or model name)::
|
||||
|
||||
> Model = migrator.orm['table_name'] # Return model in current state by name
|
||||
> Model = migrator.ModelClass # Return model in current state by name
|
||||
|
||||
> migrator.sql(sql) # Run custom SQL
|
||||
> migrator.run(func, *args, **kwargs) # Run python function with the given args
|
||||
> migrator.create_model(Model) # Create a model (could be used as decorator)
|
||||
> migrator.remove_model(model, cascade=True) # Remove a model
|
||||
> migrator.add_fields(model, **fields) # Add fields to a model
|
||||
> migrator.change_fields(model, **fields) # Change fields
|
||||
> migrator.remove_fields(model, *field_names, cascade=True)
|
||||
> migrator.rename_field(model, old_field_name, new_field_name)
|
||||
> migrator.rename_table(model, new_table_name)
|
||||
> migrator.add_index(model, *col_names, unique=False)
|
||||
> migrator.add_not_null(model, *field_names)
|
||||
> migrator.add_default(model, field_name, default)
|
||||
> migrator.add_constraint(model, name, sql)
|
||||
> migrator.drop_index(model, *col_names)
|
||||
> migrator.drop_not_null(model, *field_names)
|
||||
> migrator.drop_constraints(model, *constraints)
|
||||
|
||||
"""
|
||||
|
||||
from contextlib import suppress
|
||||
|
||||
import peewee as pw
|
||||
from peewee_migrate import Migrator
|
||||
|
||||
|
||||
with suppress(ImportError):
|
||||
import playhouse.postgres_ext as pw_pext
|
||||
|
||||
|
||||
def migrate(migrator: Migrator, database: pw.Database, *, fake=False):
|
||||
"""Write your migrations here."""
|
||||
|
||||
@migrator.create_model
|
||||
class Tool(pw.Model):
|
||||
id = pw.TextField(unique=True)
|
||||
user_id = pw.TextField()
|
||||
|
||||
name = pw.TextField()
|
||||
content = pw.TextField()
|
||||
specs = pw.TextField()
|
||||
|
||||
meta = pw.TextField()
|
||||
|
||||
created_at = pw.BigIntegerField(null=False)
|
||||
updated_at = pw.BigIntegerField(null=False)
|
||||
|
||||
class Meta:
|
||||
table_name = "tool"
|
||||
|
||||
|
||||
def rollback(migrator: Migrator, database: pw.Database, *, fake=False):
|
||||
"""Write your rollback migrations here."""
|
||||
|
||||
migrator.remove_model("tool")
|
||||
+1
-1
@@ -1,4 +1,4 @@
|
||||
"""Peewee migrations -- 011_add_user_oauth_sub.py.
|
||||
"""Peewee migrations -- 013_add_user_oauth_sub.py.
|
||||
|
||||
Some examples (model - class or model name)::
|
||||
|
||||
@@ -8,6 +8,7 @@ from apps.webui.routers import (
|
||||
users,
|
||||
chats,
|
||||
documents,
|
||||
tools,
|
||||
models,
|
||||
prompts,
|
||||
configs,
|
||||
@@ -27,9 +28,9 @@ from config import (
|
||||
WEBHOOK_URL,
|
||||
WEBUI_AUTH_TRUSTED_EMAIL_HEADER,
|
||||
JWT_EXPIRES_IN,
|
||||
AppConfig,
|
||||
ENABLE_COMMUNITY_SHARING,
|
||||
WEBUI_BANNERS,
|
||||
ENABLE_COMMUNITY_SHARING,
|
||||
AppConfig,
|
||||
)
|
||||
|
||||
app = FastAPI()
|
||||
@@ -40,6 +41,7 @@ app.state.config = AppConfig()
|
||||
|
||||
app.state.config.ENABLE_SIGNUP = ENABLE_SIGNUP
|
||||
app.state.config.JWT_EXPIRES_IN = JWT_EXPIRES_IN
|
||||
app.state.AUTH_TRUSTED_EMAIL_HEADER = WEBUI_AUTH_TRUSTED_EMAIL_HEADER
|
||||
|
||||
|
||||
app.state.config.SHOW_ADMIN_DETAILS = SHOW_ADMIN_DETAILS
|
||||
@@ -56,7 +58,7 @@ app.state.config.BANNERS = WEBUI_BANNERS
|
||||
app.state.config.ENABLE_COMMUNITY_SHARING = ENABLE_COMMUNITY_SHARING
|
||||
|
||||
app.state.MODELS = {}
|
||||
app.state.AUTH_TRUSTED_EMAIL_HEADER = WEBUI_AUTH_TRUSTED_EMAIL_HEADER
|
||||
app.state.TOOLS = {}
|
||||
|
||||
|
||||
app.add_middleware(
|
||||
@@ -72,6 +74,7 @@ app.include_router(users.router, prefix="/users", tags=["users"])
|
||||
app.include_router(chats.router, prefix="/chats", tags=["chats"])
|
||||
|
||||
app.include_router(documents.router, prefix="/documents", tags=["documents"])
|
||||
app.include_router(tools.router, prefix="/tools", tags=["tools"])
|
||||
app.include_router(models.router, prefix="/models", tags=["models"])
|
||||
app.include_router(prompts.router, prefix="/prompts", tags=["prompts"])
|
||||
app.include_router(memories.router, prefix="/memories", tags=["memories"])
|
||||
|
||||
@@ -65,6 +65,20 @@ class MemoriesTable:
|
||||
else:
|
||||
return None
|
||||
|
||||
def update_memory_by_id(
|
||||
self,
|
||||
id: str,
|
||||
content: str,
|
||||
) -> Optional[MemoryModel]:
|
||||
try:
|
||||
memory = Memory.get(Memory.id == id)
|
||||
memory.content = content
|
||||
memory.updated_at = int(time.time())
|
||||
memory.save()
|
||||
return MemoryModel(**model_to_dict(memory))
|
||||
except:
|
||||
return None
|
||||
|
||||
def get_memories(self) -> List[MemoryModel]:
|
||||
try:
|
||||
memories = Memory.select()
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
from pydantic import BaseModel
|
||||
from peewee import *
|
||||
from playhouse.shortcuts import model_to_dict
|
||||
from typing import List, Union, Optional
|
||||
import time
|
||||
import logging
|
||||
from apps.webui.internal.db import DB, JSONField
|
||||
|
||||
import json
|
||||
|
||||
from config import SRC_LOG_LEVELS
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["MODELS"])
|
||||
|
||||
####################
|
||||
# Tools DB Schema
|
||||
####################
|
||||
|
||||
|
||||
class Tool(Model):
|
||||
id = CharField(unique=True)
|
||||
user_id = CharField()
|
||||
name = TextField()
|
||||
content = TextField()
|
||||
specs = JSONField()
|
||||
meta = JSONField()
|
||||
updated_at = BigIntegerField()
|
||||
created_at = BigIntegerField()
|
||||
|
||||
class Meta:
|
||||
database = DB
|
||||
|
||||
|
||||
class ToolMeta(BaseModel):
|
||||
description: Optional[str] = None
|
||||
|
||||
|
||||
class ToolModel(BaseModel):
|
||||
id: str
|
||||
user_id: str
|
||||
name: str
|
||||
content: str
|
||||
specs: List[dict]
|
||||
meta: ToolMeta
|
||||
updated_at: int # timestamp in epoch
|
||||
created_at: int # timestamp in epoch
|
||||
|
||||
|
||||
####################
|
||||
# Forms
|
||||
####################
|
||||
|
||||
|
||||
class ToolResponse(BaseModel):
|
||||
id: str
|
||||
user_id: str
|
||||
name: str
|
||||
meta: ToolMeta
|
||||
updated_at: int # timestamp in epoch
|
||||
created_at: int # timestamp in epoch
|
||||
|
||||
|
||||
class ToolForm(BaseModel):
|
||||
id: str
|
||||
name: str
|
||||
content: str
|
||||
meta: ToolMeta
|
||||
|
||||
|
||||
class ToolsTable:
|
||||
def __init__(self, db):
|
||||
self.db = db
|
||||
self.db.create_tables([Tool])
|
||||
|
||||
def insert_new_tool(
|
||||
self, user_id: str, form_data: ToolForm, specs: List[dict]
|
||||
) -> Optional[ToolModel]:
|
||||
tool = ToolModel(
|
||||
**{
|
||||
**form_data.model_dump(),
|
||||
"specs": specs,
|
||||
"user_id": user_id,
|
||||
"updated_at": int(time.time()),
|
||||
"created_at": int(time.time()),
|
||||
}
|
||||
)
|
||||
|
||||
try:
|
||||
result = Tool.create(**tool.model_dump())
|
||||
if result:
|
||||
return tool
|
||||
else:
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f"Error creating tool: {e}")
|
||||
return None
|
||||
|
||||
def get_tool_by_id(self, id: str) -> Optional[ToolModel]:
|
||||
try:
|
||||
tool = Tool.get(Tool.id == id)
|
||||
return ToolModel(**model_to_dict(tool))
|
||||
except:
|
||||
return None
|
||||
|
||||
def get_tools(self) -> List[ToolModel]:
|
||||
return [ToolModel(**model_to_dict(tool)) for tool in Tool.select()]
|
||||
|
||||
def update_tool_by_id(self, id: str, updated: dict) -> Optional[ToolModel]:
|
||||
try:
|
||||
query = Tool.update(
|
||||
**updated,
|
||||
updated_at=int(time.time()),
|
||||
).where(Tool.id == id)
|
||||
query.execute()
|
||||
|
||||
tool = Tool.get(Tool.id == id)
|
||||
return ToolModel(**model_to_dict(tool))
|
||||
except:
|
||||
return None
|
||||
|
||||
def delete_tool_by_id(self, id: str) -> bool:
|
||||
try:
|
||||
query = Tool.delete().where((Tool.id == id))
|
||||
query.execute() # Remove the rows, return number of rows removed.
|
||||
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
|
||||
|
||||
Tools = ToolsTable(DB)
|
||||
@@ -161,7 +161,7 @@ async def get_archived_session_user_chat_list(
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/archive/all", response_model=List[ChatTitleIdResponse])
|
||||
@router.post("/archive/all", response_model=bool)
|
||||
async def archive_all_chats(user=Depends(get_current_user)):
|
||||
return Chats.archive_all_chats_by_user_id(user.id)
|
||||
|
||||
|
||||
@@ -73,7 +73,7 @@ async def create_new_doc(form_data: DocumentForm, user=Depends(get_admin_user)):
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/name/{name}", response_model=Optional[DocumentResponse])
|
||||
@router.get("/doc", response_model=Optional[DocumentResponse])
|
||||
async def get_doc_by_name(name: str, user=Depends(get_current_user)):
|
||||
doc = Documents.get_doc_by_name(name)
|
||||
|
||||
@@ -105,7 +105,7 @@ class TagDocumentForm(BaseModel):
|
||||
tags: List[dict]
|
||||
|
||||
|
||||
@router.post("/name/{name}/tags", response_model=Optional[DocumentResponse])
|
||||
@router.post("/doc/tags", response_model=Optional[DocumentResponse])
|
||||
async def tag_doc_by_name(form_data: TagDocumentForm, user=Depends(get_current_user)):
|
||||
doc = Documents.update_doc_content_by_name(form_data.name, {"tags": form_data.tags})
|
||||
|
||||
@@ -128,7 +128,7 @@ async def tag_doc_by_name(form_data: TagDocumentForm, user=Depends(get_current_u
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/name/{name}/update", response_model=Optional[DocumentResponse])
|
||||
@router.post("/doc/update", response_model=Optional[DocumentResponse])
|
||||
async def update_doc_by_name(
|
||||
name: str, form_data: DocumentUpdateForm, user=Depends(get_admin_user)
|
||||
):
|
||||
@@ -152,7 +152,7 @@ async def update_doc_by_name(
|
||||
############################
|
||||
|
||||
|
||||
@router.delete("/name/{name}/delete", response_model=bool)
|
||||
@router.delete("/doc/delete", response_model=bool)
|
||||
async def delete_doc_by_name(name: str, user=Depends(get_admin_user)):
|
||||
result = Documents.delete_doc_by_name(name)
|
||||
return result
|
||||
|
||||
@@ -44,6 +44,10 @@ class AddMemoryForm(BaseModel):
|
||||
content: str
|
||||
|
||||
|
||||
class MemoryUpdateModel(BaseModel):
|
||||
content: Optional[str] = None
|
||||
|
||||
|
||||
@router.post("/add", response_model=Optional[MemoryModel])
|
||||
async def add_memory(
|
||||
request: Request, form_data: AddMemoryForm, user=Depends(get_verified_user)
|
||||
@@ -62,6 +66,34 @@ async def add_memory(
|
||||
return memory
|
||||
|
||||
|
||||
@router.post("/{memory_id}/update", response_model=Optional[MemoryModel])
|
||||
async def update_memory_by_id(
|
||||
memory_id: str,
|
||||
request: Request,
|
||||
form_data: MemoryUpdateModel,
|
||||
user=Depends(get_verified_user),
|
||||
):
|
||||
memory = Memories.update_memory_by_id(memory_id, form_data.content)
|
||||
if memory is None:
|
||||
raise HTTPException(status_code=404, detail="Memory not found")
|
||||
|
||||
if form_data.content is not None:
|
||||
memory_embedding = request.app.state.EMBEDDING_FUNCTION(form_data.content)
|
||||
collection = CHROMA_CLIENT.get_or_create_collection(
|
||||
name=f"user-memory-{user.id}"
|
||||
)
|
||||
collection.upsert(
|
||||
documents=[form_data.content],
|
||||
ids=[memory.id],
|
||||
embeddings=[memory_embedding],
|
||||
metadatas=[
|
||||
{"created_at": memory.created_at, "updated_at": memory.updated_at}
|
||||
],
|
||||
)
|
||||
|
||||
return memory
|
||||
|
||||
|
||||
############################
|
||||
# QueryMemory
|
||||
############################
|
||||
|
||||
@@ -0,0 +1,183 @@
|
||||
from fastapi import Depends, FastAPI, HTTPException, status, Request
|
||||
from datetime import datetime, timedelta
|
||||
from typing import List, Union, Optional
|
||||
|
||||
from fastapi import APIRouter
|
||||
from pydantic import BaseModel
|
||||
import json
|
||||
|
||||
from apps.webui.models.tools import Tools, ToolForm, ToolModel, ToolResponse
|
||||
from apps.webui.utils import load_toolkit_module_by_id
|
||||
|
||||
from utils.utils import get_current_user, get_admin_user
|
||||
from utils.tools import get_tools_specs
|
||||
from constants import ERROR_MESSAGES
|
||||
|
||||
from importlib import util
|
||||
import os
|
||||
|
||||
from config import DATA_DIR
|
||||
|
||||
|
||||
TOOLS_DIR = f"{DATA_DIR}/tools"
|
||||
os.makedirs(TOOLS_DIR, exist_ok=True)
|
||||
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
############################
|
||||
# GetToolkits
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/", response_model=List[ToolResponse])
|
||||
async def get_toolkits(user=Depends(get_current_user)):
|
||||
toolkits = [toolkit for toolkit in Tools.get_tools()]
|
||||
return toolkits
|
||||
|
||||
|
||||
############################
|
||||
# ExportToolKits
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/export", response_model=List[ToolModel])
|
||||
async def get_toolkits(user=Depends(get_admin_user)):
|
||||
toolkits = [toolkit for toolkit in Tools.get_tools()]
|
||||
return toolkits
|
||||
|
||||
|
||||
############################
|
||||
# CreateNewToolKit
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/create", response_model=Optional[ToolResponse])
|
||||
async def create_new_toolkit(
|
||||
request: Request, form_data: ToolForm, user=Depends(get_admin_user)
|
||||
):
|
||||
if not form_data.id.isidentifier():
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Only alphanumeric characters and underscores are allowed in the id",
|
||||
)
|
||||
|
||||
form_data.id = form_data.id.lower()
|
||||
|
||||
toolkit = Tools.get_tool_by_id(form_data.id)
|
||||
if toolkit == None:
|
||||
toolkit_path = os.path.join(TOOLS_DIR, f"{form_data.id}.py")
|
||||
try:
|
||||
with open(toolkit_path, "w") as tool_file:
|
||||
tool_file.write(form_data.content)
|
||||
|
||||
toolkit_module = load_toolkit_module_by_id(form_data.id)
|
||||
|
||||
TOOLS = request.app.state.TOOLS
|
||||
TOOLS[form_data.id] = toolkit_module
|
||||
|
||||
specs = get_tools_specs(TOOLS[form_data.id])
|
||||
toolkit = Tools.insert_new_tool(user.id, form_data, specs)
|
||||
|
||||
if toolkit:
|
||||
return toolkit
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error creating toolkit"),
|
||||
)
|
||||
except Exception as e:
|
||||
print(e)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(e),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.ID_TAKEN,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# GetToolkitById
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/id/{id}", response_model=Optional[ToolModel])
|
||||
async def get_toolkit_by_id(id: str, user=Depends(get_admin_user)):
|
||||
toolkit = Tools.get_tool_by_id(id)
|
||||
|
||||
if toolkit:
|
||||
return toolkit
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# UpdateToolkitById
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/id/{id}/update", response_model=Optional[ToolModel])
|
||||
async def update_toolkit_by_id(
|
||||
request: Request, id: str, form_data: ToolForm, user=Depends(get_admin_user)
|
||||
):
|
||||
toolkit_path = os.path.join(TOOLS_DIR, f"{id}.py")
|
||||
|
||||
try:
|
||||
with open(toolkit_path, "w") as tool_file:
|
||||
tool_file.write(form_data.content)
|
||||
|
||||
toolkit_module = load_toolkit_module_by_id(id)
|
||||
|
||||
TOOLS = request.app.state.TOOLS
|
||||
TOOLS[id] = toolkit_module
|
||||
|
||||
specs = get_tools_specs(TOOLS[id])
|
||||
|
||||
updated = {
|
||||
**form_data.model_dump(exclude={"id"}),
|
||||
"specs": specs,
|
||||
}
|
||||
|
||||
print(updated)
|
||||
toolkit = Tools.update_tool_by_id(id, updated)
|
||||
|
||||
if toolkit:
|
||||
return toolkit
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error updating toolkit"),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT(e),
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# DeleteToolkitById
|
||||
############################
|
||||
|
||||
|
||||
@router.delete("/id/{id}/delete", response_model=bool)
|
||||
async def delete_toolkit_by_id(request: Request, id: str, user=Depends(get_admin_user)):
|
||||
result = Tools.delete_tool_by_id(id)
|
||||
|
||||
if result:
|
||||
TOOLS = request.app.state.TOOLS
|
||||
if id in TOOLS:
|
||||
del TOOLS[id]
|
||||
|
||||
# delete the toolkit file
|
||||
toolkit_path = os.path.join(TOOLS_DIR, f"{id}.py")
|
||||
os.remove(toolkit_path)
|
||||
|
||||
return result
|
||||
@@ -7,6 +7,8 @@ from pydantic import BaseModel
|
||||
|
||||
from fpdf import FPDF
|
||||
import markdown
|
||||
import black
|
||||
|
||||
|
||||
from apps.webui.internal.db import DB
|
||||
from utils.utils import get_admin_user
|
||||
@@ -26,6 +28,21 @@ async def get_gravatar(
|
||||
return get_gravatar_url(email)
|
||||
|
||||
|
||||
class CodeFormatRequest(BaseModel):
|
||||
code: str
|
||||
|
||||
|
||||
@router.post("/code/format")
|
||||
async def format_code(request: CodeFormatRequest):
|
||||
try:
|
||||
formatted_code = black.format_str(request.code, mode=black.Mode())
|
||||
return {"code": formatted_code}
|
||||
except black.NothingChanged:
|
||||
return {"code": request.code}
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
|
||||
|
||||
class MarkdownForm(BaseModel):
|
||||
md: str
|
||||
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
from importlib import util
|
||||
import os
|
||||
|
||||
from config import TOOLS_DIR
|
||||
|
||||
|
||||
def load_toolkit_module_by_id(toolkit_id):
|
||||
toolkit_path = os.path.join(TOOLS_DIR, f"{toolkit_id}.py")
|
||||
spec = util.spec_from_file_location(toolkit_id, toolkit_path)
|
||||
module = util.module_from_spec(spec)
|
||||
|
||||
try:
|
||||
spec.loader.exec_module(module)
|
||||
print(f"Loaded module: {module.__name__}")
|
||||
if hasattr(module, "Tools"):
|
||||
return module.Tools()
|
||||
else:
|
||||
raise Exception("No Tools class found")
|
||||
except Exception as e:
|
||||
print(f"Error loading module: {toolkit_id}")
|
||||
# Move the file to the error folder
|
||||
os.rename(toolkit_path, f"{toolkit_path}.error")
|
||||
raise e
|
||||
Reference in New Issue
Block a user