main.py 37 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231
  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. "SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD": app.state.config.SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD,
  414. }
  415. class TaskConfigForm(BaseModel):
  416. TASK_MODEL: Optional[str]
  417. TASK_MODEL_EXTERNAL: Optional[str]
  418. TITLE_GENERATION_PROMPT_TEMPLATE: str
  419. SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE: str
  420. SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD: int
  421. @app.post("/api/task/config/update")
  422. async def update_task_config(form_data: TaskConfigForm, user=Depends(get_admin_user)):
  423. app.state.config.TASK_MODEL = form_data.TASK_MODEL
  424. app.state.config.TASK_MODEL_EXTERNAL = form_data.TASK_MODEL_EXTERNAL
  425. app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE = (
  426. form_data.TITLE_GENERATION_PROMPT_TEMPLATE
  427. )
  428. app.state.config.SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE = (
  429. form_data.SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE
  430. )
  431. app.state.config.SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD = (
  432. form_data.SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD
  433. )
  434. return {
  435. "TASK_MODEL": app.state.config.TASK_MODEL,
  436. "TASK_MODEL_EXTERNAL": app.state.config.TASK_MODEL_EXTERNAL,
  437. "TITLE_GENERATION_PROMPT_TEMPLATE": app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE,
  438. "SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE": app.state.config.SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE,
  439. "SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD": app.state.config.SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD,
  440. }
  441. @app.post("/api/task/title/completions")
  442. async def generate_title(form_data: dict, user=Depends(get_verified_user)):
  443. print("generate_title")
  444. model_id = form_data["model"]
  445. if model_id not in app.state.MODELS:
  446. raise HTTPException(
  447. status_code=status.HTTP_404_NOT_FOUND,
  448. detail="Model not found",
  449. )
  450. # Check if the user has a custom task model
  451. # If the user has a custom task model, use that model
  452. if app.state.MODELS[model_id]["owned_by"] == "ollama":
  453. if app.state.config.TASK_MODEL:
  454. task_model_id = app.state.config.TASK_MODEL
  455. if task_model_id in app.state.MODELS:
  456. model_id = task_model_id
  457. else:
  458. if app.state.config.TASK_MODEL_EXTERNAL:
  459. task_model_id = app.state.config.TASK_MODEL_EXTERNAL
  460. if task_model_id in app.state.MODELS:
  461. model_id = task_model_id
  462. print(model_id)
  463. model = app.state.MODELS[model_id]
  464. template = app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE
  465. content = title_generation_template(
  466. template, form_data["prompt"], user.model_dump()
  467. )
  468. payload = {
  469. "model": model_id,
  470. "messages": [{"role": "user", "content": content}],
  471. "stream": False,
  472. "max_tokens": 50,
  473. "chat_id": form_data.get("chat_id", None),
  474. "title": True,
  475. }
  476. print(payload)
  477. payload = filter_pipeline(payload, user)
  478. if model["owned_by"] == "ollama":
  479. return await generate_ollama_chat_completion(
  480. OpenAIChatCompletionForm(**payload), user=user
  481. )
  482. else:
  483. return await generate_openai_chat_completion(payload, user=user)
  484. @app.post("/api/task/query/completions")
  485. async def generate_search_query(form_data: dict, user=Depends(get_verified_user)):
  486. print("generate_search_query")
  487. if len(form_data["prompt"]) < app.state.config.SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD:
  488. raise HTTPException(
  489. status_code=status.HTTP_400_BAD_REQUEST,
  490. detail=f"Skip search query generation for short prompts (< {app.state.config.SEARCH_QUERY_PROMPT_LENGTH_THRESHOLD} characters)",
  491. )
  492. model_id = form_data["model"]
  493. if model_id not in app.state.MODELS:
  494. raise HTTPException(
  495. status_code=status.HTTP_404_NOT_FOUND,
  496. detail="Model not found",
  497. )
  498. # Check if the user has a custom task model
  499. # If the user has a custom task model, use that model
  500. if app.state.MODELS[model_id]["owned_by"] == "ollama":
  501. if app.state.config.TASK_MODEL:
  502. task_model_id = app.state.config.TASK_MODEL
  503. if task_model_id in app.state.MODELS:
  504. model_id = task_model_id
  505. else:
  506. if app.state.config.TASK_MODEL_EXTERNAL:
  507. task_model_id = app.state.config.TASK_MODEL_EXTERNAL
  508. if task_model_id in app.state.MODELS:
  509. model_id = task_model_id
  510. print(model_id)
  511. model = app.state.MODELS[model_id]
  512. template = app.state.config.SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE
  513. content = search_query_generation_template(
  514. template, form_data["prompt"], user.model_dump()
  515. )
  516. payload = {
  517. "model": model_id,
  518. "messages": [{"role": "user", "content": content}],
  519. "stream": False,
  520. "max_tokens": 30,
  521. }
  522. print(payload)
  523. payload = filter_pipeline(payload, user)
  524. if model["owned_by"] == "ollama":
  525. return await generate_ollama_chat_completion(
  526. OpenAIChatCompletionForm(**payload), user=user
  527. )
  528. else:
  529. return await generate_openai_chat_completion(payload, user=user)
  530. @app.post("/api/chat/completions")
  531. async def generate_chat_completions(form_data: dict, user=Depends(get_verified_user)):
  532. model_id = form_data["model"]
  533. if model_id not in app.state.MODELS:
  534. raise HTTPException(
  535. status_code=status.HTTP_404_NOT_FOUND,
  536. detail="Model not found",
  537. )
  538. model = app.state.MODELS[model_id]
  539. print(model)
  540. if model["owned_by"] == "ollama":
  541. return await generate_ollama_chat_completion(
  542. OpenAIChatCompletionForm(**form_data), user=user
  543. )
  544. else:
  545. return await generate_openai_chat_completion(form_data, user=user)
  546. @app.post("/api/chat/completed")
  547. async def chat_completed(form_data: dict, user=Depends(get_verified_user)):
  548. data = form_data
  549. model_id = data["model"]
  550. filters = [
  551. model
  552. for model in app.state.MODELS.values()
  553. if "pipeline" in model
  554. and "type" in model["pipeline"]
  555. and model["pipeline"]["type"] == "filter"
  556. and (
  557. model["pipeline"]["pipelines"] == ["*"]
  558. or any(
  559. model_id == target_model_id
  560. for target_model_id in model["pipeline"]["pipelines"]
  561. )
  562. )
  563. ]
  564. sorted_filters = sorted(filters, key=lambda x: x["pipeline"]["priority"])
  565. print(model_id)
  566. if model_id in app.state.MODELS:
  567. model = app.state.MODELS[model_id]
  568. if "pipeline" in model:
  569. sorted_filters = [model] + sorted_filters
  570. for filter in sorted_filters:
  571. r = None
  572. try:
  573. urlIdx = filter["urlIdx"]
  574. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  575. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  576. if key != "":
  577. headers = {"Authorization": f"Bearer {key}"}
  578. r = requests.post(
  579. f"{url}/{filter['id']}/filter/outlet",
  580. headers=headers,
  581. json={
  582. "user": {"id": user.id, "name": user.name, "role": user.role},
  583. "body": data,
  584. },
  585. )
  586. r.raise_for_status()
  587. data = r.json()
  588. except Exception as e:
  589. # Handle connection error here
  590. print(f"Connection error: {e}")
  591. if r is not None:
  592. try:
  593. res = r.json()
  594. if "detail" in res:
  595. return JSONResponse(
  596. status_code=r.status_code,
  597. content=res,
  598. )
  599. except:
  600. pass
  601. else:
  602. pass
  603. return data
  604. @app.get("/api/pipelines/list")
  605. async def get_pipelines_list(user=Depends(get_admin_user)):
  606. responses = await get_openai_models(raw=True)
  607. print(responses)
  608. urlIdxs = [
  609. idx
  610. for idx, response in enumerate(responses)
  611. if response != None and "pipelines" in response
  612. ]
  613. return {
  614. "data": [
  615. {
  616. "url": openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx],
  617. "idx": urlIdx,
  618. }
  619. for urlIdx in urlIdxs
  620. ]
  621. }
  622. @app.post("/api/pipelines/upload")
  623. async def upload_pipeline(
  624. urlIdx: int = Form(...), file: UploadFile = File(...), user=Depends(get_admin_user)
  625. ):
  626. print("upload_pipeline", urlIdx, file.filename)
  627. # Check if the uploaded file is a python file
  628. if not file.filename.endswith(".py"):
  629. raise HTTPException(
  630. status_code=status.HTTP_400_BAD_REQUEST,
  631. detail="Only Python (.py) files are allowed.",
  632. )
  633. upload_folder = f"{CACHE_DIR}/pipelines"
  634. os.makedirs(upload_folder, exist_ok=True)
  635. file_path = os.path.join(upload_folder, file.filename)
  636. try:
  637. # Save the uploaded file
  638. with open(file_path, "wb") as buffer:
  639. shutil.copyfileobj(file.file, buffer)
  640. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  641. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  642. headers = {"Authorization": f"Bearer {key}"}
  643. with open(file_path, "rb") as f:
  644. files = {"file": f}
  645. r = requests.post(f"{url}/pipelines/upload", headers=headers, files=files)
  646. r.raise_for_status()
  647. data = r.json()
  648. return {**data}
  649. except Exception as e:
  650. # Handle connection error here
  651. print(f"Connection error: {e}")
  652. detail = "Pipeline not found"
  653. if r is not None:
  654. try:
  655. res = r.json()
  656. if "detail" in res:
  657. detail = res["detail"]
  658. except:
  659. pass
  660. raise HTTPException(
  661. status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
  662. detail=detail,
  663. )
  664. finally:
  665. # Ensure the file is deleted after the upload is completed or on failure
  666. if os.path.exists(file_path):
  667. os.remove(file_path)
  668. class AddPipelineForm(BaseModel):
  669. url: str
  670. urlIdx: int
  671. @app.post("/api/pipelines/add")
  672. async def add_pipeline(form_data: AddPipelineForm, user=Depends(get_admin_user)):
  673. r = None
  674. try:
  675. urlIdx = form_data.urlIdx
  676. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  677. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  678. headers = {"Authorization": f"Bearer {key}"}
  679. r = requests.post(
  680. f"{url}/pipelines/add", headers=headers, json={"url": form_data.url}
  681. )
  682. r.raise_for_status()
  683. data = r.json()
  684. return {**data}
  685. except Exception as e:
  686. # Handle connection error here
  687. print(f"Connection error: {e}")
  688. detail = "Pipeline not found"
  689. if r is not None:
  690. try:
  691. res = r.json()
  692. if "detail" in res:
  693. detail = res["detail"]
  694. except:
  695. pass
  696. raise HTTPException(
  697. status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
  698. detail=detail,
  699. )
  700. class DeletePipelineForm(BaseModel):
  701. id: str
  702. urlIdx: int
  703. @app.delete("/api/pipelines/delete")
  704. async def delete_pipeline(form_data: DeletePipelineForm, user=Depends(get_admin_user)):
  705. r = None
  706. try:
  707. urlIdx = form_data.urlIdx
  708. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  709. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  710. headers = {"Authorization": f"Bearer {key}"}
  711. r = requests.delete(
  712. f"{url}/pipelines/delete", headers=headers, json={"id": form_data.id}
  713. )
  714. r.raise_for_status()
  715. data = r.json()
  716. return {**data}
  717. except Exception as e:
  718. # Handle connection error here
  719. print(f"Connection error: {e}")
  720. detail = "Pipeline not found"
  721. if r is not None:
  722. try:
  723. res = r.json()
  724. if "detail" in res:
  725. detail = res["detail"]
  726. except:
  727. pass
  728. raise HTTPException(
  729. status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
  730. detail=detail,
  731. )
  732. @app.get("/api/pipelines")
  733. async def get_pipelines(urlIdx: Optional[int] = None, user=Depends(get_admin_user)):
  734. r = None
  735. try:
  736. urlIdx
  737. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  738. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  739. headers = {"Authorization": f"Bearer {key}"}
  740. r = requests.get(f"{url}/pipelines", headers=headers)
  741. r.raise_for_status()
  742. data = r.json()
  743. return {**data}
  744. except Exception as e:
  745. # Handle connection error here
  746. print(f"Connection error: {e}")
  747. detail = "Pipeline not found"
  748. if r is not None:
  749. try:
  750. res = r.json()
  751. if "detail" in res:
  752. detail = res["detail"]
  753. except:
  754. pass
  755. raise HTTPException(
  756. status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
  757. detail=detail,
  758. )
  759. @app.get("/api/pipelines/{pipeline_id}/valves")
  760. async def get_pipeline_valves(
  761. urlIdx: Optional[int], pipeline_id: str, user=Depends(get_admin_user)
  762. ):
  763. models = await get_all_models()
  764. r = None
  765. try:
  766. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  767. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  768. headers = {"Authorization": f"Bearer {key}"}
  769. r = requests.get(f"{url}/{pipeline_id}/valves", headers=headers)
  770. r.raise_for_status()
  771. data = r.json()
  772. return {**data}
  773. except Exception as e:
  774. # Handle connection error here
  775. print(f"Connection error: {e}")
  776. detail = "Pipeline not found"
  777. if r is not None:
  778. try:
  779. res = r.json()
  780. if "detail" in res:
  781. detail = res["detail"]
  782. except:
  783. pass
  784. raise HTTPException(
  785. status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
  786. detail=detail,
  787. )
  788. @app.get("/api/pipelines/{pipeline_id}/valves/spec")
  789. async def get_pipeline_valves_spec(
  790. urlIdx: Optional[int], pipeline_id: str, user=Depends(get_admin_user)
  791. ):
  792. models = await get_all_models()
  793. r = None
  794. try:
  795. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  796. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  797. headers = {"Authorization": f"Bearer {key}"}
  798. r = requests.get(f"{url}/{pipeline_id}/valves/spec", headers=headers)
  799. r.raise_for_status()
  800. data = r.json()
  801. return {**data}
  802. except Exception as e:
  803. # Handle connection error here
  804. print(f"Connection error: {e}")
  805. detail = "Pipeline not found"
  806. if r is not None:
  807. try:
  808. res = r.json()
  809. if "detail" in res:
  810. detail = res["detail"]
  811. except:
  812. pass
  813. raise HTTPException(
  814. status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
  815. detail=detail,
  816. )
  817. @app.post("/api/pipelines/{pipeline_id}/valves/update")
  818. async def update_pipeline_valves(
  819. urlIdx: Optional[int],
  820. pipeline_id: str,
  821. form_data: dict,
  822. user=Depends(get_admin_user),
  823. ):
  824. models = await get_all_models()
  825. r = None
  826. try:
  827. url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
  828. key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
  829. headers = {"Authorization": f"Bearer {key}"}
  830. r = requests.post(
  831. f"{url}/{pipeline_id}/valves/update",
  832. headers=headers,
  833. json={**form_data},
  834. )
  835. r.raise_for_status()
  836. data = r.json()
  837. return {**data}
  838. except Exception as e:
  839. # Handle connection error here
  840. print(f"Connection error: {e}")
  841. detail = "Pipeline not found"
  842. if r is not None:
  843. try:
  844. res = r.json()
  845. if "detail" in res:
  846. detail = res["detail"]
  847. except:
  848. pass
  849. raise HTTPException(
  850. status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
  851. detail=detail,
  852. )
  853. @app.get("/api/config")
  854. async def get_app_config():
  855. # Checking and Handling the Absence of 'ui' in CONFIG_DATA
  856. default_locale = "en-US"
  857. if "ui" in CONFIG_DATA:
  858. default_locale = CONFIG_DATA["ui"].get("default_locale", "en-US")
  859. # The Rest of the Function Now Uses the Variables Defined Above
  860. return {
  861. "status": True,
  862. "name": WEBUI_NAME,
  863. "version": VERSION,
  864. "default_locale": default_locale,
  865. "default_models": webui_app.state.config.DEFAULT_MODELS,
  866. "default_prompt_suggestions": webui_app.state.config.DEFAULT_PROMPT_SUGGESTIONS,
  867. "features": {
  868. "auth": WEBUI_AUTH,
  869. "auth_trusted_header": bool(webui_app.state.AUTH_TRUSTED_EMAIL_HEADER),
  870. "enable_signup": webui_app.state.config.ENABLE_SIGNUP,
  871. "enable_web_search": rag_app.state.config.ENABLE_RAG_WEB_SEARCH,
  872. "enable_image_generation": images_app.state.config.ENABLED,
  873. "enable_community_sharing": webui_app.state.config.ENABLE_COMMUNITY_SHARING,
  874. "enable_admin_export": ENABLE_ADMIN_EXPORT,
  875. },
  876. "audio": {
  877. "tts": {
  878. "engine": audio_app.state.config.TTS_ENGINE,
  879. "voice": audio_app.state.config.TTS_VOICE,
  880. },
  881. "stt": {
  882. "engine": audio_app.state.config.STT_ENGINE,
  883. },
  884. },
  885. }
  886. @app.get("/api/config/model/filter")
  887. async def get_model_filter_config(user=Depends(get_admin_user)):
  888. return {
  889. "enabled": app.state.config.ENABLE_MODEL_FILTER,
  890. "models": app.state.config.MODEL_FILTER_LIST,
  891. }
  892. class ModelFilterConfigForm(BaseModel):
  893. enabled: bool
  894. models: List[str]
  895. @app.post("/api/config/model/filter")
  896. async def update_model_filter_config(
  897. form_data: ModelFilterConfigForm, user=Depends(get_admin_user)
  898. ):
  899. app.state.config.ENABLE_MODEL_FILTER = form_data.enabled
  900. app.state.config.MODEL_FILTER_LIST = form_data.models
  901. return {
  902. "enabled": app.state.config.ENABLE_MODEL_FILTER,
  903. "models": app.state.config.MODEL_FILTER_LIST,
  904. }
  905. @app.get("/api/webhook")
  906. async def get_webhook_url(user=Depends(get_admin_user)):
  907. return {
  908. "url": app.state.config.WEBHOOK_URL,
  909. }
  910. class UrlForm(BaseModel):
  911. url: str
  912. @app.post("/api/webhook")
  913. async def update_webhook_url(form_data: UrlForm, user=Depends(get_admin_user)):
  914. app.state.config.WEBHOOK_URL = form_data.url
  915. webui_app.state.WEBHOOK_URL = app.state.config.WEBHOOK_URL
  916. return {"url": app.state.config.WEBHOOK_URL}
  917. @app.get("/api/version")
  918. async def get_app_config():
  919. return {
  920. "version": VERSION,
  921. }
  922. @app.get("/api/changelog")
  923. async def get_app_changelog():
  924. return {key: CHANGELOG[key] for idx, key in enumerate(CHANGELOG) if idx < 5}
  925. @app.get("/api/version/updates")
  926. async def get_app_latest_release_version():
  927. try:
  928. async with aiohttp.ClientSession(trust_env=True) as session:
  929. async with session.get(
  930. "https://api.github.com/repos/open-webui/open-webui/releases/latest"
  931. ) as response:
  932. response.raise_for_status()
  933. data = await response.json()
  934. latest_version = data["tag_name"]
  935. return {"current": VERSION, "latest": latest_version[1:]}
  936. except aiohttp.ClientError as e:
  937. raise HTTPException(
  938. status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
  939. detail=ERROR_MESSAGES.RATE_LIMIT_EXCEEDED,
  940. )
  941. @app.get("/manifest.json")
  942. async def get_manifest_json():
  943. return {
  944. "name": WEBUI_NAME,
  945. "short_name": WEBUI_NAME,
  946. "start_url": "/",
  947. "display": "standalone",
  948. "background_color": "#343541",
  949. "theme_color": "#343541",
  950. "orientation": "portrait-primary",
  951. "icons": [{"src": "/static/logo.png", "type": "image/png", "sizes": "500x500"}],
  952. }
  953. @app.get("/opensearch.xml")
  954. async def get_opensearch_xml():
  955. xml_content = rf"""
  956. <OpenSearchDescription xmlns="http://a9.com/-/spec/opensearch/1.1/" xmlns:moz="http://www.mozilla.org/2006/browser/search/">
  957. <ShortName>{WEBUI_NAME}</ShortName>
  958. <Description>Search {WEBUI_NAME}</Description>
  959. <InputEncoding>UTF-8</InputEncoding>
  960. <Image width="16" height="16" type="image/x-icon">{WEBUI_URL}/favicon.png</Image>
  961. <Url type="text/html" method="get" template="{WEBUI_URL}/?q={"{searchTerms}"}"/>
  962. <moz:SearchForm>{WEBUI_URL}</moz:SearchForm>
  963. </OpenSearchDescription>
  964. """
  965. return Response(content=xml_content, media_type="application/xml")
  966. @app.get("/health")
  967. async def healthcheck():
  968. return {"status": True}
  969. app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
  970. app.mount("/cache", StaticFiles(directory=CACHE_DIR), name="cache")
  971. if os.path.exists(FRONTEND_BUILD_DIR):
  972. mimetypes.add_type("text/javascript", ".js")
  973. app.mount(
  974. "/",
  975. SPAStaticFiles(directory=FRONTEND_BUILD_DIR, html=True),
  976. name="spa-static-files",
  977. )
  978. else:
  979. log.warning(
  980. f"Frontend build directory not found at '{FRONTEND_BUILD_DIR}'. Serving API only."
  981. )