main.py 52 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661
  1. import uuid
  2. from contextlib import asynccontextmanager
  3. from authlib.integrations.starlette_client import OAuth
  4. from authlib.oidc.core import UserInfo
  5. from bs4 import BeautifulSoup
  6. import json
  7. import markdown
  8. import time
  9. import os
  10. import sys
  11. import logging
  12. import aiohttp
  13. import requests
  14. import mimetypes
  15. import shutil
  16. import os
  17. import inspect
  18. import asyncio
  19. from fastapi import FastAPI, Request, Depends, status, UploadFile, File, Form
  20. from fastapi.staticfiles import StaticFiles
  21. from fastapi.responses import JSONResponse
  22. from fastapi import HTTPException
  23. from fastapi.middleware.wsgi import WSGIMiddleware
  24. from fastapi.middleware.cors import CORSMiddleware
  25. from starlette.exceptions import HTTPException as StarletteHTTPException
  26. from starlette.middleware.base import BaseHTTPMiddleware
  27. from starlette.middleware.sessions import SessionMiddleware
  28. from starlette.responses import StreamingResponse, Response, RedirectResponse
  29. from apps.socket.main import app as socket_app
  30. from apps.ollama.main import (
  31. app as ollama_app,
  32. OpenAIChatCompletionForm,
  33. get_all_models as get_ollama_models,
  34. generate_openai_chat_completion as generate_ollama_chat_completion,
  35. )
  36. from apps.openai.main import (
  37. app as openai_app,
  38. get_all_models as get_openai_models,
  39. generate_chat_completion as generate_openai_chat_completion,
  40. )
  41. from apps.audio.main import app as audio_app
  42. from apps.images.main import app as images_app
  43. from apps.rag.main import app as rag_app
  44. from apps.webui.main import app as webui_app
  45. from pydantic import BaseModel
  46. from typing import List, Optional
  47. from apps.webui.models.auths import Auths
  48. from apps.webui.models.models import Models, ModelModel
  49. from apps.webui.models.tools import Tools
  50. from apps.webui.models.users import Users
  51. from apps.webui.utils import load_toolkit_module_by_id
  52. from utils.misc import parse_duration
  53. from utils.utils import (
  54. get_admin_user,
  55. get_verified_user,
  56. get_current_user,
  57. get_http_authorization_cred,
  58. get_password_hash,
  59. create_token,
  60. )
  61. from utils.task import (
  62. title_generation_template,
  63. search_query_generation_template,
  64. tools_function_calling_generation_template,
  65. )
  66. from utils.misc import get_last_user_message, add_or_update_system_message
  67. from apps.rag.utils import get_rag_context, rag_template
  68. from config import (
  69. CONFIG_DATA,
  70. WEBUI_NAME,
  71. WEBUI_URL,
  72. WEBUI_AUTH,
  73. ENV,
  74. VERSION,
  75. CHANGELOG,
  76. FRONTEND_BUILD_DIR,
  77. CACHE_DIR,
  78. STATIC_DIR,
  79. ENABLE_OPENAI_API,
  80. ENABLE_OLLAMA_API,
  81. ENABLE_MODEL_FILTER,
  82. MODEL_FILTER_LIST,
  83. GLOBAL_LOG_LEVEL,
  84. SRC_LOG_LEVELS,
  85. WEBHOOK_URL,
  86. ENABLE_ADMIN_EXPORT,
  87. WEBUI_BUILD_HASH,
  88. TASK_MODEL,
  89. TASK_MODEL_EXTERNAL,
  90. TITLE_GENERATION_PROMPT_TEMPLATE,
  91. SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE,
  92. SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD,
  93. TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE,
  94. OAUTH_PROVIDERS,
  95. ENABLE_OAUTH_SIGNUP,
  96. OAUTH_MERGE_ACCOUNTS_BY_EMAIL,
  97. WEBUI_SECRET_KEY,
  98. WEBUI_SESSION_COOKIE_SAME_SITE,
  99. WEBUI_SESSION_COOKIE_SECURE,
  100. AppConfig,
  101. )
  102. from constants import ERROR_MESSAGES, WEBHOOK_MESSAGES
  103. from utils.webhook import post_webhook
  104. logging.basicConfig(stream=sys.stdout, level=GLOBAL_LOG_LEVEL)
  105. log = logging.getLogger(__name__)
  106. log.setLevel(SRC_LOG_LEVELS["MAIN"])
  107. class SPAStaticFiles(StaticFiles):
  108. async def get_response(self, path: str, scope):
  109. try:
  110. return await super().get_response(path, scope)
  111. except (HTTPException, StarletteHTTPException) as ex:
  112. if ex.status_code == 404:
  113. return await super().get_response("index.html", scope)
  114. else:
  115. raise ex
  116. print(
  117. rf"""
  118. ___ __ __ _ _ _ ___
  119. / _ \ _ __ ___ _ __ \ \ / /__| |__ | | | |_ _|
  120. | | | | '_ \ / _ \ '_ \ \ \ /\ / / _ \ '_ \| | | || |
  121. | |_| | |_) | __/ | | | \ V V / __/ |_) | |_| || |
  122. \___/| .__/ \___|_| |_| \_/\_/ \___|_.__/ \___/|___|
  123. |_|
  124. v{VERSION} - building the best open-source AI user interface.
  125. {f"Commit: {WEBUI_BUILD_HASH}" if WEBUI_BUILD_HASH != "dev-build" else ""}
  126. https://github.com/open-webui/open-webui
  127. """
  128. )
  129. @asynccontextmanager
  130. async def lifespan(app: FastAPI):
  131. yield
  132. app = FastAPI(
  133. docs_url="/docs" if ENV == "dev" else None, redoc_url=None, lifespan=lifespan
  134. )
  135. app.state.config = AppConfig()
  136. app.state.config.ENABLE_OPENAI_API = ENABLE_OPENAI_API
  137. app.state.config.ENABLE_OLLAMA_API = ENABLE_OLLAMA_API
  138. app.state.config.ENABLE_MODEL_FILTER = ENABLE_MODEL_FILTER
  139. app.state.config.MODEL_FILTER_LIST = MODEL_FILTER_LIST
  140. app.state.config.WEBHOOK_URL = WEBHOOK_URL
  141. app.state.config.TASK_MODEL = TASK_MODEL
  142. app.state.config.TASK_MODEL_EXTERNAL = TASK_MODEL_EXTERNAL
  143. app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE = TITLE_GENERATION_PROMPT_TEMPLATE
  144. app.state.config.SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE = (
  145. SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE
  146. )
  147. app.state.config.SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD = (
  148. SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD
  149. )
  150. app.state.config.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE = (
  151. TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE
  152. )
  153. app.state.MODELS = {}
  154. origins = ["*"]
  155. async def get_function_call_response(messages, tool_id, template, task_model_id, user):
  156. tool = Tools.get_tool_by_id(tool_id)
  157. tools_specs = json.dumps(tool.specs, indent=2)
  158. content = tools_function_calling_generation_template(template, tools_specs)
  159. user_message = get_last_user_message(messages)
  160. prompt = (
  161. "History:\n"
  162. + "\n".join(
  163. [
  164. f"{message['role'].upper()}: \"\"\"{message['content']}\"\"\""
  165. for message in messages[::-1][:4]
  166. ]
  167. )
  168. + f"\nQuery: {user_message}"
  169. )
  170. print(prompt)
  171. payload = {
  172. "model": task_model_id,
  173. "messages": [
  174. {"role": "system", "content": content},
  175. {"role": "user", "content": f"Query: {prompt}"},
  176. ],
  177. "stream": False,
  178. }
  179. try:
  180. payload = filter_pipeline(payload, user)
  181. except Exception as e:
  182. raise e
  183. model = app.state.MODELS[task_model_id]
  184. response = None
  185. try:
  186. if model["owned_by"] == "ollama":
  187. response = await generate_ollama_chat_completion(
  188. OpenAIChatCompletionForm(**payload), user=user
  189. )
  190. else:
  191. response = await generate_openai_chat_completion(payload, user=user)
  192. content = None
  193. if hasattr(response, "body_iterator"):
  194. async for chunk in response.body_iterator:
  195. data = json.loads(chunk.decode("utf-8"))
  196. content = data["choices"][0]["message"]["content"]
  197. # Cleanup any remaining background tasks if necessary
  198. if response.background is not None:
  199. await response.background()
  200. else:
  201. content = response["choices"][0]["message"]["content"]
  202. # Parse the function response
  203. if content is not None:
  204. print(f"content: {content}")
  205. result = json.loads(content)
  206. print(result)
  207. # Call the function
  208. if "name" in result:
  209. if tool_id in webui_app.state.TOOLS:
  210. toolkit_module = webui_app.state.TOOLS[tool_id]
  211. else:
  212. toolkit_module = load_toolkit_module_by_id(tool_id)
  213. webui_app.state.TOOLS[tool_id] = toolkit_module
  214. function = getattr(toolkit_module, result["name"])
  215. function_result = None
  216. try:
  217. # Get the signature of the function
  218. sig = inspect.signature(function)
  219. # Check if '__user__' is a parameter of the function
  220. if "__user__" in sig.parameters:
  221. # Call the function with the '__user__' parameter included
  222. function_result = function(
  223. **{
  224. **result["parameters"],
  225. "__user__": {
  226. "id": user.id,
  227. "email": user.email,
  228. "name": user.name,
  229. "role": user.role,
  230. },
  231. }
  232. )
  233. else:
  234. # Call the function without modifying the parameters
  235. function_result = function(**result["parameters"])
  236. except Exception as e:
  237. print(e)
  238. # Add the function result to the system prompt
  239. if function_result:
  240. return function_result
  241. except Exception as e:
  242. print(f"Error: {e}")
  243. return None
  244. class ChatCompletionMiddleware(BaseHTTPMiddleware):
  245. async def dispatch(self, request: Request, call_next):
  246. return_citations = False
  247. if request.method == "POST" and (
  248. "/ollama/api/chat" in request.url.path
  249. or "/chat/completions" in request.url.path
  250. ):
  251. log.debug(f"request.url.path: {request.url.path}")
  252. # Read the original request body
  253. body = await request.body()
  254. # Decode body to string
  255. body_str = body.decode("utf-8")
  256. # Parse string to JSON
  257. data = json.loads(body_str) if body_str else {}
  258. user = get_current_user(
  259. get_http_authorization_cred(request.headers.get("Authorization"))
  260. )
  261. # Remove the citations from the body
  262. return_citations = data.get("citations", False)
  263. if "citations" in data:
  264. del data["citations"]
  265. # Set the task model
  266. task_model_id = data["model"]
  267. if task_model_id not in app.state.MODELS:
  268. raise HTTPException(
  269. status_code=status.HTTP_404_NOT_FOUND,
  270. detail="Model not found",
  271. )
  272. # Check if the user has a custom task model
  273. # If the user has a custom task model, use that model
  274. if app.state.MODELS[task_model_id]["owned_by"] == "ollama":
  275. if (
  276. app.state.config.TASK_MODEL
  277. and app.state.config.TASK_MODEL in app.state.MODELS
  278. ):
  279. task_model_id = app.state.config.TASK_MODEL
  280. else:
  281. if (
  282. app.state.config.TASK_MODEL_EXTERNAL
  283. and app.state.config.TASK_MODEL_EXTERNAL in app.state.MODELS
  284. ):
  285. task_model_id = app.state.config.TASK_MODEL_EXTERNAL
  286. prompt = get_last_user_message(data["messages"])
  287. context = ""
  288. # If tool_ids field is present, call the functions
  289. if "tool_ids" in data:
  290. print(data["tool_ids"])
  291. for tool_id in data["tool_ids"]:
  292. print(tool_id)
  293. try:
  294. response = await get_function_call_response(
  295. messages=data["messages"],
  296. tool_id=tool_id,
  297. template=app.state.config.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE,
  298. task_model_id=task_model_id,
  299. user=user,
  300. )
  301. if response:
  302. context += ("\n" if context != "" else "") + response
  303. except Exception as e:
  304. print(f"Error: {e}")
  305. del data["tool_ids"]
  306. print(f"tool_context: {context}")
  307. # If docs field is present, generate RAG completions
  308. if "docs" in data:
  309. data = {**data}
  310. rag_context, citations = get_rag_context(
  311. docs=data["docs"],
  312. messages=data["messages"],
  313. embedding_function=rag_app.state.EMBEDDING_FUNCTION,
  314. k=rag_app.state.config.TOP_K,
  315. reranking_function=rag_app.state.sentence_transformer_rf,
  316. r=rag_app.state.config.RELEVANCE_THRESHOLD,
  317. hybrid_search=rag_app.state.config.ENABLE_RAG_HYBRID_SEARCH,
  318. )
  319. if rag_context:
  320. context += ("\n" if context != "" else "") + rag_context
  321. del data["docs"]
  322. log.debug(f"rag_context: {rag_context}, citations: {citations}")
  323. if context != "":
  324. system_prompt = rag_template(
  325. rag_app.state.config.RAG_TEMPLATE, context, prompt
  326. )
  327. print(system_prompt)
  328. data["messages"] = add_or_update_system_message(
  329. f"\n{system_prompt}", data["messages"]
  330. )
  331. modified_body_bytes = json.dumps(data).encode("utf-8")
  332. # Replace the request body with the modified one
  333. request._body = modified_body_bytes
  334. # Set custom header to ensure content-length matches new body length
  335. request.headers.__dict__["_list"] = [
  336. (b"content-length", str(len(modified_body_bytes)).encode("utf-8")),
  337. *[
  338. (k, v)
  339. for k, v in request.headers.raw
  340. if k.lower() != b"content-length"
  341. ],
  342. ]
  343. response = await call_next(request)
  344. if return_citations:
  345. # Inject the citations into the response
  346. if isinstance(response, StreamingResponse):
  347. # If it's a streaming response, inject it as SSE event or NDJSON line
  348. content_type = response.headers.get("Content-Type")
  349. if "text/event-stream" in content_type:
  350. return StreamingResponse(
  351. self.openai_stream_wrapper(response.body_iterator, citations),
  352. )
  353. if "application/x-ndjson" in content_type:
  354. return StreamingResponse(
  355. self.ollama_stream_wrapper(response.body_iterator, citations),
  356. )
  357. return response
  358. async def _receive(self, body: bytes):
  359. return {"type": "http.request", "body": body, "more_body": False}
  360. async def openai_stream_wrapper(self, original_generator, citations):
  361. yield f"data: {json.dumps({'citations': citations})}\n\n"
  362. async for data in original_generator:
  363. yield data
  364. async def ollama_stream_wrapper(self, original_generator, citations):
  365. yield f"{json.dumps({'citations': citations})}\n"
  366. async for data in original_generator:
  367. yield data
  368. app.add_middleware(ChatCompletionMiddleware)
  369. def filter_pipeline(payload, user):
  370. user = {"id": user.id, "name": user.name, "role": user.role}
  371. model_id = payload["model"]
  372. filters = [
  373. model
  374. for model in app.state.MODELS.values()
  375. if "pipeline" in model
  376. and "type" in model["pipeline"]
  377. and model["pipeline"]["type"] == "filter"
  378. and (
  379. model["pipeline"]["pipelines"] == ["*"]
  380. or any(
  381. model_id == target_model_id
  382. for target_model_id in model["pipeline"]["pipelines"]
  383. )
  384. )
  385. ]
  386. sorted_filters = sorted(filters, key=lambda x: x["pipeline"]["priority"])
  387. model = app.state.MODELS[model_id]
  388. if "pipeline" in model:
  389. sorted_filters.append(model)
  390. for filter in sorted_filters:
  391. r = None
  392. try:
  393. urlIdx = filter["urlIdx"]
  394. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  395. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  396. if key != "":
  397. headers = {"Authorization": f"Bearer {key}"}
  398. r = requests.post(
  399. f"{url}/{filter['id']}/filter/inlet",
  400. headers=headers,
  401. json={
  402. "user": user,
  403. "body": payload,
  404. },
  405. )
  406. r.raise_for_status()
  407. payload = r.json()
  408. except Exception as e:
  409. # Handle connection error here
  410. print(f"Connection error: {e}")
  411. if r is not None:
  412. try:
  413. res = r.json()
  414. except:
  415. pass
  416. if "detail" in res:
  417. raise Exception(r.status_code, res["detail"])
  418. else:
  419. pass
  420. if "pipeline" not in app.state.MODELS[model_id]:
  421. if "chat_id" in payload:
  422. del payload["chat_id"]
  423. if "title" in payload:
  424. del payload["title"]
  425. if "task" in payload:
  426. del payload["task"]
  427. return payload
  428. class PipelineMiddleware(BaseHTTPMiddleware):
  429. async def dispatch(self, request: Request, call_next):
  430. if request.method == "POST" and (
  431. "/ollama/api/chat" in request.url.path
  432. or "/chat/completions" in request.url.path
  433. ):
  434. log.debug(f"request.url.path: {request.url.path}")
  435. # Read the original request body
  436. body = await request.body()
  437. # Decode body to string
  438. body_str = body.decode("utf-8")
  439. # Parse string to JSON
  440. data = json.loads(body_str) if body_str else {}
  441. user = get_current_user(
  442. get_http_authorization_cred(request.headers.get("Authorization"))
  443. )
  444. try:
  445. data = filter_pipeline(data, user)
  446. except Exception as e:
  447. return JSONResponse(
  448. status_code=e.args[0],
  449. content={"detail": e.args[1]},
  450. )
  451. modified_body_bytes = json.dumps(data).encode("utf-8")
  452. # Replace the request body with the modified one
  453. request._body = modified_body_bytes
  454. # Set custom header to ensure content-length matches new body length
  455. request.headers.__dict__["_list"] = [
  456. (b"content-length", str(len(modified_body_bytes)).encode("utf-8")),
  457. *[
  458. (k, v)
  459. for k, v in request.headers.raw
  460. if k.lower() != b"content-length"
  461. ],
  462. ]
  463. response = await call_next(request)
  464. return response
  465. async def _receive(self, body: bytes):
  466. return {"type": "http.request", "body": body, "more_body": False}
  467. app.add_middleware(PipelineMiddleware)
  468. app.add_middleware(
  469. CORSMiddleware,
  470. allow_origins=origins,
  471. allow_credentials=True,
  472. allow_methods=["*"],
  473. allow_headers=["*"],
  474. )
  475. @app.middleware("http")
  476. async def check_url(request: Request, call_next):
  477. if len(app.state.MODELS) == 0:
  478. await get_all_models()
  479. else:
  480. pass
  481. start_time = int(time.time())
  482. response = await call_next(request)
  483. process_time = int(time.time()) - start_time
  484. response.headers["X-Process-Time"] = str(process_time)
  485. return response
  486. @app.middleware("http")
  487. async def update_embedding_function(request: Request, call_next):
  488. response = await call_next(request)
  489. if "/embedding/update" in request.url.path:
  490. webui_app.state.EMBEDDING_FUNCTION = rag_app.state.EMBEDDING_FUNCTION
  491. return response
  492. app.mount("/ws", socket_app)
  493. app.mount("/ollama", ollama_app)
  494. app.mount("/openai", openai_app)
  495. app.mount("/images/api/v1", images_app)
  496. app.mount("/audio/api/v1", audio_app)
  497. app.mount("/rag/api/v1", rag_app)
  498. app.mount("/api/v1", webui_app)
  499. webui_app.state.EMBEDDING_FUNCTION = rag_app.state.EMBEDDING_FUNCTION
  500. async def get_all_models():
  501. openai_models = []
  502. ollama_models = []
  503. if app.state.config.ENABLE_OPENAI_API:
  504. openai_models = await get_openai_models()
  505. openai_models = openai_models["data"]
  506. if app.state.config.ENABLE_OLLAMA_API:
  507. ollama_models = await get_ollama_models()
  508. ollama_models = [
  509. {
  510. "id": model["model"],
  511. "name": model["name"],
  512. "object": "model",
  513. "created": int(time.time()),
  514. "owned_by": "ollama",
  515. "ollama": model,
  516. }
  517. for model in ollama_models["models"]
  518. ]
  519. models = openai_models + ollama_models
  520. custom_models = Models.get_all_models()
  521. for custom_model in custom_models:
  522. if custom_model.base_model_id == None:
  523. for model in models:
  524. if (
  525. custom_model.id == model["id"]
  526. or custom_model.id == model["id"].split(":")[0]
  527. ):
  528. model["name"] = custom_model.name
  529. model["info"] = custom_model.model_dump()
  530. else:
  531. owned_by = "openai"
  532. for model in models:
  533. if (
  534. custom_model.base_model_id == model["id"]
  535. or custom_model.base_model_id == model["id"].split(":")[0]
  536. ):
  537. owned_by = model["owned_by"]
  538. break
  539. models.append(
  540. {
  541. "id": custom_model.id,
  542. "name": custom_model.name,
  543. "object": "model",
  544. "created": custom_model.created_at,
  545. "owned_by": owned_by,
  546. "info": custom_model.model_dump(),
  547. "preset": True,
  548. }
  549. )
  550. app.state.MODELS = {model["id"]: model for model in models}
  551. webui_app.state.MODELS = app.state.MODELS
  552. return models
  553. @app.get("/api/models")
  554. async def get_models(user=Depends(get_verified_user)):
  555. models = await get_all_models()
  556. # Filter out filter pipelines
  557. models = [
  558. model
  559. for model in models
  560. if "pipeline" not in model or model["pipeline"].get("type", None) != "filter"
  561. ]
  562. if app.state.config.ENABLE_MODEL_FILTER:
  563. if user.role == "user":
  564. models = list(
  565. filter(
  566. lambda model: model["id"] in app.state.config.MODEL_FILTER_LIST,
  567. models,
  568. )
  569. )
  570. return {"data": models}
  571. return {"data": models}
  572. @app.get("/api/task/config")
  573. async def get_task_config(user=Depends(get_verified_user)):
  574. return {
  575. "TASK_MODEL": app.state.config.TASK_MODEL,
  576. "TASK_MODEL_EXTERNAL": app.state.config.TASK_MODEL_EXTERNAL,
  577. "TITLE_GENERATION_PROMPT_TEMPLATE": app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE,
  578. "SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE": app.state.config.SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE,
  579. "SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD": app.state.config.SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD,
  580. "TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE": app.state.config.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE,
  581. }
  582. class TaskConfigForm(BaseModel):
  583. TASK_MODEL: Optional[str]
  584. TASK_MODEL_EXTERNAL: Optional[str]
  585. TITLE_GENERATION_PROMPT_TEMPLATE: str
  586. SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE: str
  587. SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD: int
  588. TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE: str
  589. @app.post("/api/task/config/update")
  590. async def update_task_config(form_data: TaskConfigForm, user=Depends(get_admin_user)):
  591. app.state.config.TASK_MODEL = form_data.TASK_MODEL
  592. app.state.config.TASK_MODEL_EXTERNAL = form_data.TASK_MODEL_EXTERNAL
  593. app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE = (
  594. form_data.TITLE_GENERATION_PROMPT_TEMPLATE
  595. )
  596. app.state.config.SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE = (
  597. form_data.SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE
  598. )
  599. app.state.config.SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD = (
  600. form_data.SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD
  601. )
  602. app.state.config.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE = (
  603. form_data.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE
  604. )
  605. return {
  606. "TASK_MODEL": app.state.config.TASK_MODEL,
  607. "TASK_MODEL_EXTERNAL": app.state.config.TASK_MODEL_EXTERNAL,
  608. "TITLE_GENERATION_PROMPT_TEMPLATE": app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE,
  609. "SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE": app.state.config.SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE,
  610. "SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD": app.state.config.SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD,
  611. "TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE": app.state.config.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE,
  612. }
  613. @app.post("/api/task/title/completions")
  614. async def generate_title(form_data: dict, user=Depends(get_verified_user)):
  615. print("generate_title")
  616. model_id = form_data["model"]
  617. if model_id not in app.state.MODELS:
  618. raise HTTPException(
  619. status_code=status.HTTP_404_NOT_FOUND,
  620. detail="Model not found",
  621. )
  622. # Check if the user has a custom task model
  623. # If the user has a custom task model, use that model
  624. if app.state.MODELS[model_id]["owned_by"] == "ollama":
  625. if app.state.config.TASK_MODEL:
  626. task_model_id = app.state.config.TASK_MODEL
  627. if task_model_id in app.state.MODELS:
  628. model_id = task_model_id
  629. else:
  630. if app.state.config.TASK_MODEL_EXTERNAL:
  631. task_model_id = app.state.config.TASK_MODEL_EXTERNAL
  632. if task_model_id in app.state.MODELS:
  633. model_id = task_model_id
  634. print(model_id)
  635. model = app.state.MODELS[model_id]
  636. template = app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE
  637. content = title_generation_template(
  638. template, form_data["prompt"], user.model_dump()
  639. )
  640. payload = {
  641. "model": model_id,
  642. "messages": [{"role": "user", "content": content}],
  643. "stream": False,
  644. "max_tokens": 50,
  645. "chat_id": form_data.get("chat_id", None),
  646. "title": True,
  647. }
  648. print(payload)
  649. try:
  650. payload = filter_pipeline(payload, user)
  651. except Exception as e:
  652. return JSONResponse(
  653. status_code=e.args[0],
  654. content={"detail": e.args[1]},
  655. )
  656. if model["owned_by"] == "ollama":
  657. return await generate_ollama_chat_completion(
  658. OpenAIChatCompletionForm(**payload), user=user
  659. )
  660. else:
  661. return await generate_openai_chat_completion(payload, user=user)
  662. @app.post("/api/task/query/completions")
  663. async def generate_search_query(form_data: dict, user=Depends(get_verified_user)):
  664. print("generate_search_query")
  665. if len(form_data["prompt"]) < app.state.config.SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD:
  666. raise HTTPException(
  667. status_code=status.HTTP_400_BAD_REQUEST,
  668. detail=f"Skip search query generation for short prompts (< {app.state.config.SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD} characters)",
  669. )
  670. model_id = form_data["model"]
  671. if model_id not in app.state.MODELS:
  672. raise HTTPException(
  673. status_code=status.HTTP_404_NOT_FOUND,
  674. detail="Model not found",
  675. )
  676. # Check if the user has a custom task model
  677. # If the user has a custom task model, use that model
  678. if app.state.MODELS[model_id]["owned_by"] == "ollama":
  679. if app.state.config.TASK_MODEL:
  680. task_model_id = app.state.config.TASK_MODEL
  681. if task_model_id in app.state.MODELS:
  682. model_id = task_model_id
  683. else:
  684. if app.state.config.TASK_MODEL_EXTERNAL:
  685. task_model_id = app.state.config.TASK_MODEL_EXTERNAL
  686. if task_model_id in app.state.MODELS:
  687. model_id = task_model_id
  688. print(model_id)
  689. model = app.state.MODELS[model_id]
  690. template = app.state.config.SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE
  691. content = search_query_generation_template(
  692. template, form_data["prompt"], user.model_dump()
  693. )
  694. payload = {
  695. "model": model_id,
  696. "messages": [{"role": "user", "content": content}],
  697. "stream": False,
  698. "max_tokens": 30,
  699. "task": True,
  700. }
  701. print(payload)
  702. try:
  703. payload = filter_pipeline(payload, user)
  704. except Exception as e:
  705. return JSONResponse(
  706. status_code=e.args[0],
  707. content={"detail": e.args[1]},
  708. )
  709. if model["owned_by"] == "ollama":
  710. return await generate_ollama_chat_completion(
  711. OpenAIChatCompletionForm(**payload), user=user
  712. )
  713. else:
  714. return await generate_openai_chat_completion(payload, user=user)
  715. @app.post("/api/task/emoji/completions")
  716. async def generate_emoji(form_data: dict, user=Depends(get_verified_user)):
  717. print("generate_emoji")
  718. model_id = form_data["model"]
  719. if model_id not in app.state.MODELS:
  720. raise HTTPException(
  721. status_code=status.HTTP_404_NOT_FOUND,
  722. detail="Model not found",
  723. )
  724. # Check if the user has a custom task model
  725. # If the user has a custom task model, use that model
  726. if app.state.MODELS[model_id]["owned_by"] == "ollama":
  727. if app.state.config.TASK_MODEL:
  728. task_model_id = app.state.config.TASK_MODEL
  729. if task_model_id in app.state.MODELS:
  730. model_id = task_model_id
  731. else:
  732. if app.state.config.TASK_MODEL_EXTERNAL:
  733. task_model_id = app.state.config.TASK_MODEL_EXTERNAL
  734. if task_model_id in app.state.MODELS:
  735. model_id = task_model_id
  736. print(model_id)
  737. model = app.state.MODELS[model_id]
  738. template = '''
  739. 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., 😊, 😢, 😡, 😱).
  740. Message: """{{prompt}}"""
  741. '''
  742. content = title_generation_template(
  743. template, form_data["prompt"], user.model_dump()
  744. )
  745. payload = {
  746. "model": model_id,
  747. "messages": [{"role": "user", "content": content}],
  748. "stream": False,
  749. "max_tokens": 4,
  750. "chat_id": form_data.get("chat_id", None),
  751. "task": True,
  752. }
  753. print(payload)
  754. try:
  755. payload = filter_pipeline(payload, user)
  756. except Exception as e:
  757. return JSONResponse(
  758. status_code=e.args[0],
  759. content={"detail": e.args[1]},
  760. )
  761. if model["owned_by"] == "ollama":
  762. return await generate_ollama_chat_completion(
  763. OpenAIChatCompletionForm(**payload), user=user
  764. )
  765. else:
  766. return await generate_openai_chat_completion(payload, user=user)
  767. @app.post("/api/task/tools/completions")
  768. async def get_tools_function_calling(form_data: dict, user=Depends(get_verified_user)):
  769. print("get_tools_function_calling")
  770. model_id = form_data["model"]
  771. if model_id not in app.state.MODELS:
  772. raise HTTPException(
  773. status_code=status.HTTP_404_NOT_FOUND,
  774. detail="Model not found",
  775. )
  776. # Check if the user has a custom task model
  777. # If the user has a custom task model, use that model
  778. if app.state.MODELS[model_id]["owned_by"] == "ollama":
  779. if app.state.config.TASK_MODEL:
  780. task_model_id = app.state.config.TASK_MODEL
  781. if task_model_id in app.state.MODELS:
  782. model_id = task_model_id
  783. else:
  784. if app.state.config.TASK_MODEL_EXTERNAL:
  785. task_model_id = app.state.config.TASK_MODEL_EXTERNAL
  786. if task_model_id in app.state.MODELS:
  787. model_id = task_model_id
  788. print(model_id)
  789. template = app.state.config.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE
  790. try:
  791. context = await get_function_call_response(
  792. form_data["messages"], form_data["tool_id"], template, model_id, user
  793. )
  794. return context
  795. except Exception as e:
  796. return JSONResponse(
  797. status_code=e.args[0],
  798. content={"detail": e.args[1]},
  799. )
  800. @app.post("/api/chat/completions")
  801. async def generate_chat_completions(form_data: dict, user=Depends(get_verified_user)):
  802. model_id = form_data["model"]
  803. if model_id not in app.state.MODELS:
  804. raise HTTPException(
  805. status_code=status.HTTP_404_NOT_FOUND,
  806. detail="Model not found",
  807. )
  808. model = app.state.MODELS[model_id]
  809. print(model)
  810. if model["owned_by"] == "ollama":
  811. return await generate_ollama_chat_completion(
  812. OpenAIChatCompletionForm(**form_data), user=user
  813. )
  814. else:
  815. return await generate_openai_chat_completion(form_data, user=user)
  816. @app.post("/api/chat/completed")
  817. async def chat_completed(form_data: dict, user=Depends(get_verified_user)):
  818. data = form_data
  819. model_id = data["model"]
  820. filters = [
  821. model
  822. for model in app.state.MODELS.values()
  823. if "pipeline" in model
  824. and "type" in model["pipeline"]
  825. and model["pipeline"]["type"] == "filter"
  826. and (
  827. model["pipeline"]["pipelines"] == ["*"]
  828. or any(
  829. model_id == target_model_id
  830. for target_model_id in model["pipeline"]["pipelines"]
  831. )
  832. )
  833. ]
  834. sorted_filters = sorted(filters, key=lambda x: x["pipeline"]["priority"])
  835. print(model_id)
  836. if model_id in app.state.MODELS:
  837. model = app.state.MODELS[model_id]
  838. if "pipeline" in model:
  839. sorted_filters = [model] + sorted_filters
  840. for filter in sorted_filters:
  841. r = None
  842. try:
  843. urlIdx = filter["urlIdx"]
  844. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  845. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  846. if key != "":
  847. headers = {"Authorization": f"Bearer {key}"}
  848. r = requests.post(
  849. f"{url}/{filter['id']}/filter/outlet",
  850. headers=headers,
  851. json={
  852. "user": {"id": user.id, "name": user.name, "role": user.role},
  853. "body": data,
  854. },
  855. )
  856. r.raise_for_status()
  857. data = r.json()
  858. except Exception as e:
  859. # Handle connection error here
  860. print(f"Connection error: {e}")
  861. if r is not None:
  862. try:
  863. res = r.json()
  864. if "detail" in res:
  865. return JSONResponse(
  866. status_code=r.status_code,
  867. content=res,
  868. )
  869. except:
  870. pass
  871. else:
  872. pass
  873. return data
  874. @app.get("/api/pipelines/list")
  875. async def get_pipelines_list(user=Depends(get_admin_user)):
  876. responses = await get_openai_models(raw=True)
  877. print(responses)
  878. urlIdxs = [
  879. idx
  880. for idx, response in enumerate(responses)
  881. if response != None and "pipelines" in response
  882. ]
  883. return {
  884. "data": [
  885. {
  886. "url": openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx],
  887. "idx": urlIdx,
  888. }
  889. for urlIdx in urlIdxs
  890. ]
  891. }
  892. @app.post("/api/pipelines/upload")
  893. async def upload_pipeline(
  894. urlIdx: int = Form(...), file: UploadFile = File(...), user=Depends(get_admin_user)
  895. ):
  896. print("upload_pipeline", urlIdx, file.filename)
  897. # Check if the uploaded file is a python file
  898. if not file.filename.endswith(".py"):
  899. raise HTTPException(
  900. status_code=status.HTTP_400_BAD_REQUEST,
  901. detail="Only Python (.py) files are allowed.",
  902. )
  903. upload_folder = f"{CACHE_DIR}/pipelines"
  904. os.makedirs(upload_folder, exist_ok=True)
  905. file_path = os.path.join(upload_folder, file.filename)
  906. try:
  907. # Save the uploaded file
  908. with open(file_path, "wb") as buffer:
  909. shutil.copyfileobj(file.file, buffer)
  910. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  911. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  912. headers = {"Authorization": f"Bearer {key}"}
  913. with open(file_path, "rb") as f:
  914. files = {"file": f}
  915. r = requests.post(f"{url}/pipelines/upload", headers=headers, files=files)
  916. r.raise_for_status()
  917. data = r.json()
  918. return {**data}
  919. except Exception as e:
  920. # Handle connection error here
  921. print(f"Connection error: {e}")
  922. detail = "Pipeline not found"
  923. if r is not None:
  924. try:
  925. res = r.json()
  926. if "detail" in res:
  927. detail = res["detail"]
  928. except:
  929. pass
  930. raise HTTPException(
  931. status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
  932. detail=detail,
  933. )
  934. finally:
  935. # Ensure the file is deleted after the upload is completed or on failure
  936. if os.path.exists(file_path):
  937. os.remove(file_path)
  938. class AddPipelineForm(BaseModel):
  939. url: str
  940. urlIdx: int
  941. @app.post("/api/pipelines/add")
  942. async def add_pipeline(form_data: AddPipelineForm, user=Depends(get_admin_user)):
  943. r = None
  944. try:
  945. urlIdx = form_data.urlIdx
  946. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  947. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  948. headers = {"Authorization": f"Bearer {key}"}
  949. r = requests.post(
  950. f"{url}/pipelines/add", headers=headers, json={"url": form_data.url}
  951. )
  952. r.raise_for_status()
  953. data = r.json()
  954. return {**data}
  955. except Exception as e:
  956. # Handle connection error here
  957. print(f"Connection error: {e}")
  958. detail = "Pipeline not found"
  959. if r is not None:
  960. try:
  961. res = r.json()
  962. if "detail" in res:
  963. detail = res["detail"]
  964. except:
  965. pass
  966. raise HTTPException(
  967. status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
  968. detail=detail,
  969. )
  970. class DeletePipelineForm(BaseModel):
  971. id: str
  972. urlIdx: int
  973. @app.delete("/api/pipelines/delete")
  974. async def delete_pipeline(form_data: DeletePipelineForm, user=Depends(get_admin_user)):
  975. r = None
  976. try:
  977. urlIdx = form_data.urlIdx
  978. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  979. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  980. headers = {"Authorization": f"Bearer {key}"}
  981. r = requests.delete(
  982. f"{url}/pipelines/delete", headers=headers, json={"id": form_data.id}
  983. )
  984. r.raise_for_status()
  985. data = r.json()
  986. return {**data}
  987. except Exception as e:
  988. # Handle connection error here
  989. print(f"Connection error: {e}")
  990. detail = "Pipeline not found"
  991. if r is not None:
  992. try:
  993. res = r.json()
  994. if "detail" in res:
  995. detail = res["detail"]
  996. except:
  997. pass
  998. raise HTTPException(
  999. status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
  1000. detail=detail,
  1001. )
  1002. @app.get("/api/pipelines")
  1003. async def get_pipelines(urlIdx: Optional[int] = None, user=Depends(get_admin_user)):
  1004. r = None
  1005. try:
  1006. urlIdx
  1007. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  1008. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  1009. headers = {"Authorization": f"Bearer {key}"}
  1010. r = requests.get(f"{url}/pipelines", headers=headers)
  1011. r.raise_for_status()
  1012. data = r.json()
  1013. return {**data}
  1014. except Exception as e:
  1015. # Handle connection error here
  1016. print(f"Connection error: {e}")
  1017. detail = "Pipeline not found"
  1018. if r is not None:
  1019. try:
  1020. res = r.json()
  1021. if "detail" in res:
  1022. detail = res["detail"]
  1023. except:
  1024. pass
  1025. raise HTTPException(
  1026. status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
  1027. detail=detail,
  1028. )
  1029. @app.get("/api/pipelines/{pipeline_id}/valves")
  1030. async def get_pipeline_valves(
  1031. urlIdx: Optional[int], pipeline_id: str, user=Depends(get_admin_user)
  1032. ):
  1033. models = await get_all_models()
  1034. r = None
  1035. try:
  1036. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  1037. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  1038. headers = {"Authorization": f"Bearer {key}"}
  1039. r = requests.get(f"{url}/{pipeline_id}/valves", headers=headers)
  1040. r.raise_for_status()
  1041. data = r.json()
  1042. return {**data}
  1043. except Exception as e:
  1044. # Handle connection error here
  1045. print(f"Connection error: {e}")
  1046. detail = "Pipeline not found"
  1047. if r is not None:
  1048. try:
  1049. res = r.json()
  1050. if "detail" in res:
  1051. detail = res["detail"]
  1052. except:
  1053. pass
  1054. raise HTTPException(
  1055. status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
  1056. detail=detail,
  1057. )
  1058. @app.get("/api/pipelines/{pipeline_id}/valves/spec")
  1059. async def get_pipeline_valves_spec(
  1060. urlIdx: Optional[int], pipeline_id: str, user=Depends(get_admin_user)
  1061. ):
  1062. models = await get_all_models()
  1063. r = None
  1064. try:
  1065. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  1066. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  1067. headers = {"Authorization": f"Bearer {key}"}
  1068. r = requests.get(f"{url}/{pipeline_id}/valves/spec", headers=headers)
  1069. r.raise_for_status()
  1070. data = r.json()
  1071. return {**data}
  1072. except Exception as e:
  1073. # Handle connection error here
  1074. print(f"Connection error: {e}")
  1075. detail = "Pipeline not found"
  1076. if r is not None:
  1077. try:
  1078. res = r.json()
  1079. if "detail" in res:
  1080. detail = res["detail"]
  1081. except:
  1082. pass
  1083. raise HTTPException(
  1084. status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
  1085. detail=detail,
  1086. )
  1087. @app.post("/api/pipelines/{pipeline_id}/valves/update")
  1088. async def update_pipeline_valves(
  1089. urlIdx: Optional[int],
  1090. pipeline_id: str,
  1091. form_data: dict,
  1092. user=Depends(get_admin_user),
  1093. ):
  1094. models = await get_all_models()
  1095. r = None
  1096. try:
  1097. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  1098. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  1099. headers = {"Authorization": f"Bearer {key}"}
  1100. r = requests.post(
  1101. f"{url}/{pipeline_id}/valves/update",
  1102. headers=headers,
  1103. json={**form_data},
  1104. )
  1105. r.raise_for_status()
  1106. data = r.json()
  1107. return {**data}
  1108. except Exception as e:
  1109. # Handle connection error here
  1110. print(f"Connection error: {e}")
  1111. detail = "Pipeline not found"
  1112. if r is not None:
  1113. try:
  1114. res = r.json()
  1115. if "detail" in res:
  1116. detail = res["detail"]
  1117. except:
  1118. pass
  1119. raise HTTPException(
  1120. status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
  1121. detail=detail,
  1122. )
  1123. @app.get("/api/config")
  1124. async def get_app_config():
  1125. # Checking and Handling the Absence of 'ui' in CONFIG_DATA
  1126. default_locale = "en-US"
  1127. if "ui" in CONFIG_DATA:
  1128. default_locale = CONFIG_DATA["ui"].get("default_locale", "en-US")
  1129. # The Rest of the Function Now Uses the Variables Defined Above
  1130. return {
  1131. "status": True,
  1132. "name": WEBUI_NAME,
  1133. "version": VERSION,
  1134. "default_locale": default_locale,
  1135. "default_models": webui_app.state.config.DEFAULT_MODELS,
  1136. "default_prompt_suggestions": webui_app.state.config.DEFAULT_PROMPT_SUGGESTIONS,
  1137. "features": {
  1138. "auth": WEBUI_AUTH,
  1139. "auth_trusted_header": bool(webui_app.state.AUTH_TRUSTED_EMAIL_HEADER),
  1140. "enable_signup": webui_app.state.config.ENABLE_SIGNUP,
  1141. "enable_web_search": rag_app.state.config.ENABLE_RAG_WEB_SEARCH,
  1142. "enable_image_generation": images_app.state.config.ENABLED,
  1143. "enable_community_sharing": webui_app.state.config.ENABLE_COMMUNITY_SHARING,
  1144. "enable_admin_export": ENABLE_ADMIN_EXPORT,
  1145. },
  1146. "audio": {
  1147. "tts": {
  1148. "engine": audio_app.state.config.TTS_ENGINE,
  1149. "voice": audio_app.state.config.TTS_VOICE,
  1150. },
  1151. "stt": {
  1152. "engine": audio_app.state.config.STT_ENGINE,
  1153. },
  1154. },
  1155. "oauth": {
  1156. "providers": {
  1157. name: config.get("name", name)
  1158. for name, config in OAUTH_PROVIDERS.items()
  1159. }
  1160. },
  1161. }
  1162. @app.get("/api/config/model/filter")
  1163. async def get_model_filter_config(user=Depends(get_admin_user)):
  1164. return {
  1165. "enabled": app.state.config.ENABLE_MODEL_FILTER,
  1166. "models": app.state.config.MODEL_FILTER_LIST,
  1167. }
  1168. class ModelFilterConfigForm(BaseModel):
  1169. enabled: bool
  1170. models: List[str]
  1171. @app.post("/api/config/model/filter")
  1172. async def update_model_filter_config(
  1173. form_data: ModelFilterConfigForm, user=Depends(get_admin_user)
  1174. ):
  1175. app.state.config.ENABLE_MODEL_FILTER = form_data.enabled
  1176. app.state.config.MODEL_FILTER_LIST = form_data.models
  1177. return {
  1178. "enabled": app.state.config.ENABLE_MODEL_FILTER,
  1179. "models": app.state.config.MODEL_FILTER_LIST,
  1180. }
  1181. @app.get("/api/webhook")
  1182. async def get_webhook_url(user=Depends(get_admin_user)):
  1183. return {
  1184. "url": app.state.config.WEBHOOK_URL,
  1185. }
  1186. class UrlForm(BaseModel):
  1187. url: str
  1188. @app.post("/api/webhook")
  1189. async def update_webhook_url(form_data: UrlForm, user=Depends(get_admin_user)):
  1190. app.state.config.WEBHOOK_URL = form_data.url
  1191. webui_app.state.WEBHOOK_URL = app.state.config.WEBHOOK_URL
  1192. return {"url": app.state.config.WEBHOOK_URL}
  1193. @app.get("/api/version")
  1194. async def get_app_config():
  1195. return {
  1196. "version": VERSION,
  1197. }
  1198. @app.get("/api/changelog")
  1199. async def get_app_changelog():
  1200. return {key: CHANGELOG[key] for idx, key in enumerate(CHANGELOG) if idx < 5}
  1201. @app.get("/api/version/updates")
  1202. async def get_app_latest_release_version():
  1203. try:
  1204. async with aiohttp.ClientSession(trust_env=True) as session:
  1205. async with session.get(
  1206. "https://api.github.com/repos/open-webui/open-webui/releases/latest"
  1207. ) as response:
  1208. response.raise_for_status()
  1209. data = await response.json()
  1210. latest_version = data["tag_name"]
  1211. return {"current": VERSION, "latest": latest_version[1:]}
  1212. except aiohttp.ClientError as e:
  1213. raise HTTPException(
  1214. status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
  1215. detail=ERROR_MESSAGES.RATE_LIMIT_EXCEEDED,
  1216. )
  1217. ############################
  1218. # OAuth Login & Callback
  1219. ############################
  1220. oauth = OAuth()
  1221. for provider_name, provider_config in OAUTH_PROVIDERS.items():
  1222. oauth.register(
  1223. name=provider_name,
  1224. client_id=provider_config["client_id"],
  1225. client_secret=provider_config["client_secret"],
  1226. server_metadata_url=provider_config["server_metadata_url"],
  1227. client_kwargs={
  1228. "scope": provider_config["scope"],
  1229. },
  1230. )
  1231. # SessionMiddleware is used by authlib for oauth
  1232. if len(OAUTH_PROVIDERS) > 0:
  1233. app.add_middleware(
  1234. SessionMiddleware,
  1235. secret_key=WEBUI_SECRET_KEY,
  1236. session_cookie="oui-session",
  1237. same_site=WEBUI_SESSION_COOKIE_SAME_SITE,
  1238. https_only=WEBUI_SESSION_COOKIE_SECURE,
  1239. )
  1240. @app.get("/oauth/{provider}/login")
  1241. async def oauth_login(provider: str, request: Request):
  1242. if provider not in OAUTH_PROVIDERS:
  1243. raise HTTPException(404)
  1244. redirect_uri = request.url_for("oauth_callback", provider=provider)
  1245. return await oauth.create_client(provider).authorize_redirect(request, redirect_uri)
  1246. @app.get("/oauth/{provider}/callback")
  1247. async def oauth_callback(provider: str, request: Request):
  1248. if provider not in OAUTH_PROVIDERS:
  1249. raise HTTPException(404)
  1250. client = oauth.create_client(provider)
  1251. try:
  1252. token = await client.authorize_access_token(request)
  1253. except Exception as e:
  1254. log.error(f"OAuth callback error: {e}")
  1255. raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)
  1256. user_data: UserInfo = token["userinfo"]
  1257. sub = user_data.get("sub")
  1258. if not sub:
  1259. raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)
  1260. provider_sub = f"{provider}@{sub}"
  1261. # Check if the user exists
  1262. user = Users.get_user_by_oauth_sub(provider_sub)
  1263. if not user:
  1264. # If the user does not exist, check if merging is enabled
  1265. if OAUTH_MERGE_ACCOUNTS_BY_EMAIL.value:
  1266. # Check if the user exists by email
  1267. email = user_data.get("email", "").lower()
  1268. if not email:
  1269. raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)
  1270. user = Users.get_user_by_email(user_data.get("email", "").lower(), True)
  1271. if user:
  1272. # Update the user with the new oauth sub
  1273. Users.update_user_oauth_sub_by_id(user.id, provider_sub)
  1274. if not user:
  1275. # If the user does not exist, check if signups are enabled
  1276. if ENABLE_OAUTH_SIGNUP.value:
  1277. user = Auths.insert_new_auth(
  1278. email=user_data.get("email", "").lower(),
  1279. password=get_password_hash(
  1280. str(uuid.uuid4())
  1281. ), # Random password, not used
  1282. name=user_data.get("name", "User"),
  1283. profile_image_url=user_data.get("picture", "/user.png"),
  1284. role=webui_app.state.config.DEFAULT_USER_ROLE,
  1285. oauth_sub=provider_sub,
  1286. )
  1287. if webui_app.state.config.WEBHOOK_URL:
  1288. post_webhook(
  1289. webui_app.state.config.WEBHOOK_URL,
  1290. WEBHOOK_MESSAGES.USER_SIGNUP(user.name),
  1291. {
  1292. "action": "signup",
  1293. "message": WEBHOOK_MESSAGES.USER_SIGNUP(user.name),
  1294. "user": user.model_dump_json(exclude_none=True),
  1295. },
  1296. )
  1297. else:
  1298. raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)
  1299. jwt_token = create_token(
  1300. data={"id": user.id},
  1301. expires_delta=parse_duration(webui_app.state.config.JWT_EXPIRES_IN),
  1302. )
  1303. # Redirect back to the frontend with the JWT token
  1304. redirect_url = f"{request.base_url}auth#token={jwt_token}"
  1305. return RedirectResponse(url=redirect_url)
  1306. @app.get("/manifest.json")
  1307. async def get_manifest_json():
  1308. return {
  1309. "name": WEBUI_NAME,
  1310. "short_name": WEBUI_NAME,
  1311. "start_url": "/",
  1312. "display": "standalone",
  1313. "background_color": "#343541",
  1314. "theme_color": "#343541",
  1315. "orientation": "portrait-primary",
  1316. "icons": [{"src": "/static/logo.png", "type": "image/png", "sizes": "500x500"}],
  1317. }
  1318. @app.get("/opensearch.xml")
  1319. async def get_opensearch_xml():
  1320. xml_content = rf"""
  1321. <OpenSearchDescription xmlns="http://a9.com/-/spec/opensearch/1.1/" xmlns:moz="http://www.mozilla.org/2006/browser/search/">
  1322. <ShortName>{WEBUI_NAME}</ShortName>
  1323. <Description>Search {WEBUI_NAME}</Description>
  1324. <InputEncoding>UTF-8</InputEncoding>
  1325. <Image width="16" height="16" type="image/x-icon">{WEBUI_URL}/favicon.png</Image>
  1326. <Url type="text/html" method="get" template="{WEBUI_URL}/?q={"{searchTerms}"}"/>
  1327. <moz:SearchForm>{WEBUI_URL}</moz:SearchForm>
  1328. </OpenSearchDescription>
  1329. """
  1330. return Response(content=xml_content, media_type="application/xml")
  1331. @app.get("/health")
  1332. async def healthcheck():
  1333. return {"status": True}
  1334. app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
  1335. app.mount("/cache", StaticFiles(directory=CACHE_DIR), name="cache")
  1336. if os.path.exists(FRONTEND_BUILD_DIR):
  1337. mimetypes.add_type("text/javascript", ".js")
  1338. app.mount(
  1339. "/",
  1340. SPAStaticFiles(directory=FRONTEND_BUILD_DIR, html=True),
  1341. name="spa-static-files",
  1342. )
  1343. else:
  1344. log.warning(
  1345. f"Frontend build directory not found at '{FRONTEND_BUILD_DIR}'. Serving API only."
  1346. )