main.py 37 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225
  1. from contextlib import asynccontextmanager
  2. from bs4 import BeautifulSoup
  3. import json
  4. import markdown
  5. import time
  6. import os
  7. import sys
  8. import logging
  9. import aiohttp
  10. import requests
  11. import mimetypes
  12. import shutil
  13. import os
  14. import asyncio
  15. from fastapi import FastAPI, Request, Depends, status, UploadFile, File, Form
  16. from fastapi.staticfiles import StaticFiles
  17. from fastapi.responses import JSONResponse
  18. from fastapi import HTTPException
  19. from fastapi.middleware.wsgi import WSGIMiddleware
  20. from fastapi.middleware.cors import CORSMiddleware
  21. from starlette.exceptions import HTTPException as StarletteHTTPException
  22. from starlette.middleware.base import BaseHTTPMiddleware
  23. from starlette.responses import StreamingResponse, Response
  24. from apps.socket.main import app as socket_app
  25. from apps.ollama.main import (
  26. app as ollama_app,
  27. OpenAIChatCompletionForm,
  28. get_all_models as get_ollama_models,
  29. generate_openai_chat_completion as generate_ollama_chat_completion,
  30. )
  31. from apps.openai.main import (
  32. app as openai_app,
  33. get_all_models as get_openai_models,
  34. generate_chat_completion as generate_openai_chat_completion,
  35. )
  36. from apps.audio.main import app as audio_app
  37. from apps.images.main import app as images_app
  38. from apps.rag.main import app as rag_app
  39. from apps.webui.main import app as webui_app
  40. from pydantic import BaseModel
  41. from typing import List, Optional
  42. from apps.webui.models.models import Models, ModelModel
  43. from utils.utils import (
  44. get_admin_user,
  45. get_verified_user,
  46. get_current_user,
  47. get_http_authorization_cred,
  48. )
  49. from utils.task import title_generation_template, search_query_generation_template
  50. from apps.rag.utils import rag_messages
  51. from config import (
  52. CONFIG_DATA,
  53. WEBUI_NAME,
  54. WEBUI_URL,
  55. WEBUI_AUTH,
  56. ENV,
  57. VERSION,
  58. CHANGELOG,
  59. FRONTEND_BUILD_DIR,
  60. CACHE_DIR,
  61. STATIC_DIR,
  62. ENABLE_OPENAI_API,
  63. ENABLE_OLLAMA_API,
  64. ENABLE_MODEL_FILTER,
  65. MODEL_FILTER_LIST,
  66. GLOBAL_LOG_LEVEL,
  67. SRC_LOG_LEVELS,
  68. WEBHOOK_URL,
  69. ENABLE_ADMIN_EXPORT,
  70. WEBUI_BUILD_HASH,
  71. TASK_MODEL,
  72. TASK_MODEL_EXTERNAL,
  73. TITLE_GENERATION_PROMPT_TEMPLATE,
  74. SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE,
  75. SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD,
  76. AppConfig,
  77. )
  78. from constants import ERROR_MESSAGES
  79. logging.basicConfig(stream=sys.stdout, level=GLOBAL_LOG_LEVEL)
  80. log = logging.getLogger(__name__)
  81. log.setLevel(SRC_LOG_LEVELS["MAIN"])
  82. class SPAStaticFiles(StaticFiles):
  83. async def get_response(self, path: str, scope):
  84. try:
  85. return await super().get_response(path, scope)
  86. except (HTTPException, StarletteHTTPException) as ex:
  87. if ex.status_code == 404:
  88. return await super().get_response("index.html", scope)
  89. else:
  90. raise ex
  91. print(
  92. rf"""
  93. ___ __ __ _ _ _ ___
  94. / _ \ _ __ ___ _ __ \ \ / /__| |__ | | | |_ _|
  95. | | | | '_ \ / _ \ '_ \ \ \ /\ / / _ \ '_ \| | | || |
  96. | |_| | |_) | __/ | | | \ V V / __/ |_) | |_| || |
  97. \___/| .__/ \___|_| |_| \_/\_/ \___|_.__/ \___/|___|
  98. |_|
  99. v{VERSION} - building the best open-source AI user interface.
  100. {f"Commit: {WEBUI_BUILD_HASH}" if WEBUI_BUILD_HASH != "dev-build" else ""}
  101. https://github.com/open-webui/open-webui
  102. """
  103. )
  104. @asynccontextmanager
  105. async def lifespan(app: FastAPI):
  106. yield
  107. app = FastAPI(
  108. docs_url="/docs" if ENV == "dev" else None, redoc_url=None, lifespan=lifespan
  109. )
  110. app.state.config = AppConfig()
  111. app.state.config.ENABLE_OPENAI_API = ENABLE_OPENAI_API
  112. app.state.config.ENABLE_OLLAMA_API = ENABLE_OLLAMA_API
  113. app.state.config.ENABLE_MODEL_FILTER = ENABLE_MODEL_FILTER
  114. app.state.config.MODEL_FILTER_LIST = MODEL_FILTER_LIST
  115. app.state.config.WEBHOOK_URL = WEBHOOK_URL
  116. app.state.config.TASK_MODEL = TASK_MODEL
  117. app.state.config.TASK_MODEL_EXTERNAL = TASK_MODEL_EXTERNAL
  118. app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE = TITLE_GENERATION_PROMPT_TEMPLATE
  119. app.state.config.SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE = (
  120. SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE
  121. )
  122. app.state.config.SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD = (
  123. SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD
  124. )
  125. app.state.MODELS = {}
  126. origins = ["*"]
  127. # Custom middleware to add security headers
  128. # class SecurityHeadersMiddleware(BaseHTTPMiddleware):
  129. # async def dispatch(self, request: Request, call_next):
  130. # response: Response = await call_next(request)
  131. # response.headers["Cross-Origin-Opener-Policy"] = "same-origin"
  132. # response.headers["Cross-Origin-Embedder-Policy"] = "require-corp"
  133. # return response
  134. # app.add_middleware(SecurityHeadersMiddleware)
  135. class RAGMiddleware(BaseHTTPMiddleware):
  136. async def dispatch(self, request: Request, call_next):
  137. return_citations = False
  138. if request.method == "POST" and (
  139. "/ollama/api/chat" in request.url.path
  140. or "/chat/completions" in request.url.path
  141. ):
  142. log.debug(f"request.url.path: {request.url.path}")
  143. # Read the original request body
  144. body = await request.body()
  145. # Decode body to string
  146. body_str = body.decode("utf-8")
  147. # Parse string to JSON
  148. data = json.loads(body_str) if body_str else {}
  149. return_citations = data.get("citations", False)
  150. if "citations" in data:
  151. del data["citations"]
  152. # Example: Add a new key-value pair or modify existing ones
  153. # data["modified"] = True # Example modification
  154. if "docs" in data:
  155. data = {**data}
  156. data["messages"], citations = rag_messages(
  157. docs=data["docs"],
  158. messages=data["messages"],
  159. template=rag_app.state.config.RAG_TEMPLATE,
  160. embedding_function=rag_app.state.EMBEDDING_FUNCTION,
  161. k=rag_app.state.config.TOP_K,
  162. reranking_function=rag_app.state.sentence_transformer_rf,
  163. r=rag_app.state.config.RELEVANCE_THRESHOLD,
  164. hybrid_search=rag_app.state.config.ENABLE_RAG_HYBRID_SEARCH,
  165. )
  166. del data["docs"]
  167. log.debug(
  168. f"data['messages']: {data['messages']}, citations: {citations}"
  169. )
  170. modified_body_bytes = json.dumps(data).encode("utf-8")
  171. # Replace the request body with the modified one
  172. request._body = modified_body_bytes
  173. # Set custom header to ensure content-length matches new body length
  174. request.headers.__dict__["_list"] = [
  175. (b"content-length", str(len(modified_body_bytes)).encode("utf-8")),
  176. *[
  177. (k, v)
  178. for k, v in request.headers.raw
  179. if k.lower() != b"content-length"
  180. ],
  181. ]
  182. response = await call_next(request)
  183. if return_citations:
  184. # Inject the citations into the response
  185. if isinstance(response, StreamingResponse):
  186. # If it's a streaming response, inject it as SSE event or NDJSON line
  187. content_type = response.headers.get("Content-Type")
  188. if "text/event-stream" in content_type:
  189. return StreamingResponse(
  190. self.openai_stream_wrapper(response.body_iterator, citations),
  191. )
  192. if "application/x-ndjson" in content_type:
  193. return StreamingResponse(
  194. self.ollama_stream_wrapper(response.body_iterator, citations),
  195. )
  196. return response
  197. async def _receive(self, body: bytes):
  198. return {"type": "http.request", "body": body, "more_body": False}
  199. async def openai_stream_wrapper(self, original_generator, citations):
  200. yield f"data: {json.dumps({'citations': citations})}\n\n"
  201. async for data in original_generator:
  202. yield data
  203. async def ollama_stream_wrapper(self, original_generator, citations):
  204. yield f"{json.dumps({'citations': citations})}\n"
  205. async for data in original_generator:
  206. yield data
  207. app.add_middleware(RAGMiddleware)
  208. def filter_pipeline(payload, user):
  209. user = {"id": user.id, "name": user.name, "role": user.role}
  210. model_id = payload["model"]
  211. filters = [
  212. model
  213. for model in app.state.MODELS.values()
  214. if "pipeline" in model
  215. and "type" in model["pipeline"]
  216. and model["pipeline"]["type"] == "filter"
  217. and (
  218. model["pipeline"]["pipelines"] == ["*"]
  219. or any(
  220. model_id == target_model_id
  221. for target_model_id in model["pipeline"]["pipelines"]
  222. )
  223. )
  224. ]
  225. sorted_filters = sorted(filters, key=lambda x: x["pipeline"]["priority"])
  226. model = app.state.MODELS[model_id]
  227. if "pipeline" in model:
  228. sorted_filters.append(model)
  229. for filter in sorted_filters:
  230. r = None
  231. try:
  232. urlIdx = filter["urlIdx"]
  233. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  234. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  235. if key != "":
  236. headers = {"Authorization": f"Bearer {key}"}
  237. r = requests.post(
  238. f"{url}/{filter['id']}/filter/inlet",
  239. headers=headers,
  240. json={
  241. "user": user,
  242. "body": payload,
  243. },
  244. )
  245. r.raise_for_status()
  246. payload = r.json()
  247. except Exception as e:
  248. # Handle connection error here
  249. print(f"Connection error: {e}")
  250. if r is not None:
  251. try:
  252. res = r.json()
  253. if "detail" in res:
  254. return JSONResponse(
  255. status_code=r.status_code,
  256. content=res,
  257. )
  258. except:
  259. pass
  260. else:
  261. pass
  262. if "pipeline" not in app.state.MODELS[model_id]:
  263. if "chat_id" in payload:
  264. del payload["chat_id"]
  265. if "title" in payload:
  266. del payload["title"]
  267. return payload
  268. class PipelineMiddleware(BaseHTTPMiddleware):
  269. async def dispatch(self, request: Request, call_next):
  270. if request.method == "POST" and (
  271. "/ollama/api/chat" in request.url.path
  272. or "/chat/completions" in request.url.path
  273. ):
  274. log.debug(f"request.url.path: {request.url.path}")
  275. # Read the original request body
  276. body = await request.body()
  277. # Decode body to string
  278. body_str = body.decode("utf-8")
  279. # Parse string to JSON
  280. data = json.loads(body_str) if body_str else {}
  281. user = get_current_user(
  282. get_http_authorization_cred(request.headers.get("Authorization"))
  283. )
  284. data = filter_pipeline(data, user)
  285. modified_body_bytes = json.dumps(data).encode("utf-8")
  286. # Replace the request body with the modified one
  287. request._body = modified_body_bytes
  288. # Set custom header to ensure content-length matches new body length
  289. request.headers.__dict__["_list"] = [
  290. (b"content-length", str(len(modified_body_bytes)).encode("utf-8")),
  291. *[
  292. (k, v)
  293. for k, v in request.headers.raw
  294. if k.lower() != b"content-length"
  295. ],
  296. ]
  297. response = await call_next(request)
  298. return response
  299. async def _receive(self, body: bytes):
  300. return {"type": "http.request", "body": body, "more_body": False}
  301. app.add_middleware(PipelineMiddleware)
  302. app.add_middleware(
  303. CORSMiddleware,
  304. allow_origins=origins,
  305. allow_credentials=True,
  306. allow_methods=["*"],
  307. allow_headers=["*"],
  308. )
  309. @app.middleware("http")
  310. async def check_url(request: Request, call_next):
  311. if len(app.state.MODELS) == 0:
  312. await get_all_models()
  313. else:
  314. pass
  315. start_time = int(time.time())
  316. response = await call_next(request)
  317. process_time = int(time.time()) - start_time
  318. response.headers["X-Process-Time"] = str(process_time)
  319. return response
  320. @app.middleware("http")
  321. async def update_embedding_function(request: Request, call_next):
  322. response = await call_next(request)
  323. if "/embedding/update" in request.url.path:
  324. webui_app.state.EMBEDDING_FUNCTION = rag_app.state.EMBEDDING_FUNCTION
  325. return response
  326. app.mount("/ws", socket_app)
  327. app.mount("/ollama", ollama_app)
  328. app.mount("/openai", openai_app)
  329. app.mount("/images/api/v1", images_app)
  330. app.mount("/audio/api/v1", audio_app)
  331. app.mount("/rag/api/v1", rag_app)
  332. app.mount("/api/v1", webui_app)
  333. webui_app.state.EMBEDDING_FUNCTION = rag_app.state.EMBEDDING_FUNCTION
  334. async def get_all_models():
  335. openai_models = []
  336. ollama_models = []
  337. if app.state.config.ENABLE_OPENAI_API:
  338. openai_models = await get_openai_models()
  339. openai_models = openai_models["data"]
  340. if app.state.config.ENABLE_OLLAMA_API:
  341. ollama_models = await get_ollama_models()
  342. ollama_models = [
  343. {
  344. "id": model["model"],
  345. "name": model["name"],
  346. "object": "model",
  347. "created": int(time.time()),
  348. "owned_by": "ollama",
  349. "ollama": model,
  350. }
  351. for model in ollama_models["models"]
  352. ]
  353. models = openai_models + ollama_models
  354. custom_models = Models.get_all_models()
  355. for custom_model in custom_models:
  356. if custom_model.base_model_id == None:
  357. for model in models:
  358. if (
  359. custom_model.id == model["id"]
  360. or custom_model.id == model["id"].split(":")[0]
  361. ):
  362. model["name"] = custom_model.name
  363. model["info"] = custom_model.model_dump()
  364. else:
  365. owned_by = "openai"
  366. for model in models:
  367. if (
  368. custom_model.base_model_id == model["id"]
  369. or custom_model.base_model_id == model["id"].split(":")[0]
  370. ):
  371. owned_by = model["owned_by"]
  372. break
  373. models.append(
  374. {
  375. "id": custom_model.id,
  376. "name": custom_model.name,
  377. "object": "model",
  378. "created": custom_model.created_at,
  379. "owned_by": owned_by,
  380. "info": custom_model.model_dump(),
  381. "preset": True,
  382. }
  383. )
  384. app.state.MODELS = {model["id"]: model for model in models}
  385. webui_app.state.MODELS = app.state.MODELS
  386. return models
  387. @app.get("/api/models")
  388. async def get_models(user=Depends(get_verified_user)):
  389. models = await get_all_models()
  390. # Filter out filter pipelines
  391. models = [
  392. model
  393. for model in models
  394. if "pipeline" not in model or model["pipeline"].get("type", None) != "filter"
  395. ]
  396. if app.state.config.ENABLE_MODEL_FILTER:
  397. if user.role == "user":
  398. models = list(
  399. filter(
  400. lambda model: model["id"] in app.state.config.MODEL_FILTER_LIST,
  401. models,
  402. )
  403. )
  404. return {"data": models}
  405. return {"data": models}
  406. @app.get("/api/task/config")
  407. async def get_task_config(user=Depends(get_verified_user)):
  408. return {
  409. "TASK_MODEL": app.state.config.TASK_MODEL,
  410. "TASK_MODEL_EXTERNAL": app.state.config.TASK_MODEL_EXTERNAL,
  411. "TITLE_GENERATION_PROMPT_TEMPLATE": app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE,
  412. "SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE": app.state.config.SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE,
  413. }
  414. class TaskConfigForm(BaseModel):
  415. TASK_MODEL: Optional[str]
  416. TASK_MODEL_EXTERNAL: Optional[str]
  417. TITLE_GENERATION_PROMPT_TEMPLATE: str
  418. SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE: str
  419. @app.post("/api/task/config/update")
  420. async def update_task_config(form_data: TaskConfigForm, user=Depends(get_admin_user)):
  421. app.state.config.TASK_MODEL = form_data.TASK_MODEL
  422. app.state.config.TASK_MODEL_EXTERNAL = form_data.TASK_MODEL_EXTERNAL
  423. app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE = (
  424. form_data.TITLE_GENERATION_PROMPT_TEMPLATE
  425. )
  426. app.state.config.SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE = (
  427. form_data.SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE
  428. )
  429. return {
  430. "TASK_MODEL": app.state.config.TASK_MODEL,
  431. "TASK_MODEL_EXTERNAL": app.state.config.TASK_MODEL_EXTERNAL,
  432. "TITLE_GENERATION_PROMPT_TEMPLATE": app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE,
  433. "SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE": app.state.config.SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE,
  434. }
  435. @app.post("/api/task/title/completions")
  436. async def generate_title(form_data: dict, user=Depends(get_verified_user)):
  437. print("generate_title")
  438. model_id = form_data["model"]
  439. if model_id not in app.state.MODELS:
  440. raise HTTPException(
  441. status_code=status.HTTP_404_NOT_FOUND,
  442. detail="Model not found",
  443. )
  444. # Check if the user has a custom task model
  445. # If the user has a custom task model, use that model
  446. if app.state.MODELS[model_id]["owned_by"] == "ollama":
  447. if app.state.config.TASK_MODEL:
  448. task_model_id = app.state.config.TASK_MODEL
  449. if task_model_id in app.state.MODELS:
  450. model_id = task_model_id
  451. else:
  452. if app.state.config.TASK_MODEL_EXTERNAL:
  453. task_model_id = app.state.config.TASK_MODEL_EXTERNAL
  454. if task_model_id in app.state.MODELS:
  455. model_id = task_model_id
  456. print(model_id)
  457. model = app.state.MODELS[model_id]
  458. template = app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE
  459. content = title_generation_template(
  460. template, form_data["prompt"], user.model_dump()
  461. )
  462. payload = {
  463. "model": model_id,
  464. "messages": [{"role": "user", "content": content}],
  465. "stream": False,
  466. "max_tokens": 50,
  467. "chat_id": form_data.get("chat_id", None),
  468. "title": True,
  469. }
  470. print(payload)
  471. payload = filter_pipeline(payload, user)
  472. if model["owned_by"] == "ollama":
  473. return await generate_ollama_chat_completion(
  474. OpenAIChatCompletionForm(**payload), user=user
  475. )
  476. else:
  477. return await generate_openai_chat_completion(payload, user=user)
  478. @app.post("/api/task/query/completions")
  479. async def generate_search_query(form_data: dict, user=Depends(get_verified_user)):
  480. print("generate_search_query")
  481. if len(form_data["prompt"]) < app.state.config.SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD:
  482. raise HTTPException(
  483. status_code=status.HTTP_400_BAD_REQUEST,
  484. detail=f"Skip search query generation for short prompts (< {app.state.config.SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD} characters)",
  485. )
  486. model_id = form_data["model"]
  487. if model_id not in app.state.MODELS:
  488. raise HTTPException(
  489. status_code=status.HTTP_404_NOT_FOUND,
  490. detail="Model not found",
  491. )
  492. # Check if the user has a custom task model
  493. # If the user has a custom task model, use that model
  494. if app.state.MODELS[model_id]["owned_by"] == "ollama":
  495. if app.state.config.TASK_MODEL:
  496. task_model_id = app.state.config.TASK_MODEL
  497. if task_model_id in app.state.MODELS:
  498. model_id = task_model_id
  499. else:
  500. if app.state.config.TASK_MODEL_EXTERNAL:
  501. task_model_id = app.state.config.TASK_MODEL_EXTERNAL
  502. if task_model_id in app.state.MODELS:
  503. model_id = task_model_id
  504. print(model_id)
  505. model = app.state.MODELS[model_id]
  506. template = app.state.config.SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE
  507. content = search_query_generation_template(
  508. template, form_data["prompt"], user.model_dump()
  509. )
  510. payload = {
  511. "model": model_id,
  512. "messages": [{"role": "user", "content": content}],
  513. "stream": False,
  514. "max_tokens": 30,
  515. }
  516. print(payload)
  517. payload = filter_pipeline(payload, user)
  518. if model["owned_by"] == "ollama":
  519. return await generate_ollama_chat_completion(
  520. OpenAIChatCompletionForm(**payload), user=user
  521. )
  522. else:
  523. return await generate_openai_chat_completion(payload, user=user)
  524. @app.post("/api/chat/completions")
  525. async def generate_chat_completions(form_data: dict, user=Depends(get_verified_user)):
  526. model_id = form_data["model"]
  527. if model_id not in app.state.MODELS:
  528. raise HTTPException(
  529. status_code=status.HTTP_404_NOT_FOUND,
  530. detail="Model not found",
  531. )
  532. model = app.state.MODELS[model_id]
  533. print(model)
  534. if model["owned_by"] == "ollama":
  535. return await generate_ollama_chat_completion(
  536. OpenAIChatCompletionForm(**form_data), user=user
  537. )
  538. else:
  539. return await generate_openai_chat_completion(form_data, user=user)
  540. @app.post("/api/chat/completed")
  541. async def chat_completed(form_data: dict, user=Depends(get_verified_user)):
  542. data = form_data
  543. model_id = data["model"]
  544. filters = [
  545. model
  546. for model in app.state.MODELS.values()
  547. if "pipeline" in model
  548. and "type" in model["pipeline"]
  549. and model["pipeline"]["type"] == "filter"
  550. and (
  551. model["pipeline"]["pipelines"] == ["*"]
  552. or any(
  553. model_id == target_model_id
  554. for target_model_id in model["pipeline"]["pipelines"]
  555. )
  556. )
  557. ]
  558. sorted_filters = sorted(filters, key=lambda x: x["pipeline"]["priority"])
  559. print(model_id)
  560. if model_id in app.state.MODELS:
  561. model = app.state.MODELS[model_id]
  562. if "pipeline" in model:
  563. sorted_filters = [model] + sorted_filters
  564. for filter in sorted_filters:
  565. r = None
  566. try:
  567. urlIdx = filter["urlIdx"]
  568. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  569. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  570. if key != "":
  571. headers = {"Authorization": f"Bearer {key}"}
  572. r = requests.post(
  573. f"{url}/{filter['id']}/filter/outlet",
  574. headers=headers,
  575. json={
  576. "user": {"id": user.id, "name": user.name, "role": user.role},
  577. "body": data,
  578. },
  579. )
  580. r.raise_for_status()
  581. data = r.json()
  582. except Exception as e:
  583. # Handle connection error here
  584. print(f"Connection error: {e}")
  585. if r is not None:
  586. try:
  587. res = r.json()
  588. if "detail" in res:
  589. return JSONResponse(
  590. status_code=r.status_code,
  591. content=res,
  592. )
  593. except:
  594. pass
  595. else:
  596. pass
  597. return data
  598. @app.get("/api/pipelines/list")
  599. async def get_pipelines_list(user=Depends(get_admin_user)):
  600. responses = await get_openai_models(raw=True)
  601. print(responses)
  602. urlIdxs = [
  603. idx
  604. for idx, response in enumerate(responses)
  605. if response != None and "pipelines" in response
  606. ]
  607. return {
  608. "data": [
  609. {
  610. "url": openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx],
  611. "idx": urlIdx,
  612. }
  613. for urlIdx in urlIdxs
  614. ]
  615. }
  616. @app.post("/api/pipelines/upload")
  617. async def upload_pipeline(
  618. urlIdx: int = Form(...), file: UploadFile = File(...), user=Depends(get_admin_user)
  619. ):
  620. print("upload_pipeline", urlIdx, file.filename)
  621. # Check if the uploaded file is a python file
  622. if not file.filename.endswith(".py"):
  623. raise HTTPException(
  624. status_code=status.HTTP_400_BAD_REQUEST,
  625. detail="Only Python (.py) files are allowed.",
  626. )
  627. upload_folder = f"{CACHE_DIR}/pipelines"
  628. os.makedirs(upload_folder, exist_ok=True)
  629. file_path = os.path.join(upload_folder, file.filename)
  630. try:
  631. # Save the uploaded file
  632. with open(file_path, "wb") as buffer:
  633. shutil.copyfileobj(file.file, buffer)
  634. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  635. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  636. headers = {"Authorization": f"Bearer {key}"}
  637. with open(file_path, "rb") as f:
  638. files = {"file": f}
  639. r = requests.post(f"{url}/pipelines/upload", headers=headers, files=files)
  640. r.raise_for_status()
  641. data = r.json()
  642. return {**data}
  643. except Exception as e:
  644. # Handle connection error here
  645. print(f"Connection error: {e}")
  646. detail = "Pipeline not found"
  647. if r is not None:
  648. try:
  649. res = r.json()
  650. if "detail" in res:
  651. detail = res["detail"]
  652. except:
  653. pass
  654. raise HTTPException(
  655. status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
  656. detail=detail,
  657. )
  658. finally:
  659. # Ensure the file is deleted after the upload is completed or on failure
  660. if os.path.exists(file_path):
  661. os.remove(file_path)
  662. class AddPipelineForm(BaseModel):
  663. url: str
  664. urlIdx: int
  665. @app.post("/api/pipelines/add")
  666. async def add_pipeline(form_data: AddPipelineForm, user=Depends(get_admin_user)):
  667. r = None
  668. try:
  669. urlIdx = form_data.urlIdx
  670. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  671. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  672. headers = {"Authorization": f"Bearer {key}"}
  673. r = requests.post(
  674. f"{url}/pipelines/add", headers=headers, json={"url": form_data.url}
  675. )
  676. r.raise_for_status()
  677. data = r.json()
  678. return {**data}
  679. except Exception as e:
  680. # Handle connection error here
  681. print(f"Connection error: {e}")
  682. detail = "Pipeline not found"
  683. if r is not None:
  684. try:
  685. res = r.json()
  686. if "detail" in res:
  687. detail = res["detail"]
  688. except:
  689. pass
  690. raise HTTPException(
  691. status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
  692. detail=detail,
  693. )
  694. class DeletePipelineForm(BaseModel):
  695. id: str
  696. urlIdx: int
  697. @app.delete("/api/pipelines/delete")
  698. async def delete_pipeline(form_data: DeletePipelineForm, user=Depends(get_admin_user)):
  699. r = None
  700. try:
  701. urlIdx = form_data.urlIdx
  702. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  703. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  704. headers = {"Authorization": f"Bearer {key}"}
  705. r = requests.delete(
  706. f"{url}/pipelines/delete", headers=headers, json={"id": form_data.id}
  707. )
  708. r.raise_for_status()
  709. data = r.json()
  710. return {**data}
  711. except Exception as e:
  712. # Handle connection error here
  713. print(f"Connection error: {e}")
  714. detail = "Pipeline not found"
  715. if r is not None:
  716. try:
  717. res = r.json()
  718. if "detail" in res:
  719. detail = res["detail"]
  720. except:
  721. pass
  722. raise HTTPException(
  723. status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
  724. detail=detail,
  725. )
  726. @app.get("/api/pipelines")
  727. async def get_pipelines(urlIdx: Optional[int] = None, user=Depends(get_admin_user)):
  728. r = None
  729. try:
  730. urlIdx
  731. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  732. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  733. headers = {"Authorization": f"Bearer {key}"}
  734. r = requests.get(f"{url}/pipelines", headers=headers)
  735. r.raise_for_status()
  736. data = r.json()
  737. return {**data}
  738. except Exception as e:
  739. # Handle connection error here
  740. print(f"Connection error: {e}")
  741. detail = "Pipeline not found"
  742. if r is not None:
  743. try:
  744. res = r.json()
  745. if "detail" in res:
  746. detail = res["detail"]
  747. except:
  748. pass
  749. raise HTTPException(
  750. status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
  751. detail=detail,
  752. )
  753. @app.get("/api/pipelines/{pipeline_id}/valves")
  754. async def get_pipeline_valves(
  755. urlIdx: Optional[int], pipeline_id: str, user=Depends(get_admin_user)
  756. ):
  757. models = await get_all_models()
  758. r = None
  759. try:
  760. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  761. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  762. headers = {"Authorization": f"Bearer {key}"}
  763. r = requests.get(f"{url}/{pipeline_id}/valves", headers=headers)
  764. r.raise_for_status()
  765. data = r.json()
  766. return {**data}
  767. except Exception as e:
  768. # Handle connection error here
  769. print(f"Connection error: {e}")
  770. detail = "Pipeline not found"
  771. if r is not None:
  772. try:
  773. res = r.json()
  774. if "detail" in res:
  775. detail = res["detail"]
  776. except:
  777. pass
  778. raise HTTPException(
  779. status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
  780. detail=detail,
  781. )
  782. @app.get("/api/pipelines/{pipeline_id}/valves/spec")
  783. async def get_pipeline_valves_spec(
  784. urlIdx: Optional[int], pipeline_id: str, user=Depends(get_admin_user)
  785. ):
  786. models = await get_all_models()
  787. r = None
  788. try:
  789. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  790. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  791. headers = {"Authorization": f"Bearer {key}"}
  792. r = requests.get(f"{url}/{pipeline_id}/valves/spec", headers=headers)
  793. r.raise_for_status()
  794. data = r.json()
  795. return {**data}
  796. except Exception as e:
  797. # Handle connection error here
  798. print(f"Connection error: {e}")
  799. detail = "Pipeline not found"
  800. if r is not None:
  801. try:
  802. res = r.json()
  803. if "detail" in res:
  804. detail = res["detail"]
  805. except:
  806. pass
  807. raise HTTPException(
  808. status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
  809. detail=detail,
  810. )
  811. @app.post("/api/pipelines/{pipeline_id}/valves/update")
  812. async def update_pipeline_valves(
  813. urlIdx: Optional[int],
  814. pipeline_id: str,
  815. form_data: dict,
  816. user=Depends(get_admin_user),
  817. ):
  818. models = await get_all_models()
  819. r = None
  820. try:
  821. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  822. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  823. headers = {"Authorization": f"Bearer {key}"}
  824. r = requests.post(
  825. f"{url}/{pipeline_id}/valves/update",
  826. headers=headers,
  827. json={**form_data},
  828. )
  829. r.raise_for_status()
  830. data = r.json()
  831. return {**data}
  832. except Exception as e:
  833. # Handle connection error here
  834. print(f"Connection error: {e}")
  835. detail = "Pipeline not found"
  836. if r is not None:
  837. try:
  838. res = r.json()
  839. if "detail" in res:
  840. detail = res["detail"]
  841. except:
  842. pass
  843. raise HTTPException(
  844. status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
  845. detail=detail,
  846. )
  847. @app.get("/api/config")
  848. async def get_app_config():
  849. # Checking and Handling the Absence of 'ui' in CONFIG_DATA
  850. default_locale = "en-US"
  851. if "ui" in CONFIG_DATA:
  852. default_locale = CONFIG_DATA["ui"].get("default_locale", "en-US")
  853. # The Rest of the Function Now Uses the Variables Defined Above
  854. return {
  855. "status": True,
  856. "name": WEBUI_NAME,
  857. "version": VERSION,
  858. "default_locale": default_locale,
  859. "default_models": webui_app.state.config.DEFAULT_MODELS,
  860. "default_prompt_suggestions": webui_app.state.config.DEFAULT_PROMPT_SUGGESTIONS,
  861. "features": {
  862. "auth": WEBUI_AUTH,
  863. "auth_trusted_header": bool(webui_app.state.AUTH_TRUSTED_EMAIL_HEADER),
  864. "enable_signup": webui_app.state.config.ENABLE_SIGNUP,
  865. "enable_web_search": rag_app.state.config.ENABLE_RAG_WEB_SEARCH,
  866. "enable_image_generation": images_app.state.config.ENABLED,
  867. "enable_community_sharing": webui_app.state.config.ENABLE_COMMUNITY_SHARING,
  868. "enable_admin_export": ENABLE_ADMIN_EXPORT,
  869. },
  870. "audio": {
  871. "tts": {
  872. "engine": audio_app.state.config.TTS_ENGINE,
  873. "voice": audio_app.state.config.TTS_VOICE,
  874. },
  875. "stt": {
  876. "engine": audio_app.state.config.STT_ENGINE,
  877. },
  878. },
  879. }
  880. @app.get("/api/config/model/filter")
  881. async def get_model_filter_config(user=Depends(get_admin_user)):
  882. return {
  883. "enabled": app.state.config.ENABLE_MODEL_FILTER,
  884. "models": app.state.config.MODEL_FILTER_LIST,
  885. }
  886. class ModelFilterConfigForm(BaseModel):
  887. enabled: bool
  888. models: List[str]
  889. @app.post("/api/config/model/filter")
  890. async def update_model_filter_config(
  891. form_data: ModelFilterConfigForm, user=Depends(get_admin_user)
  892. ):
  893. app.state.config.ENABLE_MODEL_FILTER = form_data.enabled
  894. app.state.config.MODEL_FILTER_LIST = form_data.models
  895. return {
  896. "enabled": app.state.config.ENABLE_MODEL_FILTER,
  897. "models": app.state.config.MODEL_FILTER_LIST,
  898. }
  899. @app.get("/api/webhook")
  900. async def get_webhook_url(user=Depends(get_admin_user)):
  901. return {
  902. "url": app.state.config.WEBHOOK_URL,
  903. }
  904. class UrlForm(BaseModel):
  905. url: str
  906. @app.post("/api/webhook")
  907. async def update_webhook_url(form_data: UrlForm, user=Depends(get_admin_user)):
  908. app.state.config.WEBHOOK_URL = form_data.url
  909. webui_app.state.WEBHOOK_URL = app.state.config.WEBHOOK_URL
  910. return {"url": app.state.config.WEBHOOK_URL}
  911. @app.get("/api/version")
  912. async def get_app_config():
  913. return {
  914. "version": VERSION,
  915. }
  916. @app.get("/api/changelog")
  917. async def get_app_changelog():
  918. return {key: CHANGELOG[key] for idx, key in enumerate(CHANGELOG) if idx < 5}
  919. @app.get("/api/version/updates")
  920. async def get_app_latest_release_version():
  921. try:
  922. async with aiohttp.ClientSession(trust_env=True) as session:
  923. async with session.get(
  924. "https://api.github.com/repos/open-webui/open-webui/releases/latest"
  925. ) as response:
  926. response.raise_for_status()
  927. data = await response.json()
  928. latest_version = data["tag_name"]
  929. return {"current": VERSION, "latest": latest_version[1:]}
  930. except aiohttp.ClientError as e:
  931. raise HTTPException(
  932. status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
  933. detail=ERROR_MESSAGES.RATE_LIMIT_EXCEEDED,
  934. )
  935. @app.get("/manifest.json")
  936. async def get_manifest_json():
  937. return {
  938. "name": WEBUI_NAME,
  939. "short_name": WEBUI_NAME,
  940. "start_url": "/",
  941. "display": "standalone",
  942. "background_color": "#343541",
  943. "theme_color": "#343541",
  944. "orientation": "portrait-primary",
  945. "icons": [{"src": "/static/logo.png", "type": "image/png", "sizes": "500x500"}],
  946. }
  947. @app.get("/opensearch.xml")
  948. async def get_opensearch_xml():
  949. xml_content = rf"""
  950. <OpenSearchDescription xmlns="http://a9.com/-/spec/opensearch/1.1/" xmlns:moz="http://www.mozilla.org/2006/browser/search/">
  951. <ShortName>{WEBUI_NAME}</ShortName>
  952. <Description>Search {WEBUI_NAME}</Description>
  953. <InputEncoding>UTF-8</InputEncoding>
  954. <Image width="16" height="16" type="image/x-icon">{WEBUI_URL}/favicon.png</Image>
  955. <Url type="text/html" method="get" template="{WEBUI_URL}/?q={"{searchTerms}"}"/>
  956. <moz:SearchForm>{WEBUI_URL}</moz:SearchForm>
  957. </OpenSearchDescription>
  958. """
  959. return Response(content=xml_content, media_type="application/xml")
  960. @app.get("/health")
  961. async def healthcheck():
  962. return {"status": True}
  963. app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
  964. app.mount("/cache", StaticFiles(directory=CACHE_DIR), name="cache")
  965. if os.path.exists(FRONTEND_BUILD_DIR):
  966. mimetypes.add_type("text/javascript", ".js")
  967. app.mount(
  968. "/",
  969. SPAStaticFiles(directory=FRONTEND_BUILD_DIR, html=True),
  970. name="spa-static-files",
  971. )
  972. else:
  973. log.warning(
  974. f"Frontend build directory not found at '{FRONTEND_BUILD_DIR}'. Serving API only."
  975. )