main.py 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996
  1. from fastapi import FastAPI, Request, Response, HTTPException, Depends, status
  2. from fastapi.middleware.cors import CORSMiddleware
  3. from fastapi.responses import StreamingResponse
  4. from fastapi.concurrency import run_in_threadpool
  5. from pydantic import BaseModel, ConfigDict
  6. import random
  7. import requests
  8. import json
  9. import uuid
  10. import aiohttp
  11. import asyncio
  12. import logging
  13. from apps.web.models.users import Users
  14. from constants import ERROR_MESSAGES
  15. from utils.utils import decode_token, get_current_user, get_admin_user
  16. from config import SRC_LOG_LEVELS, OLLAMA_BASE_URLS, MODEL_FILTER_ENABLED, MODEL_FILTER_LIST
  17. from typing import Optional, List, Union
  18. log = logging.getLogger(__name__)
  19. log.setLevel(SRC_LOG_LEVELS["OLLAMA"])
  20. app = FastAPI()
  21. app.add_middleware(
  22. CORSMiddleware,
  23. allow_origins=["*"],
  24. allow_credentials=True,
  25. allow_methods=["*"],
  26. allow_headers=["*"],
  27. )
  28. app.state.MODEL_FILTER_ENABLED = MODEL_FILTER_ENABLED
  29. app.state.MODEL_FILTER_LIST = MODEL_FILTER_LIST
  30. app.state.OLLAMA_BASE_URLS = OLLAMA_BASE_URLS
  31. app.state.MODELS = {}
  32. REQUEST_POOL = []
  33. # TODO: Implement a more intelligent load balancing mechanism for distributing requests among multiple backend instances.
  34. # Current implementation uses a simple round-robin approach (random.choice). Consider incorporating algorithms like weighted round-robin,
  35. # least connections, or least response time for better resource utilization and performance optimization.
  36. @app.middleware("http")
  37. async def check_url(request: Request, call_next):
  38. if len(app.state.MODELS) == 0:
  39. await get_all_models()
  40. else:
  41. pass
  42. response = await call_next(request)
  43. return response
  44. @app.get("/urls")
  45. async def get_ollama_api_urls(user=Depends(get_admin_user)):
  46. return {"OLLAMA_BASE_URLS": app.state.OLLAMA_BASE_URLS}
  47. class UrlUpdateForm(BaseModel):
  48. urls: List[str]
  49. @app.post("/urls/update")
  50. async def update_ollama_api_url(form_data: UrlUpdateForm, user=Depends(get_admin_user)):
  51. app.state.OLLAMA_BASE_URLS = form_data.urls
  52. log.info(f"app.state.OLLAMA_BASE_URLS: {app.state.OLLAMA_BASE_URLS}")
  53. return {"OLLAMA_BASE_URLS": app.state.OLLAMA_BASE_URLS}
  54. @app.get("/cancel/{request_id}")
  55. async def cancel_ollama_request(request_id: str, user=Depends(get_current_user)):
  56. if user:
  57. if request_id in REQUEST_POOL:
  58. REQUEST_POOL.remove(request_id)
  59. return True
  60. else:
  61. raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
  62. async def fetch_url(url):
  63. try:
  64. async with aiohttp.ClientSession() as session:
  65. async with session.get(url) as response:
  66. return await response.json()
  67. except Exception as e:
  68. # Handle connection error here
  69. log.error(f"Connection error: {e}")
  70. return None
  71. def merge_models_lists(model_lists):
  72. merged_models = {}
  73. for idx, model_list in enumerate(model_lists):
  74. if model_list is not None:
  75. for model in model_list:
  76. digest = model["digest"]
  77. if digest not in merged_models:
  78. model["urls"] = [idx]
  79. merged_models[digest] = model
  80. else:
  81. merged_models[digest]["urls"].append(idx)
  82. return list(merged_models.values())
  83. # user=Depends(get_current_user)
  84. async def get_all_models():
  85. log.info("get_all_models()")
  86. tasks = [fetch_url(f"{url}/api/tags") for url in app.state.OLLAMA_BASE_URLS]
  87. responses = await asyncio.gather(*tasks)
  88. models = {
  89. "models": merge_models_lists(
  90. map(lambda response: response["models"] if response else None, responses)
  91. )
  92. }
  93. app.state.MODELS = {model["model"]: model for model in models["models"]}
  94. return models
  95. @app.get("/api/tags")
  96. @app.get("/api/tags/{url_idx}")
  97. async def get_ollama_tags(
  98. url_idx: Optional[int] = None, user=Depends(get_current_user)
  99. ):
  100. if url_idx == None:
  101. models = await get_all_models()
  102. if app.state.MODEL_FILTER_ENABLED:
  103. if user.role == "user":
  104. models["models"] = list(
  105. filter(
  106. lambda model: model["name"] in app.state.MODEL_FILTER_LIST,
  107. models["models"],
  108. )
  109. )
  110. return models
  111. return models
  112. else:
  113. url = app.state.OLLAMA_BASE_URLS[url_idx]
  114. try:
  115. r = requests.request(method="GET", url=f"{url}/api/tags")
  116. r.raise_for_status()
  117. return r.json()
  118. except Exception as e:
  119. log.exception(e)
  120. error_detail = "Open WebUI: Server Connection Error"
  121. if r is not None:
  122. try:
  123. res = r.json()
  124. if "error" in res:
  125. error_detail = f"Ollama: {res['error']}"
  126. except:
  127. error_detail = f"Ollama: {e}"
  128. raise HTTPException(
  129. status_code=r.status_code if r else 500,
  130. detail=error_detail,
  131. )
  132. @app.get("/api/version")
  133. @app.get("/api/version/{url_idx}")
  134. async def get_ollama_versions(url_idx: Optional[int] = None):
  135. if url_idx == None:
  136. # returns lowest version
  137. tasks = [fetch_url(f"{url}/api/version") for url in app.state.OLLAMA_BASE_URLS]
  138. responses = await asyncio.gather(*tasks)
  139. responses = list(filter(lambda x: x is not None, responses))
  140. if len(responses) > 0:
  141. lowest_version = min(
  142. responses, key=lambda x: tuple(map(int, x["version"].split(".")))
  143. )
  144. return {"version": lowest_version["version"]}
  145. else:
  146. raise HTTPException(
  147. status_code=500,
  148. detail=ERROR_MESSAGES.OLLAMA_NOT_FOUND,
  149. )
  150. else:
  151. url = app.state.OLLAMA_BASE_URLS[url_idx]
  152. try:
  153. r = requests.request(method="GET", url=f"{url}/api/version")
  154. r.raise_for_status()
  155. return r.json()
  156. except Exception as e:
  157. log.exception(e)
  158. error_detail = "Open WebUI: Server Connection Error"
  159. if r is not None:
  160. try:
  161. res = r.json()
  162. if "error" in res:
  163. error_detail = f"Ollama: {res['error']}"
  164. except:
  165. error_detail = f"Ollama: {e}"
  166. raise HTTPException(
  167. status_code=r.status_code if r else 500,
  168. detail=error_detail,
  169. )
  170. class ModelNameForm(BaseModel):
  171. name: str
  172. @app.post("/api/pull")
  173. @app.post("/api/pull/{url_idx}")
  174. async def pull_model(
  175. form_data: ModelNameForm, url_idx: int = 0, user=Depends(get_admin_user)
  176. ):
  177. url = app.state.OLLAMA_BASE_URLS[url_idx]
  178. log.info(f"url: {url}")
  179. r = None
  180. def get_request():
  181. nonlocal url
  182. nonlocal r
  183. try:
  184. def stream_content():
  185. for chunk in r.iter_content(chunk_size=8192):
  186. yield chunk
  187. r = requests.request(
  188. method="POST",
  189. url=f"{url}/api/pull",
  190. data=form_data.model_dump_json(exclude_none=True).encode(),
  191. stream=True,
  192. )
  193. r.raise_for_status()
  194. return StreamingResponse(
  195. stream_content(),
  196. status_code=r.status_code,
  197. headers=dict(r.headers),
  198. )
  199. except Exception as e:
  200. raise e
  201. try:
  202. return await run_in_threadpool(get_request)
  203. except Exception as e:
  204. log.exception(e)
  205. error_detail = "Open WebUI: Server Connection Error"
  206. if r is not None:
  207. try:
  208. res = r.json()
  209. if "error" in res:
  210. error_detail = f"Ollama: {res['error']}"
  211. except:
  212. error_detail = f"Ollama: {e}"
  213. raise HTTPException(
  214. status_code=r.status_code if r else 500,
  215. detail=error_detail,
  216. )
  217. class PushModelForm(BaseModel):
  218. name: str
  219. insecure: Optional[bool] = None
  220. stream: Optional[bool] = None
  221. @app.delete("/api/push")
  222. @app.delete("/api/push/{url_idx}")
  223. async def push_model(
  224. form_data: PushModelForm,
  225. url_idx: Optional[int] = None,
  226. user=Depends(get_admin_user),
  227. ):
  228. if url_idx == None:
  229. if form_data.name in app.state.MODELS:
  230. url_idx = app.state.MODELS[form_data.name]["urls"][0]
  231. else:
  232. raise HTTPException(
  233. status_code=400,
  234. detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.name),
  235. )
  236. url = app.state.OLLAMA_BASE_URLS[url_idx]
  237. log.debug(f"url: {url}")
  238. r = None
  239. def get_request():
  240. nonlocal url
  241. nonlocal r
  242. try:
  243. def stream_content():
  244. for chunk in r.iter_content(chunk_size=8192):
  245. yield chunk
  246. r = requests.request(
  247. method="POST",
  248. url=f"{url}/api/push",
  249. data=form_data.model_dump_json(exclude_none=True).encode(),
  250. )
  251. r.raise_for_status()
  252. return StreamingResponse(
  253. stream_content(),
  254. status_code=r.status_code,
  255. headers=dict(r.headers),
  256. )
  257. except Exception as e:
  258. raise e
  259. try:
  260. return await run_in_threadpool(get_request)
  261. except Exception as e:
  262. log.exception(e)
  263. error_detail = "Open WebUI: Server Connection Error"
  264. if r is not None:
  265. try:
  266. res = r.json()
  267. if "error" in res:
  268. error_detail = f"Ollama: {res['error']}"
  269. except:
  270. error_detail = f"Ollama: {e}"
  271. raise HTTPException(
  272. status_code=r.status_code if r else 500,
  273. detail=error_detail,
  274. )
  275. class CreateModelForm(BaseModel):
  276. name: str
  277. modelfile: Optional[str] = None
  278. stream: Optional[bool] = None
  279. path: Optional[str] = None
  280. @app.post("/api/create")
  281. @app.post("/api/create/{url_idx}")
  282. async def create_model(
  283. form_data: CreateModelForm, url_idx: int = 0, user=Depends(get_admin_user)
  284. ):
  285. log.debug(f"form_data: {form_data}")
  286. url = app.state.OLLAMA_BASE_URLS[url_idx]
  287. log.info(f"url: {url}")
  288. r = None
  289. def get_request():
  290. nonlocal url
  291. nonlocal r
  292. try:
  293. def stream_content():
  294. for chunk in r.iter_content(chunk_size=8192):
  295. yield chunk
  296. r = requests.request(
  297. method="POST",
  298. url=f"{url}/api/create",
  299. data=form_data.model_dump_json(exclude_none=True).encode(),
  300. stream=True,
  301. )
  302. r.raise_for_status()
  303. log.debug(f"r: {r}")
  304. return StreamingResponse(
  305. stream_content(),
  306. status_code=r.status_code,
  307. headers=dict(r.headers),
  308. )
  309. except Exception as e:
  310. raise e
  311. try:
  312. return await run_in_threadpool(get_request)
  313. except Exception as e:
  314. log.exception(e)
  315. error_detail = "Open WebUI: Server Connection Error"
  316. if r is not None:
  317. try:
  318. res = r.json()
  319. if "error" in res:
  320. error_detail = f"Ollama: {res['error']}"
  321. except:
  322. error_detail = f"Ollama: {e}"
  323. raise HTTPException(
  324. status_code=r.status_code if r else 500,
  325. detail=error_detail,
  326. )
  327. class CopyModelForm(BaseModel):
  328. source: str
  329. destination: str
  330. @app.post("/api/copy")
  331. @app.post("/api/copy/{url_idx}")
  332. async def copy_model(
  333. form_data: CopyModelForm,
  334. url_idx: Optional[int] = None,
  335. user=Depends(get_admin_user),
  336. ):
  337. if url_idx == None:
  338. if form_data.source in app.state.MODELS:
  339. url_idx = app.state.MODELS[form_data.source]["urls"][0]
  340. else:
  341. raise HTTPException(
  342. status_code=400,
  343. detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.source),
  344. )
  345. url = app.state.OLLAMA_BASE_URLS[url_idx]
  346. log.info(f"url: {url}")
  347. try:
  348. r = requests.request(
  349. method="POST",
  350. url=f"{url}/api/copy",
  351. data=form_data.model_dump_json(exclude_none=True).encode(),
  352. )
  353. r.raise_for_status()
  354. log.debug(f"r.text: {r.text}")
  355. return True
  356. except Exception as e:
  357. log.exception(e)
  358. error_detail = "Open WebUI: Server Connection Error"
  359. if r is not None:
  360. try:
  361. res = r.json()
  362. if "error" in res:
  363. error_detail = f"Ollama: {res['error']}"
  364. except:
  365. error_detail = f"Ollama: {e}"
  366. raise HTTPException(
  367. status_code=r.status_code if r else 500,
  368. detail=error_detail,
  369. )
  370. @app.delete("/api/delete")
  371. @app.delete("/api/delete/{url_idx}")
  372. async def delete_model(
  373. form_data: ModelNameForm,
  374. url_idx: Optional[int] = None,
  375. user=Depends(get_admin_user),
  376. ):
  377. if url_idx == None:
  378. if form_data.name in app.state.MODELS:
  379. url_idx = app.state.MODELS[form_data.name]["urls"][0]
  380. else:
  381. raise HTTPException(
  382. status_code=400,
  383. detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.name),
  384. )
  385. url = app.state.OLLAMA_BASE_URLS[url_idx]
  386. log.info(f"url: {url}")
  387. try:
  388. r = requests.request(
  389. method="DELETE",
  390. url=f"{url}/api/delete",
  391. data=form_data.model_dump_json(exclude_none=True).encode(),
  392. )
  393. r.raise_for_status()
  394. log.debug(f"r.text: {r.text}")
  395. return True
  396. except Exception as e:
  397. log.exception(e)
  398. error_detail = "Open WebUI: Server Connection Error"
  399. if r is not None:
  400. try:
  401. res = r.json()
  402. if "error" in res:
  403. error_detail = f"Ollama: {res['error']}"
  404. except:
  405. error_detail = f"Ollama: {e}"
  406. raise HTTPException(
  407. status_code=r.status_code if r else 500,
  408. detail=error_detail,
  409. )
  410. @app.post("/api/show")
  411. async def show_model_info(form_data: ModelNameForm, user=Depends(get_current_user)):
  412. if form_data.name not in app.state.MODELS:
  413. raise HTTPException(
  414. status_code=400,
  415. detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.name),
  416. )
  417. url_idx = random.choice(app.state.MODELS[form_data.name]["urls"])
  418. url = app.state.OLLAMA_BASE_URLS[url_idx]
  419. log.info(f"url: {url}")
  420. try:
  421. r = requests.request(
  422. method="POST",
  423. url=f"{url}/api/show",
  424. data=form_data.model_dump_json(exclude_none=True).encode(),
  425. )
  426. r.raise_for_status()
  427. return r.json()
  428. except Exception as e:
  429. log.exception(e)
  430. error_detail = "Open WebUI: Server Connection Error"
  431. if r is not None:
  432. try:
  433. res = r.json()
  434. if "error" in res:
  435. error_detail = f"Ollama: {res['error']}"
  436. except:
  437. error_detail = f"Ollama: {e}"
  438. raise HTTPException(
  439. status_code=r.status_code if r else 500,
  440. detail=error_detail,
  441. )
  442. class GenerateEmbeddingsForm(BaseModel):
  443. model: str
  444. prompt: str
  445. options: Optional[dict] = None
  446. keep_alive: Optional[Union[int, str]] = None
  447. @app.post("/api/embeddings")
  448. @app.post("/api/embeddings/{url_idx}")
  449. async def generate_embeddings(
  450. form_data: GenerateEmbeddingsForm,
  451. url_idx: Optional[int] = None,
  452. user=Depends(get_current_user),
  453. ):
  454. if url_idx == None:
  455. if form_data.model in app.state.MODELS:
  456. url_idx = random.choice(app.state.MODELS[form_data.model]["urls"])
  457. else:
  458. raise HTTPException(
  459. status_code=400,
  460. detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.model),
  461. )
  462. url = app.state.OLLAMA_BASE_URLS[url_idx]
  463. log.info(f"url: {url}")
  464. try:
  465. r = requests.request(
  466. method="POST",
  467. url=f"{url}/api/embeddings",
  468. data=form_data.model_dump_json(exclude_none=True).encode(),
  469. )
  470. r.raise_for_status()
  471. return r.json()
  472. except Exception as e:
  473. log.exception(e)
  474. error_detail = "Open WebUI: Server Connection Error"
  475. if r is not None:
  476. try:
  477. res = r.json()
  478. if "error" in res:
  479. error_detail = f"Ollama: {res['error']}"
  480. except:
  481. error_detail = f"Ollama: {e}"
  482. raise HTTPException(
  483. status_code=r.status_code if r else 500,
  484. detail=error_detail,
  485. )
  486. class GenerateCompletionForm(BaseModel):
  487. model: str
  488. prompt: str
  489. images: Optional[List[str]] = None
  490. format: Optional[str] = None
  491. options: Optional[dict] = None
  492. system: Optional[str] = None
  493. template: Optional[str] = None
  494. context: Optional[str] = None
  495. stream: Optional[bool] = True
  496. raw: Optional[bool] = None
  497. keep_alive: Optional[Union[int, str]] = None
  498. @app.post("/api/generate")
  499. @app.post("/api/generate/{url_idx}")
  500. async def generate_completion(
  501. form_data: GenerateCompletionForm,
  502. url_idx: Optional[int] = None,
  503. user=Depends(get_current_user),
  504. ):
  505. if url_idx == None:
  506. if form_data.model in app.state.MODELS:
  507. url_idx = random.choice(app.state.MODELS[form_data.model]["urls"])
  508. else:
  509. raise HTTPException(
  510. status_code=400,
  511. detail="error_detail",
  512. )
  513. url = app.state.OLLAMA_BASE_URLS[url_idx]
  514. log.info(f"url: {url}")
  515. r = None
  516. def get_request():
  517. nonlocal form_data
  518. nonlocal r
  519. request_id = str(uuid.uuid4())
  520. try:
  521. REQUEST_POOL.append(request_id)
  522. def stream_content():
  523. try:
  524. if form_data.stream:
  525. yield json.dumps({"id": request_id, "done": False}) + "\n"
  526. for chunk in r.iter_content(chunk_size=8192):
  527. if request_id in REQUEST_POOL:
  528. yield chunk
  529. else:
  530. log.warning("User: canceled request")
  531. break
  532. finally:
  533. if hasattr(r, "close"):
  534. r.close()
  535. if request_id in REQUEST_POOL:
  536. REQUEST_POOL.remove(request_id)
  537. r = requests.request(
  538. method="POST",
  539. url=f"{url}/api/generate",
  540. data=form_data.model_dump_json(exclude_none=True).encode(),
  541. stream=True,
  542. )
  543. r.raise_for_status()
  544. return StreamingResponse(
  545. stream_content(),
  546. status_code=r.status_code,
  547. headers=dict(r.headers),
  548. )
  549. except Exception as e:
  550. raise e
  551. try:
  552. return await run_in_threadpool(get_request)
  553. except Exception as e:
  554. error_detail = "Open WebUI: Server Connection Error"
  555. if r is not None:
  556. try:
  557. res = r.json()
  558. if "error" in res:
  559. error_detail = f"Ollama: {res['error']}"
  560. except:
  561. error_detail = f"Ollama: {e}"
  562. raise HTTPException(
  563. status_code=r.status_code if r else 500,
  564. detail=error_detail,
  565. )
  566. class ChatMessage(BaseModel):
  567. role: str
  568. content: str
  569. images: Optional[List[str]] = None
  570. class GenerateChatCompletionForm(BaseModel):
  571. model: str
  572. messages: List[ChatMessage]
  573. format: Optional[str] = None
  574. options: Optional[dict] = None
  575. template: Optional[str] = None
  576. stream: Optional[bool] = True
  577. keep_alive: Optional[Union[int, str]] = None
  578. @app.post("/api/chat")
  579. @app.post("/api/chat/{url_idx}")
  580. async def generate_chat_completion(
  581. form_data: GenerateChatCompletionForm,
  582. url_idx: Optional[int] = None,
  583. user=Depends(get_current_user),
  584. ):
  585. if url_idx == None:
  586. if form_data.model in app.state.MODELS:
  587. url_idx = random.choice(app.state.MODELS[form_data.model]["urls"])
  588. else:
  589. raise HTTPException(
  590. status_code=400,
  591. detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.model),
  592. )
  593. url = app.state.OLLAMA_BASE_URLS[url_idx]
  594. log.info(f"url: {url}")
  595. r = None
  596. log.debug("form_data.model_dump_json(exclude_none=True).encode(): {0} ".format(form_data.model_dump_json(exclude_none=True).encode()))
  597. def get_request():
  598. nonlocal form_data
  599. nonlocal r
  600. request_id = str(uuid.uuid4())
  601. try:
  602. REQUEST_POOL.append(request_id)
  603. def stream_content():
  604. try:
  605. if form_data.stream:
  606. yield json.dumps({"id": request_id, "done": False}) + "\n"
  607. for chunk in r.iter_content(chunk_size=8192):
  608. if request_id in REQUEST_POOL:
  609. yield chunk
  610. else:
  611. log.warning("User: canceled request")
  612. break
  613. finally:
  614. if hasattr(r, "close"):
  615. r.close()
  616. if request_id in REQUEST_POOL:
  617. REQUEST_POOL.remove(request_id)
  618. r = requests.request(
  619. method="POST",
  620. url=f"{url}/api/chat",
  621. data=form_data.model_dump_json(exclude_none=True).encode(),
  622. stream=True,
  623. )
  624. r.raise_for_status()
  625. return StreamingResponse(
  626. stream_content(),
  627. status_code=r.status_code,
  628. headers=dict(r.headers),
  629. )
  630. except Exception as e:
  631. log.exception(e)
  632. raise e
  633. try:
  634. return await run_in_threadpool(get_request)
  635. except Exception as e:
  636. error_detail = "Open WebUI: Server Connection Error"
  637. if r is not None:
  638. try:
  639. res = r.json()
  640. if "error" in res:
  641. error_detail = f"Ollama: {res['error']}"
  642. except:
  643. error_detail = f"Ollama: {e}"
  644. raise HTTPException(
  645. status_code=r.status_code if r else 500,
  646. detail=error_detail,
  647. )
  648. # TODO: we should update this part once Ollama supports other types
  649. class OpenAIChatMessage(BaseModel):
  650. role: str
  651. content: str
  652. model_config = ConfigDict(extra="allow")
  653. class OpenAIChatCompletionForm(BaseModel):
  654. model: str
  655. messages: List[OpenAIChatMessage]
  656. model_config = ConfigDict(extra="allow")
  657. @app.post("/v1/chat/completions")
  658. @app.post("/v1/chat/completions/{url_idx}")
  659. async def generate_openai_chat_completion(
  660. form_data: OpenAIChatCompletionForm,
  661. url_idx: Optional[int] = None,
  662. user=Depends(get_current_user),
  663. ):
  664. if url_idx == None:
  665. if form_data.model in app.state.MODELS:
  666. url_idx = random.choice(app.state.MODELS[form_data.model]["urls"])
  667. else:
  668. raise HTTPException(
  669. status_code=400,
  670. detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.model),
  671. )
  672. url = app.state.OLLAMA_BASE_URLS[url_idx]
  673. log.info(f"url: {url}")
  674. r = None
  675. def get_request():
  676. nonlocal form_data
  677. nonlocal r
  678. request_id = str(uuid.uuid4())
  679. try:
  680. REQUEST_POOL.append(request_id)
  681. def stream_content():
  682. try:
  683. if form_data.stream:
  684. yield json.dumps(
  685. {"request_id": request_id, "done": False}
  686. ) + "\n"
  687. for chunk in r.iter_content(chunk_size=8192):
  688. if request_id in REQUEST_POOL:
  689. yield chunk
  690. else:
  691. log.warning("User: canceled request")
  692. break
  693. finally:
  694. if hasattr(r, "close"):
  695. r.close()
  696. if request_id in REQUEST_POOL:
  697. REQUEST_POOL.remove(request_id)
  698. r = requests.request(
  699. method="POST",
  700. url=f"{url}/v1/chat/completions",
  701. data=form_data.model_dump_json(exclude_none=True).encode(),
  702. stream=True,
  703. )
  704. r.raise_for_status()
  705. return StreamingResponse(
  706. stream_content(),
  707. status_code=r.status_code,
  708. headers=dict(r.headers),
  709. )
  710. except Exception as e:
  711. raise e
  712. try:
  713. return await run_in_threadpool(get_request)
  714. except Exception as e:
  715. error_detail = "Open WebUI: Server Connection Error"
  716. if r is not None:
  717. try:
  718. res = r.json()
  719. if "error" in res:
  720. error_detail = f"Ollama: {res['error']}"
  721. except:
  722. error_detail = f"Ollama: {e}"
  723. raise HTTPException(
  724. status_code=r.status_code if r else 500,
  725. detail=error_detail,
  726. )
  727. @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
  728. async def deprecated_proxy(path: str, request: Request, user=Depends(get_current_user)):
  729. url = app.state.OLLAMA_BASE_URLS[0]
  730. target_url = f"{url}/{path}"
  731. body = await request.body()
  732. headers = dict(request.headers)
  733. if user.role in ["user", "admin"]:
  734. if path in ["pull", "delete", "push", "copy", "create"]:
  735. if user.role != "admin":
  736. raise HTTPException(
  737. status_code=status.HTTP_401_UNAUTHORIZED,
  738. detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
  739. )
  740. else:
  741. raise HTTPException(
  742. status_code=status.HTTP_401_UNAUTHORIZED,
  743. detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
  744. )
  745. headers.pop("host", None)
  746. headers.pop("authorization", None)
  747. headers.pop("origin", None)
  748. headers.pop("referer", None)
  749. r = None
  750. def get_request():
  751. nonlocal r
  752. request_id = str(uuid.uuid4())
  753. try:
  754. REQUEST_POOL.append(request_id)
  755. def stream_content():
  756. try:
  757. if path == "generate":
  758. data = json.loads(body.decode("utf-8"))
  759. if not ("stream" in data and data["stream"] == False):
  760. yield json.dumps({"id": request_id, "done": False}) + "\n"
  761. elif path == "chat":
  762. yield json.dumps({"id": request_id, "done": False}) + "\n"
  763. for chunk in r.iter_content(chunk_size=8192):
  764. if request_id in REQUEST_POOL:
  765. yield chunk
  766. else:
  767. log.warning("User: canceled request")
  768. break
  769. finally:
  770. if hasattr(r, "close"):
  771. r.close()
  772. if request_id in REQUEST_POOL:
  773. REQUEST_POOL.remove(request_id)
  774. r = requests.request(
  775. method=request.method,
  776. url=target_url,
  777. data=body,
  778. headers=headers,
  779. stream=True,
  780. )
  781. r.raise_for_status()
  782. # r.close()
  783. return StreamingResponse(
  784. stream_content(),
  785. status_code=r.status_code,
  786. headers=dict(r.headers),
  787. )
  788. except Exception as e:
  789. raise e
  790. try:
  791. return await run_in_threadpool(get_request)
  792. except Exception as e:
  793. error_detail = "Open WebUI: Server Connection Error"
  794. if r is not None:
  795. try:
  796. res = r.json()
  797. if "error" in res:
  798. error_detail = f"Ollama: {res['error']}"
  799. except:
  800. error_detail = f"Ollama: {e}"
  801. raise HTTPException(
  802. status_code=r.status_code if r else 500,
  803. detail=error_detail,
  804. )