middleware.py 89 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254
  1. import time
  2. import logging
  3. import sys
  4. import os
  5. import base64
  6. import asyncio
  7. from aiocache import cached
  8. from typing import Any, Optional
  9. import random
  10. import json
  11. import html
  12. import inspect
  13. import re
  14. import ast
  15. from uuid import uuid4
  16. from concurrent.futures import ThreadPoolExecutor
  17. from fastapi import Request, HTTPException
  18. from starlette.responses import Response, StreamingResponse
  19. from open_webui.models.chats import Chats
  20. from open_webui.models.users import Users
  21. from open_webui.socket.main import (
  22. get_event_call,
  23. get_event_emitter,
  24. get_active_status_by_user_id,
  25. )
  26. from open_webui.routers.tasks import (
  27. generate_queries,
  28. generate_title,
  29. generate_image_prompt,
  30. generate_chat_tags,
  31. )
  32. from open_webui.routers.retrieval import process_web_search, SearchForm
  33. from open_webui.routers.images import image_generations, GenerateImageForm
  34. from open_webui.routers.pipelines import (
  35. process_pipeline_inlet_filter,
  36. process_pipeline_outlet_filter,
  37. )
  38. from open_webui.utils.webhook import post_webhook
  39. from open_webui.models.users import UserModel
  40. from open_webui.models.functions import Functions
  41. from open_webui.models.models import Models
  42. from open_webui.retrieval.utils import get_sources_from_files
  43. from open_webui.utils.chat import generate_chat_completion
  44. from open_webui.utils.task import (
  45. get_task_model_id,
  46. rag_template,
  47. tools_function_calling_generation_template,
  48. )
  49. from open_webui.utils.misc import (
  50. deep_update,
  51. get_message_list,
  52. add_or_update_system_message,
  53. add_or_update_user_message,
  54. get_last_user_message,
  55. get_last_assistant_message,
  56. prepend_to_first_user_message_content,
  57. convert_logit_bias_input_to_json,
  58. )
  59. from open_webui.utils.tools import get_tools
  60. from open_webui.utils.plugin import load_function_module_by_id
  61. from open_webui.utils.filter import (
  62. get_sorted_filter_ids,
  63. process_filter_functions,
  64. )
  65. from open_webui.utils.code_interpreter import execute_code_jupyter
  66. from open_webui.tasks import create_task
  67. from open_webui.config import (
  68. CACHE_DIR,
  69. DEFAULT_TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE,
  70. DEFAULT_CODE_INTERPRETER_PROMPT,
  71. )
  72. from open_webui.env import (
  73. SRC_LOG_LEVELS,
  74. GLOBAL_LOG_LEVEL,
  75. BYPASS_MODEL_ACCESS_CONTROL,
  76. ENABLE_REALTIME_CHAT_SAVE,
  77. )
  78. from open_webui.constants import TASKS
  79. logging.basicConfig(stream=sys.stdout, level=GLOBAL_LOG_LEVEL)
  80. log = logging.getLogger(__name__)
  81. log.setLevel(SRC_LOG_LEVELS["MAIN"])
  82. async def chat_completion_tools_handler(
  83. request: Request, body: dict, extra_params: dict, user: UserModel, models, tools
  84. ) -> tuple[dict, dict]:
  85. async def get_content_from_response(response) -> Optional[str]:
  86. content = None
  87. if hasattr(response, "body_iterator"):
  88. async for chunk in response.body_iterator:
  89. data = json.loads(chunk.decode("utf-8"))
  90. content = data["choices"][0]["message"]["content"]
  91. # Cleanup any remaining background tasks if necessary
  92. if response.background is not None:
  93. await response.background()
  94. else:
  95. content = response["choices"][0]["message"]["content"]
  96. return content
  97. def get_tools_function_calling_payload(messages, task_model_id, content):
  98. user_message = get_last_user_message(messages)
  99. history = "\n".join(
  100. f"{message['role'].upper()}: \"\"\"{message['content']}\"\"\""
  101. for message in messages[::-1][:4]
  102. )
  103. prompt = f"History:\n{history}\nQuery: {user_message}"
  104. return {
  105. "model": task_model_id,
  106. "messages": [
  107. {"role": "system", "content": content},
  108. {"role": "user", "content": f"Query: {prompt}"},
  109. ],
  110. "stream": False,
  111. "metadata": {"task": str(TASKS.FUNCTION_CALLING)},
  112. }
  113. event_caller = extra_params["__event_call__"]
  114. metadata = extra_params["__metadata__"]
  115. task_model_id = get_task_model_id(
  116. body["model"],
  117. request.app.state.config.TASK_MODEL,
  118. request.app.state.config.TASK_MODEL_EXTERNAL,
  119. models,
  120. )
  121. skip_files = False
  122. sources = []
  123. specs = [tool["spec"] for tool in tools.values()]
  124. tools_specs = json.dumps(specs)
  125. if request.app.state.config.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE != "":
  126. template = request.app.state.config.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE
  127. else:
  128. template = DEFAULT_TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE
  129. tools_function_calling_prompt = tools_function_calling_generation_template(
  130. template, tools_specs
  131. )
  132. payload = get_tools_function_calling_payload(
  133. body["messages"], task_model_id, tools_function_calling_prompt
  134. )
  135. try:
  136. response = await generate_chat_completion(request, form_data=payload, user=user)
  137. log.debug(f"{response=}")
  138. content = await get_content_from_response(response)
  139. log.debug(f"{content=}")
  140. if not content:
  141. return body, {}
  142. try:
  143. content = content[content.find("{") : content.rfind("}") + 1]
  144. if not content:
  145. raise Exception("No JSON object found in the response")
  146. result = json.loads(content)
  147. async def tool_call_handler(tool_call):
  148. nonlocal skip_files
  149. log.debug(f"{tool_call=}")
  150. tool_function_name = tool_call.get("name", None)
  151. if tool_function_name not in tools:
  152. return body, {}
  153. tool_function_params = tool_call.get("parameters", {})
  154. try:
  155. tool = tools[tool_function_name]
  156. spec = tool.get("spec", {})
  157. allowed_params = (
  158. spec.get("parameters", {}).get("properties", {}).keys()
  159. )
  160. tool_function_params = {
  161. k: v
  162. for k, v in tool_function_params.items()
  163. if k in allowed_params
  164. }
  165. if tool.get("direct", False):
  166. tool_result = await event_caller(
  167. {
  168. "type": "execute:tool",
  169. "data": {
  170. "id": str(uuid4()),
  171. "name": tool_function_name,
  172. "params": tool_function_params,
  173. "server": tool.get("server", {}),
  174. "session_id": metadata.get("session_id", None),
  175. },
  176. }
  177. )
  178. else:
  179. tool_function = tool["callable"]
  180. tool_result = await tool_function(**tool_function_params)
  181. except Exception as e:
  182. tool_result = str(e)
  183. if isinstance(tool_result, dict) or isinstance(tool_result, list):
  184. tool_result = json.dumps(tool_result, indent=2)
  185. if isinstance(tool_result, str):
  186. tool = tools[tool_function_name]
  187. tool_id = tool.get("tool_id", "")
  188. if tool.get("metadata", {}).get("citation", False) or tool.get(
  189. "direct", False
  190. ):
  191. sources.append(
  192. {
  193. "source": {
  194. "name": (
  195. f"TOOL:" + f"{tool_id}/{tool_function_name}"
  196. if tool_id
  197. else f"{tool_function_name}"
  198. ),
  199. },
  200. "document": [tool_result],
  201. "metadata": [
  202. {
  203. "source": (
  204. f"TOOL:" + f"{tool_id}/{tool_function_name}"
  205. if tool_id
  206. else f"{tool_function_name}"
  207. )
  208. }
  209. ],
  210. }
  211. )
  212. else:
  213. sources.append(
  214. {
  215. "source": {},
  216. "document": [tool_result],
  217. "metadata": [
  218. {
  219. "source": (
  220. f"TOOL:" + f"{tool_id}/{tool_function_name}"
  221. if tool_id
  222. else f"{tool_function_name}"
  223. )
  224. }
  225. ],
  226. }
  227. )
  228. if (
  229. tools[tool_function_name]
  230. .get("metadata", {})
  231. .get("file_handler", False)
  232. ):
  233. skip_files = True
  234. # check if "tool_calls" in result
  235. if result.get("tool_calls"):
  236. for tool_call in result.get("tool_calls"):
  237. await tool_call_handler(tool_call)
  238. else:
  239. await tool_call_handler(result)
  240. except Exception as e:
  241. log.debug(f"Error: {e}")
  242. content = None
  243. except Exception as e:
  244. log.debug(f"Error: {e}")
  245. content = None
  246. log.debug(f"tool_contexts: {sources}")
  247. if skip_files and "files" in body.get("metadata", {}):
  248. del body["metadata"]["files"]
  249. return body, {"sources": sources}
  250. async def chat_web_search_handler(
  251. request: Request, form_data: dict, extra_params: dict, user
  252. ):
  253. event_emitter = extra_params["__event_emitter__"]
  254. await event_emitter(
  255. {
  256. "type": "status",
  257. "data": {
  258. "action": "web_search",
  259. "description": "Generating search query",
  260. "done": False,
  261. },
  262. }
  263. )
  264. messages = form_data["messages"]
  265. user_message = get_last_user_message(messages)
  266. queries = []
  267. try:
  268. res = await generate_queries(
  269. request,
  270. {
  271. "model": form_data["model"],
  272. "messages": messages,
  273. "prompt": user_message,
  274. "type": "web_search",
  275. },
  276. user,
  277. )
  278. response = res["choices"][0]["message"]["content"]
  279. try:
  280. bracket_start = response.find("{")
  281. bracket_end = response.rfind("}") + 1
  282. if bracket_start == -1 or bracket_end == -1:
  283. raise Exception("No JSON object found in the response")
  284. response = response[bracket_start:bracket_end]
  285. queries = json.loads(response)
  286. queries = queries.get("queries", [])
  287. except Exception as e:
  288. queries = [response]
  289. except Exception as e:
  290. log.exception(e)
  291. queries = [user_message]
  292. if len(queries) == 0:
  293. await event_emitter(
  294. {
  295. "type": "status",
  296. "data": {
  297. "action": "web_search",
  298. "description": "No search query generated",
  299. "done": True,
  300. },
  301. }
  302. )
  303. return form_data
  304. all_results = []
  305. for searchQuery in queries:
  306. await event_emitter(
  307. {
  308. "type": "status",
  309. "data": {
  310. "action": "web_search",
  311. "description": 'Searching "{{searchQuery}}"',
  312. "query": searchQuery,
  313. "done": False,
  314. },
  315. }
  316. )
  317. try:
  318. results = await process_web_search(
  319. request,
  320. SearchForm(
  321. **{
  322. "query": searchQuery,
  323. }
  324. ),
  325. user=user,
  326. )
  327. if results:
  328. all_results.append(results)
  329. files = form_data.get("files", [])
  330. if results.get("collection_name"):
  331. files.append(
  332. {
  333. "collection_name": results["collection_name"],
  334. "name": searchQuery,
  335. "type": "web_search",
  336. "urls": results["filenames"],
  337. }
  338. )
  339. elif results.get("docs"):
  340. files.append(
  341. {
  342. "docs": results.get("docs", []),
  343. "name": searchQuery,
  344. "type": "web_search",
  345. "urls": results["filenames"],
  346. }
  347. )
  348. form_data["files"] = files
  349. except Exception as e:
  350. log.exception(e)
  351. await event_emitter(
  352. {
  353. "type": "status",
  354. "data": {
  355. "action": "web_search",
  356. "description": 'Error searching "{{searchQuery}}"',
  357. "query": searchQuery,
  358. "done": True,
  359. "error": True,
  360. },
  361. }
  362. )
  363. if all_results:
  364. urls = []
  365. for results in all_results:
  366. if "filenames" in results:
  367. urls.extend(results["filenames"])
  368. await event_emitter(
  369. {
  370. "type": "status",
  371. "data": {
  372. "action": "web_search",
  373. "description": "Searched {{count}} sites",
  374. "urls": urls,
  375. "done": True,
  376. },
  377. }
  378. )
  379. else:
  380. await event_emitter(
  381. {
  382. "type": "status",
  383. "data": {
  384. "action": "web_search",
  385. "description": "No search results found",
  386. "done": True,
  387. "error": True,
  388. },
  389. }
  390. )
  391. return form_data
  392. async def chat_image_generation_handler(
  393. request: Request, form_data: dict, extra_params: dict, user
  394. ):
  395. __event_emitter__ = extra_params["__event_emitter__"]
  396. await __event_emitter__(
  397. {
  398. "type": "status",
  399. "data": {"description": "Generating an image", "done": False},
  400. }
  401. )
  402. messages = form_data["messages"]
  403. user_message = get_last_user_message(messages)
  404. prompt = user_message
  405. negative_prompt = ""
  406. if request.app.state.config.ENABLE_IMAGE_PROMPT_GENERATION:
  407. try:
  408. res = await generate_image_prompt(
  409. request,
  410. {
  411. "model": form_data["model"],
  412. "messages": messages,
  413. },
  414. user,
  415. )
  416. response = res["choices"][0]["message"]["content"]
  417. try:
  418. bracket_start = response.find("{")
  419. bracket_end = response.rfind("}") + 1
  420. if bracket_start == -1 or bracket_end == -1:
  421. raise Exception("No JSON object found in the response")
  422. response = response[bracket_start:bracket_end]
  423. response = json.loads(response)
  424. prompt = response.get("prompt", [])
  425. except Exception as e:
  426. prompt = user_message
  427. except Exception as e:
  428. log.exception(e)
  429. prompt = user_message
  430. system_message_content = ""
  431. try:
  432. images = await image_generations(
  433. request=request,
  434. form_data=GenerateImageForm(**{"prompt": prompt}),
  435. user=user,
  436. )
  437. await __event_emitter__(
  438. {
  439. "type": "status",
  440. "data": {"description": "Generated an image", "done": True},
  441. }
  442. )
  443. for image in images:
  444. await __event_emitter__(
  445. {
  446. "type": "message",
  447. "data": {"content": f"![Generated Image]({image['url']})\n"},
  448. }
  449. )
  450. system_message_content = "<context>User is shown the generated image, tell the user that the image has been generated</context>"
  451. except Exception as e:
  452. log.exception(e)
  453. await __event_emitter__(
  454. {
  455. "type": "status",
  456. "data": {
  457. "description": f"An error occurred while generating an image",
  458. "done": True,
  459. },
  460. }
  461. )
  462. system_message_content = "<context>Unable to generate an image, tell the user that an error occurred</context>"
  463. if system_message_content:
  464. form_data["messages"] = add_or_update_system_message(
  465. system_message_content, form_data["messages"]
  466. )
  467. return form_data
  468. async def chat_completion_files_handler(
  469. request: Request, body: dict, user: UserModel
  470. ) -> tuple[dict, dict[str, list]]:
  471. sources = []
  472. if files := body.get("metadata", {}).get("files", None):
  473. queries = []
  474. try:
  475. queries_response = await generate_queries(
  476. request,
  477. {
  478. "model": body["model"],
  479. "messages": body["messages"],
  480. "type": "retrieval",
  481. },
  482. user,
  483. )
  484. queries_response = queries_response["choices"][0]["message"]["content"]
  485. try:
  486. bracket_start = queries_response.find("{")
  487. bracket_end = queries_response.rfind("}") + 1
  488. if bracket_start == -1 or bracket_end == -1:
  489. raise Exception("No JSON object found in the response")
  490. queries_response = queries_response[bracket_start:bracket_end]
  491. queries_response = json.loads(queries_response)
  492. except Exception as e:
  493. queries_response = {"queries": [queries_response]}
  494. queries = queries_response.get("queries", [])
  495. except:
  496. pass
  497. if len(queries) == 0:
  498. queries = [get_last_user_message(body["messages"])]
  499. try:
  500. # Offload get_sources_from_files to a separate thread
  501. loop = asyncio.get_running_loop()
  502. with ThreadPoolExecutor() as executor:
  503. sources = await loop.run_in_executor(
  504. executor,
  505. lambda: get_sources_from_files(
  506. request=request,
  507. files=files,
  508. queries=queries,
  509. embedding_function=lambda query, prefix: request.app.state.EMBEDDING_FUNCTION(
  510. query, prefix=prefix, user=user
  511. ),
  512. k=request.app.state.config.TOP_K,
  513. reranking_function=request.app.state.rf,
  514. k_reranker=request.app.state.config.TOP_K_RERANKER,
  515. r=request.app.state.config.RELEVANCE_THRESHOLD,
  516. hybrid_search=request.app.state.config.ENABLE_RAG_HYBRID_SEARCH,
  517. full_context=request.app.state.config.RAG_FULL_CONTEXT,
  518. ),
  519. )
  520. except Exception as e:
  521. log.exception(e)
  522. log.debug(f"rag_contexts:sources: {sources}")
  523. return body, {"sources": sources}
  524. def apply_params_to_form_data(form_data, model):
  525. params = form_data.pop("params", {})
  526. if model.get("ollama"):
  527. form_data["options"] = params
  528. if "format" in params:
  529. form_data["format"] = params["format"]
  530. if "keep_alive" in params:
  531. form_data["keep_alive"] = params["keep_alive"]
  532. else:
  533. if "seed" in params and params["seed"] is not None:
  534. form_data["seed"] = params["seed"]
  535. if "stop" in params and params["stop"] is not None:
  536. form_data["stop"] = params["stop"]
  537. if "temperature" in params and params["temperature"] is not None:
  538. form_data["temperature"] = params["temperature"]
  539. if "max_tokens" in params and params["max_tokens"] is not None:
  540. form_data["max_tokens"] = params["max_tokens"]
  541. if "top_p" in params and params["top_p"] is not None:
  542. form_data["top_p"] = params["top_p"]
  543. if "frequency_penalty" in params and params["frequency_penalty"] is not None:
  544. form_data["frequency_penalty"] = params["frequency_penalty"]
  545. if "reasoning_effort" in params and params["reasoning_effort"] is not None:
  546. form_data["reasoning_effort"] = params["reasoning_effort"]
  547. if "logit_bias" in params and params["logit_bias"] is not None:
  548. try:
  549. form_data["logit_bias"] = json.loads(
  550. convert_logit_bias_input_to_json(params["logit_bias"])
  551. )
  552. except Exception as e:
  553. print(f"Error parsing logit_bias: {e}")
  554. return form_data
  555. async def process_chat_payload(request, form_data, user, metadata, model):
  556. form_data = apply_params_to_form_data(form_data, model)
  557. log.debug(f"form_data: {form_data}")
  558. event_emitter = get_event_emitter(metadata)
  559. event_call = get_event_call(metadata)
  560. extra_params = {
  561. "__event_emitter__": event_emitter,
  562. "__event_call__": event_call,
  563. "__user__": {
  564. "id": user.id,
  565. "email": user.email,
  566. "name": user.name,
  567. "role": user.role,
  568. },
  569. "__metadata__": metadata,
  570. "__request__": request,
  571. "__model__": model,
  572. }
  573. # Initialize events to store additional event to be sent to the client
  574. # Initialize contexts and citation
  575. if getattr(request.state, "direct", False) and hasattr(request.state, "model"):
  576. models = {
  577. request.state.model["id"]: request.state.model,
  578. }
  579. else:
  580. models = request.app.state.MODELS
  581. task_model_id = get_task_model_id(
  582. form_data["model"],
  583. request.app.state.config.TASK_MODEL,
  584. request.app.state.config.TASK_MODEL_EXTERNAL,
  585. models,
  586. )
  587. events = []
  588. sources = []
  589. user_message = get_last_user_message(form_data["messages"])
  590. model_knowledge = model.get("info", {}).get("meta", {}).get("knowledge", False)
  591. if model_knowledge:
  592. await event_emitter(
  593. {
  594. "type": "status",
  595. "data": {
  596. "action": "knowledge_search",
  597. "query": user_message,
  598. "done": False,
  599. },
  600. }
  601. )
  602. knowledge_files = []
  603. for item in model_knowledge:
  604. if item.get("collection_name"):
  605. knowledge_files.append(
  606. {
  607. "id": item.get("collection_name"),
  608. "name": item.get("name"),
  609. "legacy": True,
  610. }
  611. )
  612. elif item.get("collection_names"):
  613. knowledge_files.append(
  614. {
  615. "name": item.get("name"),
  616. "type": "collection",
  617. "collection_names": item.get("collection_names"),
  618. "legacy": True,
  619. }
  620. )
  621. else:
  622. knowledge_files.append(item)
  623. files = form_data.get("files", [])
  624. files.extend(knowledge_files)
  625. form_data["files"] = files
  626. variables = form_data.pop("variables", None)
  627. # Process the form_data through the pipeline
  628. try:
  629. form_data = await process_pipeline_inlet_filter(
  630. request, form_data, user, models
  631. )
  632. except Exception as e:
  633. raise e
  634. try:
  635. filter_functions = [
  636. Functions.get_function_by_id(filter_id)
  637. for filter_id in get_sorted_filter_ids(model)
  638. ]
  639. form_data, flags = await process_filter_functions(
  640. request=request,
  641. filter_functions=filter_functions,
  642. filter_type="inlet",
  643. form_data=form_data,
  644. extra_params=extra_params,
  645. )
  646. except Exception as e:
  647. raise Exception(f"Error: {e}")
  648. features = form_data.pop("features", None)
  649. if features:
  650. if "web_search" in features and features["web_search"]:
  651. form_data = await chat_web_search_handler(
  652. request, form_data, extra_params, user
  653. )
  654. if "image_generation" in features and features["image_generation"]:
  655. form_data = await chat_image_generation_handler(
  656. request, form_data, extra_params, user
  657. )
  658. if "code_interpreter" in features and features["code_interpreter"]:
  659. form_data["messages"] = add_or_update_user_message(
  660. (
  661. request.app.state.config.CODE_INTERPRETER_PROMPT_TEMPLATE
  662. if request.app.state.config.CODE_INTERPRETER_PROMPT_TEMPLATE != ""
  663. else DEFAULT_CODE_INTERPRETER_PROMPT
  664. ),
  665. form_data["messages"],
  666. )
  667. tool_ids = form_data.pop("tool_ids", None)
  668. files = form_data.pop("files", None)
  669. # Remove files duplicates
  670. if files:
  671. files = list({json.dumps(f, sort_keys=True): f for f in files}.values())
  672. metadata = {
  673. **metadata,
  674. "tool_ids": tool_ids,
  675. "files": files,
  676. }
  677. form_data["metadata"] = metadata
  678. # Server side tools
  679. tool_ids = metadata.get("tool_ids", None)
  680. # Client side tools
  681. tool_servers = metadata.get("tool_servers", None)
  682. log.debug(f"{tool_ids=}")
  683. log.debug(f"{tool_servers=}")
  684. tools_dict = {}
  685. if tool_ids:
  686. tools_dict = get_tools(
  687. request,
  688. tool_ids,
  689. user,
  690. {
  691. **extra_params,
  692. "__model__": models[task_model_id],
  693. "__messages__": form_data["messages"],
  694. "__files__": metadata.get("files", []),
  695. },
  696. )
  697. if tool_servers:
  698. for tool_server in tool_servers:
  699. tool_specs = tool_server.pop("specs", [])
  700. for tool in tool_specs:
  701. tools_dict[tool["name"]] = {
  702. "spec": tool,
  703. "direct": True,
  704. "server": tool_server,
  705. }
  706. if tools_dict:
  707. if metadata.get("function_calling") == "native":
  708. # If the function calling is native, then call the tools function calling handler
  709. metadata["tools"] = tools_dict
  710. form_data["tools"] = [
  711. {"type": "function", "function": tool.get("spec", {})}
  712. for tool in tools_dict.values()
  713. ]
  714. else:
  715. # If the function calling is not native, then call the tools function calling handler
  716. try:
  717. form_data, flags = await chat_completion_tools_handler(
  718. request, form_data, extra_params, user, models, tools_dict
  719. )
  720. sources.extend(flags.get("sources", []))
  721. except Exception as e:
  722. log.exception(e)
  723. try:
  724. form_data, flags = await chat_completion_files_handler(request, form_data, user)
  725. sources.extend(flags.get("sources", []))
  726. except Exception as e:
  727. log.exception(e)
  728. # If context is not empty, insert it into the messages
  729. if len(sources) > 0:
  730. context_string = ""
  731. for source_idx, source in enumerate(sources):
  732. if "document" in source:
  733. for doc_idx, doc_context in enumerate(source["document"]):
  734. context_string += (
  735. f'<source id="{source_idx + 1}">{doc_context}</source>\n'
  736. )
  737. context_string = context_string.strip()
  738. prompt = get_last_user_message(form_data["messages"])
  739. if prompt is None:
  740. raise Exception("No user message found")
  741. if (
  742. request.app.state.config.RELEVANCE_THRESHOLD == 0
  743. and context_string.strip() == ""
  744. ):
  745. log.debug(
  746. f"With a 0 relevancy threshold for RAG, the context cannot be empty"
  747. )
  748. # Workaround for Ollama 2.0+ system prompt issue
  749. # TODO: replace with add_or_update_system_message
  750. if model.get("owned_by") == "ollama":
  751. form_data["messages"] = prepend_to_first_user_message_content(
  752. rag_template(
  753. request.app.state.config.RAG_TEMPLATE, context_string, prompt
  754. ),
  755. form_data["messages"],
  756. )
  757. else:
  758. form_data["messages"] = add_or_update_system_message(
  759. rag_template(
  760. request.app.state.config.RAG_TEMPLATE, context_string, prompt
  761. ),
  762. form_data["messages"],
  763. )
  764. # If there are citations, add them to the data_items
  765. sources = [source for source in sources if source.get("source", {}).get("name", "")]
  766. if len(sources) > 0:
  767. events.append({"sources": sources})
  768. if model_knowledge:
  769. await event_emitter(
  770. {
  771. "type": "status",
  772. "data": {
  773. "action": "knowledge_search",
  774. "query": user_message,
  775. "done": True,
  776. "hidden": True,
  777. },
  778. }
  779. )
  780. return form_data, metadata, events
  781. async def process_chat_response(
  782. request, response, form_data, user, metadata, model, events, tasks
  783. ):
  784. async def background_tasks_handler():
  785. message_map = Chats.get_messages_by_chat_id(metadata["chat_id"])
  786. message = message_map.get(metadata["message_id"]) if message_map else None
  787. if message:
  788. messages = get_message_list(message_map, message.get("id"))
  789. if tasks and messages:
  790. if TASKS.TITLE_GENERATION in tasks:
  791. if tasks[TASKS.TITLE_GENERATION]:
  792. res = await generate_title(
  793. request,
  794. {
  795. "model": message["model"],
  796. "messages": messages,
  797. "chat_id": metadata["chat_id"],
  798. },
  799. user,
  800. )
  801. if res and isinstance(res, dict):
  802. if len(res.get("choices", [])) == 1:
  803. title_string = (
  804. res.get("choices", [])[0]
  805. .get("message", {})
  806. .get("content", message.get("content", "New Chat"))
  807. )
  808. else:
  809. title_string = ""
  810. title_string = title_string[
  811. title_string.find("{") : title_string.rfind("}") + 1
  812. ]
  813. try:
  814. title = json.loads(title_string).get(
  815. "title", "New Chat"
  816. )
  817. except Exception as e:
  818. title = ""
  819. if not title:
  820. title = messages[0].get("content", "New Chat")
  821. Chats.update_chat_title_by_id(metadata["chat_id"], title)
  822. await event_emitter(
  823. {
  824. "type": "chat:title",
  825. "data": title,
  826. }
  827. )
  828. elif len(messages) == 2:
  829. title = messages[0].get("content", "New Chat")
  830. Chats.update_chat_title_by_id(metadata["chat_id"], title)
  831. await event_emitter(
  832. {
  833. "type": "chat:title",
  834. "data": message.get("content", "New Chat"),
  835. }
  836. )
  837. if TASKS.TAGS_GENERATION in tasks and tasks[TASKS.TAGS_GENERATION]:
  838. res = await generate_chat_tags(
  839. request,
  840. {
  841. "model": message["model"],
  842. "messages": messages,
  843. "chat_id": metadata["chat_id"],
  844. },
  845. user,
  846. )
  847. if res and isinstance(res, dict):
  848. if len(res.get("choices", [])) == 1:
  849. tags_string = (
  850. res.get("choices", [])[0]
  851. .get("message", {})
  852. .get("content", "")
  853. )
  854. else:
  855. tags_string = ""
  856. tags_string = tags_string[
  857. tags_string.find("{") : tags_string.rfind("}") + 1
  858. ]
  859. try:
  860. tags = json.loads(tags_string).get("tags", [])
  861. Chats.update_chat_tags_by_id(
  862. metadata["chat_id"], tags, user
  863. )
  864. await event_emitter(
  865. {
  866. "type": "chat:tags",
  867. "data": tags,
  868. }
  869. )
  870. except Exception as e:
  871. pass
  872. event_emitter = None
  873. event_caller = None
  874. if (
  875. "session_id" in metadata
  876. and metadata["session_id"]
  877. and "chat_id" in metadata
  878. and metadata["chat_id"]
  879. and "message_id" in metadata
  880. and metadata["message_id"]
  881. ):
  882. event_emitter = get_event_emitter(metadata)
  883. event_caller = get_event_call(metadata)
  884. # Non-streaming response
  885. if not isinstance(response, StreamingResponse):
  886. if event_emitter:
  887. if "error" in response:
  888. error = response["error"].get("detail", response["error"])
  889. Chats.upsert_message_to_chat_by_id_and_message_id(
  890. metadata["chat_id"],
  891. metadata["message_id"],
  892. {
  893. "error": {"content": error},
  894. },
  895. )
  896. if "selected_model_id" in response:
  897. Chats.upsert_message_to_chat_by_id_and_message_id(
  898. metadata["chat_id"],
  899. metadata["message_id"],
  900. {
  901. "selectedModelId": response["selected_model_id"],
  902. },
  903. )
  904. choices = response.get("choices", [])
  905. if choices and choices[0].get("message", {}).get("content"):
  906. content = response["choices"][0]["message"]["content"]
  907. if content:
  908. await event_emitter(
  909. {
  910. "type": "chat:completion",
  911. "data": response,
  912. }
  913. )
  914. title = Chats.get_chat_title_by_id(metadata["chat_id"])
  915. await event_emitter(
  916. {
  917. "type": "chat:completion",
  918. "data": {
  919. "done": True,
  920. "content": content,
  921. "title": title,
  922. },
  923. }
  924. )
  925. # Save message in the database
  926. Chats.upsert_message_to_chat_by_id_and_message_id(
  927. metadata["chat_id"],
  928. metadata["message_id"],
  929. {
  930. "content": content,
  931. },
  932. )
  933. # Send a webhook notification if the user is not active
  934. if get_active_status_by_user_id(user.id) is None:
  935. webhook_url = Users.get_user_webhook_url_by_id(user.id)
  936. if webhook_url:
  937. post_webhook(
  938. request.app.state.WEBUI_NAME,
  939. webhook_url,
  940. f"{title} - {request.app.state.config.WEBUI_URL}/c/{metadata['chat_id']}\n\n{content}",
  941. {
  942. "action": "chat",
  943. "message": content,
  944. "title": title,
  945. "url": f"{request.app.state.config.WEBUI_URL}/c/{metadata['chat_id']}",
  946. },
  947. )
  948. await background_tasks_handler()
  949. return response
  950. else:
  951. return response
  952. # Non standard response
  953. if not any(
  954. content_type in response.headers["Content-Type"]
  955. for content_type in ["text/event-stream", "application/x-ndjson"]
  956. ):
  957. return response
  958. extra_params = {
  959. "__event_emitter__": event_emitter,
  960. "__event_call__": event_caller,
  961. "__user__": {
  962. "id": user.id,
  963. "email": user.email,
  964. "name": user.name,
  965. "role": user.role,
  966. },
  967. "__metadata__": metadata,
  968. "__request__": request,
  969. "__model__": model,
  970. }
  971. filter_functions = [
  972. Functions.get_function_by_id(filter_id)
  973. for filter_id in get_sorted_filter_ids(model)
  974. ]
  975. # Streaming response
  976. if event_emitter and event_caller:
  977. task_id = str(uuid4()) # Create a unique task ID.
  978. model_id = form_data.get("model", "")
  979. Chats.upsert_message_to_chat_by_id_and_message_id(
  980. metadata["chat_id"],
  981. metadata["message_id"],
  982. {
  983. "model": model_id,
  984. },
  985. )
  986. def split_content_and_whitespace(content):
  987. content_stripped = content.rstrip()
  988. original_whitespace = (
  989. content[len(content_stripped) :]
  990. if len(content) > len(content_stripped)
  991. else ""
  992. )
  993. return content_stripped, original_whitespace
  994. def is_opening_code_block(content):
  995. backtick_segments = content.split("```")
  996. # Even number of segments means the last backticks are opening a new block
  997. return len(backtick_segments) > 1 and len(backtick_segments) % 2 == 0
  998. # Handle as a background task
  999. async def post_response_handler(response, events):
  1000. def serialize_content_blocks(content_blocks, raw=False):
  1001. content = ""
  1002. for block in content_blocks:
  1003. if block["type"] == "text":
  1004. content = f"{content}{block['content'].strip()}\n"
  1005. elif block["type"] == "tool_calls":
  1006. attributes = block.get("attributes", {})
  1007. tool_calls = block.get("content", [])
  1008. results = block.get("results", [])
  1009. if results:
  1010. tool_calls_display_content = ""
  1011. for tool_call in tool_calls:
  1012. tool_call_id = tool_call.get("id", "")
  1013. tool_name = tool_call.get("function", {}).get(
  1014. "name", ""
  1015. )
  1016. tool_arguments = tool_call.get("function", {}).get(
  1017. "arguments", ""
  1018. )
  1019. tool_result = None
  1020. tool_result_files = None
  1021. for result in results:
  1022. if tool_call_id == result.get("tool_call_id", ""):
  1023. tool_result = result.get("content", None)
  1024. tool_result_files = result.get("files", None)
  1025. break
  1026. if tool_result:
  1027. tool_calls_display_content = f'{tool_calls_display_content}\n<details type="tool_calls" done="true" id="{tool_call_id}" name="{tool_name}" arguments="{html.escape(json.dumps(tool_arguments))}" result="{html.escape(json.dumps(tool_result))}" files="{html.escape(json.dumps(tool_result_files)) if tool_result_files else ""}">\n<summary>Tool Executed</summary>\n</details>\n'
  1028. else:
  1029. tool_calls_display_content = f'{tool_calls_display_content}\n<details type="tool_calls" done="false" id="{tool_call_id}" name="{tool_name}" arguments="{html.escape(json.dumps(tool_arguments))}">\n<summary>Executing...</summary>\n</details>'
  1030. if not raw:
  1031. content = f"{content}\n{tool_calls_display_content}\n\n"
  1032. else:
  1033. tool_calls_display_content = ""
  1034. for tool_call in tool_calls:
  1035. tool_call_id = tool_call.get("id", "")
  1036. tool_name = tool_call.get("function", {}).get(
  1037. "name", ""
  1038. )
  1039. tool_arguments = tool_call.get("function", {}).get(
  1040. "arguments", ""
  1041. )
  1042. tool_calls_display_content = f'{tool_calls_display_content}\n<details type="tool_calls" done="false" id="{tool_call_id}" name="{tool_name}" arguments="{html.escape(json.dumps(tool_arguments))}">\n<summary>Executing...</summary>\n</details>'
  1043. if not raw:
  1044. content = f"{content}\n{tool_calls_display_content}\n\n"
  1045. elif block["type"] == "reasoning":
  1046. reasoning_display_content = "\n".join(
  1047. (f"> {line}" if not line.startswith(">") else line)
  1048. for line in block["content"].splitlines()
  1049. )
  1050. reasoning_duration = block.get("duration", None)
  1051. if reasoning_duration is not None:
  1052. if raw:
  1053. content = f'{content}\n<{block["start_tag"]}>{block["content"]}<{block["end_tag"]}>\n'
  1054. else:
  1055. content = f'{content}\n<details type="reasoning" done="true" duration="{reasoning_duration}">\n<summary>Thought for {reasoning_duration} seconds</summary>\n{reasoning_display_content}\n</details>\n'
  1056. else:
  1057. if raw:
  1058. content = f'{content}\n<{block["start_tag"]}>{block["content"]}<{block["end_tag"]}>\n'
  1059. else:
  1060. content = f'{content}\n<details type="reasoning" done="false">\n<summary>Thinking…</summary>\n{reasoning_display_content}\n</details>\n'
  1061. elif block["type"] == "code_interpreter":
  1062. attributes = block.get("attributes", {})
  1063. output = block.get("output", None)
  1064. lang = attributes.get("lang", "")
  1065. content_stripped, original_whitespace = (
  1066. split_content_and_whitespace(content)
  1067. )
  1068. if is_opening_code_block(content_stripped):
  1069. # Remove trailing backticks that would open a new block
  1070. content = (
  1071. content_stripped.rstrip("`").rstrip()
  1072. + original_whitespace
  1073. )
  1074. else:
  1075. # Keep content as is - either closing backticks or no backticks
  1076. content = content_stripped + original_whitespace
  1077. if output:
  1078. output = html.escape(json.dumps(output))
  1079. if raw:
  1080. content = f'{content}\n<code_interpreter type="code" lang="{lang}">\n{block["content"]}\n</code_interpreter>\n```output\n{output}\n```\n'
  1081. else:
  1082. content = f'{content}\n<details type="code_interpreter" done="true" output="{output}">\n<summary>Analyzed</summary>\n```{lang}\n{block["content"]}\n```\n</details>\n'
  1083. else:
  1084. if raw:
  1085. content = f'{content}\n<code_interpreter type="code" lang="{lang}">\n{block["content"]}\n</code_interpreter>\n'
  1086. else:
  1087. content = f'{content}\n<details type="code_interpreter" done="false">\n<summary>Analyzing...</summary>\n```{lang}\n{block["content"]}\n```\n</details>\n'
  1088. else:
  1089. block_content = str(block["content"]).strip()
  1090. content = f"{content}{block['type']}: {block_content}\n"
  1091. return content.strip()
  1092. def convert_content_blocks_to_messages(content_blocks):
  1093. messages = []
  1094. temp_blocks = []
  1095. for idx, block in enumerate(content_blocks):
  1096. if block["type"] == "tool_calls":
  1097. messages.append(
  1098. {
  1099. "role": "assistant",
  1100. "content": serialize_content_blocks(temp_blocks),
  1101. "tool_calls": block.get("content"),
  1102. }
  1103. )
  1104. results = block.get("results", [])
  1105. for result in results:
  1106. messages.append(
  1107. {
  1108. "role": "tool",
  1109. "tool_call_id": result["tool_call_id"],
  1110. "content": result["content"],
  1111. }
  1112. )
  1113. temp_blocks = []
  1114. else:
  1115. temp_blocks.append(block)
  1116. if temp_blocks:
  1117. content = serialize_content_blocks(temp_blocks)
  1118. if content:
  1119. messages.append(
  1120. {
  1121. "role": "assistant",
  1122. "content": content,
  1123. }
  1124. )
  1125. return messages
  1126. def tag_content_handler(content_type, tags, content, content_blocks):
  1127. end_flag = False
  1128. def extract_attributes(tag_content):
  1129. """Extract attributes from a tag if they exist."""
  1130. attributes = {}
  1131. if not tag_content: # Ensure tag_content is not None
  1132. return attributes
  1133. # Match attributes in the format: key="value" (ignores single quotes for simplicity)
  1134. matches = re.findall(r'(\w+)\s*=\s*"([^"]+)"', tag_content)
  1135. for key, value in matches:
  1136. attributes[key] = value
  1137. return attributes
  1138. if content_blocks[-1]["type"] == "text":
  1139. for start_tag, end_tag in tags:
  1140. # Match start tag e.g., <tag> or <tag attr="value">
  1141. start_tag_pattern = rf"<{re.escape(start_tag)}(\s.*?)?>"
  1142. match = re.search(start_tag_pattern, content)
  1143. if match:
  1144. attr_content = (
  1145. match.group(1) if match.group(1) else ""
  1146. ) # Ensure it's not None
  1147. attributes = extract_attributes(
  1148. attr_content
  1149. ) # Extract attributes safely
  1150. # Capture everything before and after the matched tag
  1151. before_tag = content[
  1152. : match.start()
  1153. ] # Content before opening tag
  1154. after_tag = content[
  1155. match.end() :
  1156. ] # Content after opening tag
  1157. # Remove the start tag and after from the currently handling text block
  1158. content_blocks[-1]["content"] = content_blocks[-1][
  1159. "content"
  1160. ].replace(match.group(0) + after_tag, "")
  1161. if before_tag:
  1162. content_blocks[-1]["content"] = before_tag
  1163. if not content_blocks[-1]["content"]:
  1164. content_blocks.pop()
  1165. # Append the new block
  1166. content_blocks.append(
  1167. {
  1168. "type": content_type,
  1169. "start_tag": start_tag,
  1170. "end_tag": end_tag,
  1171. "attributes": attributes,
  1172. "content": "",
  1173. "started_at": time.time(),
  1174. }
  1175. )
  1176. if after_tag:
  1177. content_blocks[-1]["content"] = after_tag
  1178. break
  1179. elif content_blocks[-1]["type"] == content_type:
  1180. start_tag = content_blocks[-1]["start_tag"]
  1181. end_tag = content_blocks[-1]["end_tag"]
  1182. # Match end tag e.g., </tag>
  1183. end_tag_pattern = rf"<{re.escape(end_tag)}>"
  1184. # Check if the content has the end tag
  1185. if re.search(end_tag_pattern, content):
  1186. end_flag = True
  1187. block_content = content_blocks[-1]["content"]
  1188. # Strip start and end tags from the content
  1189. start_tag_pattern = rf"<{re.escape(start_tag)}(.*?)>"
  1190. block_content = re.sub(
  1191. start_tag_pattern, "", block_content
  1192. ).strip()
  1193. end_tag_regex = re.compile(end_tag_pattern, re.DOTALL)
  1194. split_content = end_tag_regex.split(block_content, maxsplit=1)
  1195. # Content inside the tag
  1196. block_content = (
  1197. split_content[0].strip() if split_content else ""
  1198. )
  1199. # Leftover content (everything after `</tag>`)
  1200. leftover_content = (
  1201. split_content[1].strip() if len(split_content) > 1 else ""
  1202. )
  1203. if block_content:
  1204. content_blocks[-1]["content"] = block_content
  1205. content_blocks[-1]["ended_at"] = time.time()
  1206. content_blocks[-1]["duration"] = int(
  1207. content_blocks[-1]["ended_at"]
  1208. - content_blocks[-1]["started_at"]
  1209. )
  1210. # Reset the content_blocks by appending a new text block
  1211. if content_type != "code_interpreter":
  1212. if leftover_content:
  1213. content_blocks.append(
  1214. {
  1215. "type": "text",
  1216. "content": leftover_content,
  1217. }
  1218. )
  1219. else:
  1220. content_blocks.append(
  1221. {
  1222. "type": "text",
  1223. "content": "",
  1224. }
  1225. )
  1226. else:
  1227. # Remove the block if content is empty
  1228. content_blocks.pop()
  1229. if leftover_content:
  1230. content_blocks.append(
  1231. {
  1232. "type": "text",
  1233. "content": leftover_content,
  1234. }
  1235. )
  1236. else:
  1237. content_blocks.append(
  1238. {
  1239. "type": "text",
  1240. "content": "",
  1241. }
  1242. )
  1243. # Clean processed content
  1244. content = re.sub(
  1245. rf"<{re.escape(start_tag)}(.*?)>(.|\n)*?<{re.escape(end_tag)}>",
  1246. "",
  1247. content,
  1248. flags=re.DOTALL,
  1249. )
  1250. return content, content_blocks, end_flag
  1251. message = Chats.get_message_by_id_and_message_id(
  1252. metadata["chat_id"], metadata["message_id"]
  1253. )
  1254. tool_calls = []
  1255. last_assistant_message = None
  1256. try:
  1257. if form_data["messages"][-1]["role"] == "assistant":
  1258. last_assistant_message = get_last_assistant_message(
  1259. form_data["messages"]
  1260. )
  1261. except Exception as e:
  1262. pass
  1263. content = (
  1264. message.get("content", "")
  1265. if message
  1266. else last_assistant_message if last_assistant_message else ""
  1267. )
  1268. content_blocks = [
  1269. {
  1270. "type": "text",
  1271. "content": content,
  1272. }
  1273. ]
  1274. # We might want to disable this by default
  1275. DETECT_REASONING = True
  1276. DETECT_SOLUTION = True
  1277. DETECT_CODE_INTERPRETER = metadata.get("features", {}).get(
  1278. "code_interpreter", False
  1279. )
  1280. reasoning_tags = [
  1281. ("think", "/think"),
  1282. ("thinking", "/thinking"),
  1283. ("reason", "/reason"),
  1284. ("reasoning", "/reasoning"),
  1285. ("thought", "/thought"),
  1286. ("Thought", "/Thought"),
  1287. ("|begin_of_thought|", "|end_of_thought|"),
  1288. ]
  1289. code_interpreter_tags = [("code_interpreter", "/code_interpreter")]
  1290. solution_tags = [("|begin_of_solution|", "|end_of_solution|")]
  1291. try:
  1292. for event in events:
  1293. await event_emitter(
  1294. {
  1295. "type": "chat:completion",
  1296. "data": event,
  1297. }
  1298. )
  1299. # Save message in the database
  1300. Chats.upsert_message_to_chat_by_id_and_message_id(
  1301. metadata["chat_id"],
  1302. metadata["message_id"],
  1303. {
  1304. **event,
  1305. },
  1306. )
  1307. async def stream_body_handler(response):
  1308. nonlocal content
  1309. nonlocal content_blocks
  1310. response_tool_calls = []
  1311. async for line in response.body_iterator:
  1312. line = line.decode("utf-8") if isinstance(line, bytes) else line
  1313. data = line
  1314. # Skip empty lines
  1315. if not data.strip():
  1316. continue
  1317. # "data:" is the prefix for each event
  1318. if not data.startswith("data:"):
  1319. continue
  1320. # Remove the prefix
  1321. data = data[len("data:") :].strip()
  1322. try:
  1323. data = json.loads(data)
  1324. data, _ = await process_filter_functions(
  1325. request=request,
  1326. filter_functions=filter_functions,
  1327. filter_type="stream",
  1328. form_data=data,
  1329. extra_params=extra_params,
  1330. )
  1331. if data:
  1332. if "selected_model_id" in data:
  1333. model_id = data["selected_model_id"]
  1334. Chats.upsert_message_to_chat_by_id_and_message_id(
  1335. metadata["chat_id"],
  1336. metadata["message_id"],
  1337. {
  1338. "selectedModelId": model_id,
  1339. },
  1340. )
  1341. else:
  1342. choices = data.get("choices", [])
  1343. if not choices:
  1344. error = data.get("error", {})
  1345. if error:
  1346. await event_emitter(
  1347. {
  1348. "type": "chat:completion",
  1349. "data": {
  1350. "error": error,
  1351. },
  1352. }
  1353. )
  1354. usage = data.get("usage", {})
  1355. if usage:
  1356. await event_emitter(
  1357. {
  1358. "type": "chat:completion",
  1359. "data": {
  1360. "usage": usage,
  1361. },
  1362. }
  1363. )
  1364. continue
  1365. delta = choices[0].get("delta", {})
  1366. delta_tool_calls = delta.get("tool_calls", None)
  1367. if delta_tool_calls:
  1368. for delta_tool_call in delta_tool_calls:
  1369. tool_call_index = delta_tool_call.get(
  1370. "index"
  1371. )
  1372. if tool_call_index is not None:
  1373. if (
  1374. len(response_tool_calls)
  1375. <= tool_call_index
  1376. ):
  1377. response_tool_calls.append(
  1378. delta_tool_call
  1379. )
  1380. else:
  1381. delta_name = delta_tool_call.get(
  1382. "function", {}
  1383. ).get("name")
  1384. delta_arguments = (
  1385. delta_tool_call.get(
  1386. "function", {}
  1387. ).get("arguments")
  1388. )
  1389. if delta_name:
  1390. response_tool_calls[
  1391. tool_call_index
  1392. ]["function"][
  1393. "name"
  1394. ] += delta_name
  1395. if delta_arguments:
  1396. response_tool_calls[
  1397. tool_call_index
  1398. ]["function"][
  1399. "arguments"
  1400. ] += delta_arguments
  1401. value = delta.get("content")
  1402. reasoning_content = delta.get(
  1403. "reasoning_content"
  1404. ) or delta.get("reasoning")
  1405. if reasoning_content:
  1406. if (
  1407. not content_blocks
  1408. or content_blocks[-1]["type"] != "reasoning"
  1409. ):
  1410. reasoning_block = {
  1411. "type": "reasoning",
  1412. "start_tag": "think",
  1413. "end_tag": "/think",
  1414. "attributes": {
  1415. "type": "reasoning_content"
  1416. },
  1417. "content": "",
  1418. "started_at": time.time(),
  1419. }
  1420. content_blocks.append(reasoning_block)
  1421. else:
  1422. reasoning_block = content_blocks[-1]
  1423. reasoning_block["content"] += reasoning_content
  1424. data = {
  1425. "content": serialize_content_blocks(
  1426. content_blocks
  1427. )
  1428. }
  1429. if value:
  1430. if (
  1431. content_blocks
  1432. and content_blocks[-1]["type"]
  1433. == "reasoning"
  1434. and content_blocks[-1]
  1435. .get("attributes", {})
  1436. .get("type")
  1437. == "reasoning_content"
  1438. ):
  1439. reasoning_block = content_blocks[-1]
  1440. reasoning_block["ended_at"] = time.time()
  1441. reasoning_block["duration"] = int(
  1442. reasoning_block["ended_at"]
  1443. - reasoning_block["started_at"]
  1444. )
  1445. content_blocks.append(
  1446. {
  1447. "type": "text",
  1448. "content": "",
  1449. }
  1450. )
  1451. content = f"{content}{value}"
  1452. if not content_blocks:
  1453. content_blocks.append(
  1454. {
  1455. "type": "text",
  1456. "content": "",
  1457. }
  1458. )
  1459. content_blocks[-1]["content"] = (
  1460. content_blocks[-1]["content"] + value
  1461. )
  1462. if DETECT_REASONING:
  1463. content, content_blocks, _ = (
  1464. tag_content_handler(
  1465. "reasoning",
  1466. reasoning_tags,
  1467. content,
  1468. content_blocks,
  1469. )
  1470. )
  1471. if DETECT_CODE_INTERPRETER:
  1472. content, content_blocks, end = (
  1473. tag_content_handler(
  1474. "code_interpreter",
  1475. code_interpreter_tags,
  1476. content,
  1477. content_blocks,
  1478. )
  1479. )
  1480. if end:
  1481. break
  1482. if DETECT_SOLUTION:
  1483. content, content_blocks, _ = (
  1484. tag_content_handler(
  1485. "solution",
  1486. solution_tags,
  1487. content,
  1488. content_blocks,
  1489. )
  1490. )
  1491. if ENABLE_REALTIME_CHAT_SAVE:
  1492. # Save message in the database
  1493. Chats.upsert_message_to_chat_by_id_and_message_id(
  1494. metadata["chat_id"],
  1495. metadata["message_id"],
  1496. {
  1497. "content": serialize_content_blocks(
  1498. content_blocks
  1499. ),
  1500. },
  1501. )
  1502. else:
  1503. data = {
  1504. "content": serialize_content_blocks(
  1505. content_blocks
  1506. ),
  1507. }
  1508. await event_emitter(
  1509. {
  1510. "type": "chat:completion",
  1511. "data": data,
  1512. }
  1513. )
  1514. except Exception as e:
  1515. done = "data: [DONE]" in line
  1516. if done:
  1517. pass
  1518. else:
  1519. log.debug("Error: ", e)
  1520. continue
  1521. if content_blocks:
  1522. # Clean up the last text block
  1523. if content_blocks[-1]["type"] == "text":
  1524. content_blocks[-1]["content"] = content_blocks[-1][
  1525. "content"
  1526. ].strip()
  1527. if not content_blocks[-1]["content"]:
  1528. content_blocks.pop()
  1529. if not content_blocks:
  1530. content_blocks.append(
  1531. {
  1532. "type": "text",
  1533. "content": "",
  1534. }
  1535. )
  1536. if response_tool_calls:
  1537. tool_calls.append(response_tool_calls)
  1538. if response.background:
  1539. await response.background()
  1540. await stream_body_handler(response)
  1541. MAX_TOOL_CALL_RETRIES = 10
  1542. tool_call_retries = 0
  1543. while len(tool_calls) > 0 and tool_call_retries < MAX_TOOL_CALL_RETRIES:
  1544. tool_call_retries += 1
  1545. response_tool_calls = tool_calls.pop(0)
  1546. content_blocks.append(
  1547. {
  1548. "type": "tool_calls",
  1549. "content": response_tool_calls,
  1550. }
  1551. )
  1552. await event_emitter(
  1553. {
  1554. "type": "chat:completion",
  1555. "data": {
  1556. "content": serialize_content_blocks(content_blocks),
  1557. },
  1558. }
  1559. )
  1560. tools = metadata.get("tools", {})
  1561. results = []
  1562. for tool_call in response_tool_calls:
  1563. tool_call_id = tool_call.get("id", "")
  1564. tool_name = tool_call.get("function", {}).get("name", "")
  1565. tool_function_params = {}
  1566. try:
  1567. # json.loads cannot be used because some models do not produce valid JSON
  1568. tool_function_params = ast.literal_eval(
  1569. tool_call.get("function", {}).get("arguments", "{}")
  1570. )
  1571. except Exception as e:
  1572. log.debug(e)
  1573. # Fallback to JSON parsing
  1574. try:
  1575. tool_function_params = json.loads(
  1576. tool_call.get("function", {}).get("arguments", "{}")
  1577. )
  1578. except Exception as e:
  1579. log.debug(
  1580. f"Error parsing tool call arguments: {tool_call.get('function', {}).get('arguments', '{}')}"
  1581. )
  1582. tool_result = None
  1583. if tool_name in tools:
  1584. tool = tools[tool_name]
  1585. spec = tool.get("spec", {})
  1586. try:
  1587. allowed_params = (
  1588. spec.get("parameters", {})
  1589. .get("properties", {})
  1590. .keys()
  1591. )
  1592. tool_function_params = {
  1593. k: v
  1594. for k, v in tool_function_params.items()
  1595. if k in allowed_params
  1596. }
  1597. if tool.get("direct", False):
  1598. tool_result = await event_caller(
  1599. {
  1600. "type": "execute:tool",
  1601. "data": {
  1602. "id": str(uuid4()),
  1603. "name": tool_name,
  1604. "params": tool_function_params,
  1605. "server": tool.get("server", {}),
  1606. "session_id": metadata.get(
  1607. "session_id", None
  1608. ),
  1609. },
  1610. }
  1611. )
  1612. else:
  1613. tool_function = tool["callable"]
  1614. tool_result = await tool_function(
  1615. **tool_function_params
  1616. )
  1617. except Exception as e:
  1618. tool_result = str(e)
  1619. tool_result_files = []
  1620. if isinstance(tool_result, list):
  1621. for item in tool_result:
  1622. if item.startswith("data:"):
  1623. tool_result_files.append(item)
  1624. tool_result.remove(item)
  1625. if isinstance(tool_result, dict) or isinstance(
  1626. tool_result, list
  1627. ):
  1628. tool_result = json.dumps(tool_result, indent=2)
  1629. results.append(
  1630. {
  1631. "tool_call_id": tool_call_id,
  1632. "content": tool_result,
  1633. **(
  1634. {"files": tool_result_files}
  1635. if tool_result_files
  1636. else {}
  1637. ),
  1638. }
  1639. )
  1640. content_blocks[-1]["results"] = results
  1641. content_blocks.append(
  1642. {
  1643. "type": "text",
  1644. "content": "",
  1645. }
  1646. )
  1647. await event_emitter(
  1648. {
  1649. "type": "chat:completion",
  1650. "data": {
  1651. "content": serialize_content_blocks(content_blocks),
  1652. },
  1653. }
  1654. )
  1655. try:
  1656. res = await generate_chat_completion(
  1657. request,
  1658. {
  1659. "model": model_id,
  1660. "stream": True,
  1661. "tools": form_data["tools"],
  1662. "messages": [
  1663. *form_data["messages"],
  1664. *convert_content_blocks_to_messages(content_blocks),
  1665. ],
  1666. },
  1667. user,
  1668. )
  1669. if isinstance(res, StreamingResponse):
  1670. await stream_body_handler(res)
  1671. else:
  1672. break
  1673. except Exception as e:
  1674. log.debug(e)
  1675. break
  1676. if DETECT_CODE_INTERPRETER:
  1677. MAX_RETRIES = 5
  1678. retries = 0
  1679. while (
  1680. content_blocks[-1]["type"] == "code_interpreter"
  1681. and retries < MAX_RETRIES
  1682. ):
  1683. await event_emitter(
  1684. {
  1685. "type": "chat:completion",
  1686. "data": {
  1687. "content": serialize_content_blocks(content_blocks),
  1688. },
  1689. }
  1690. )
  1691. retries += 1
  1692. log.debug(f"Attempt count: {retries}")
  1693. output = ""
  1694. try:
  1695. if content_blocks[-1]["attributes"].get("type") == "code":
  1696. code = content_blocks[-1]["content"]
  1697. if (
  1698. request.app.state.config.CODE_INTERPRETER_ENGINE
  1699. == "pyodide"
  1700. ):
  1701. output = await event_caller(
  1702. {
  1703. "type": "execute:python",
  1704. "data": {
  1705. "id": str(uuid4()),
  1706. "code": code,
  1707. "session_id": metadata.get(
  1708. "session_id", None
  1709. ),
  1710. },
  1711. }
  1712. )
  1713. elif (
  1714. request.app.state.config.CODE_INTERPRETER_ENGINE
  1715. == "jupyter"
  1716. ):
  1717. output = await execute_code_jupyter(
  1718. request.app.state.config.CODE_INTERPRETER_JUPYTER_URL,
  1719. code,
  1720. (
  1721. request.app.state.config.CODE_INTERPRETER_JUPYTER_AUTH_TOKEN
  1722. if request.app.state.config.CODE_INTERPRETER_JUPYTER_AUTH
  1723. == "token"
  1724. else None
  1725. ),
  1726. (
  1727. request.app.state.config.CODE_INTERPRETER_JUPYTER_AUTH_PASSWORD
  1728. if request.app.state.config.CODE_INTERPRETER_JUPYTER_AUTH
  1729. == "password"
  1730. else None
  1731. ),
  1732. request.app.state.config.CODE_INTERPRETER_JUPYTER_TIMEOUT,
  1733. )
  1734. else:
  1735. output = {
  1736. "stdout": "Code interpreter engine not configured."
  1737. }
  1738. log.debug(f"Code interpreter output: {output}")
  1739. if isinstance(output, dict):
  1740. stdout = output.get("stdout", "")
  1741. if isinstance(stdout, str):
  1742. stdoutLines = stdout.split("\n")
  1743. for idx, line in enumerate(stdoutLines):
  1744. if "data:image/png;base64" in line:
  1745. id = str(uuid4())
  1746. # ensure the path exists
  1747. os.makedirs(
  1748. os.path.join(CACHE_DIR, "images"),
  1749. exist_ok=True,
  1750. )
  1751. image_path = os.path.join(
  1752. CACHE_DIR,
  1753. f"images/{id}.png",
  1754. )
  1755. with open(image_path, "wb") as f:
  1756. f.write(
  1757. base64.b64decode(
  1758. line.split(",")[1]
  1759. )
  1760. )
  1761. stdoutLines[idx] = (
  1762. f"![Output Image {idx}](/cache/images/{id}.png)"
  1763. )
  1764. output["stdout"] = "\n".join(stdoutLines)
  1765. result = output.get("result", "")
  1766. if isinstance(result, str):
  1767. resultLines = result.split("\n")
  1768. for idx, line in enumerate(resultLines):
  1769. if "data:image/png;base64" in line:
  1770. id = str(uuid4())
  1771. # ensure the path exists
  1772. os.makedirs(
  1773. os.path.join(CACHE_DIR, "images"),
  1774. exist_ok=True,
  1775. )
  1776. image_path = os.path.join(
  1777. CACHE_DIR,
  1778. f"images/{id}.png",
  1779. )
  1780. with open(image_path, "wb") as f:
  1781. f.write(
  1782. base64.b64decode(
  1783. line.split(",")[1]
  1784. )
  1785. )
  1786. resultLines[idx] = (
  1787. f"![Output Image {idx}](/cache/images/{id}.png)"
  1788. )
  1789. output["result"] = "\n".join(resultLines)
  1790. except Exception as e:
  1791. output = str(e)
  1792. content_blocks[-1]["output"] = output
  1793. content_blocks.append(
  1794. {
  1795. "type": "text",
  1796. "content": "",
  1797. }
  1798. )
  1799. await event_emitter(
  1800. {
  1801. "type": "chat:completion",
  1802. "data": {
  1803. "content": serialize_content_blocks(content_blocks),
  1804. },
  1805. }
  1806. )
  1807. try:
  1808. res = await generate_chat_completion(
  1809. request,
  1810. {
  1811. "model": model_id,
  1812. "stream": True,
  1813. "messages": [
  1814. *form_data["messages"],
  1815. {
  1816. "role": "assistant",
  1817. "content": serialize_content_blocks(
  1818. content_blocks, raw=True
  1819. ),
  1820. },
  1821. ],
  1822. },
  1823. user,
  1824. )
  1825. if isinstance(res, StreamingResponse):
  1826. await stream_body_handler(res)
  1827. else:
  1828. break
  1829. except Exception as e:
  1830. log.debug(e)
  1831. break
  1832. title = Chats.get_chat_title_by_id(metadata["chat_id"])
  1833. data = {
  1834. "done": True,
  1835. "content": serialize_content_blocks(content_blocks),
  1836. "title": title,
  1837. }
  1838. if not ENABLE_REALTIME_CHAT_SAVE:
  1839. # Save message in the database
  1840. Chats.upsert_message_to_chat_by_id_and_message_id(
  1841. metadata["chat_id"],
  1842. metadata["message_id"],
  1843. {
  1844. "content": serialize_content_blocks(content_blocks),
  1845. },
  1846. )
  1847. # Send a webhook notification if the user is not active
  1848. if get_active_status_by_user_id(user.id) is None:
  1849. webhook_url = Users.get_user_webhook_url_by_id(user.id)
  1850. if webhook_url:
  1851. post_webhook(
  1852. request.app.state.WEBUI_NAME,
  1853. webhook_url,
  1854. f"{title} - {request.app.state.config.WEBUI_URL}/c/{metadata['chat_id']}\n\n{content}",
  1855. {
  1856. "action": "chat",
  1857. "message": content,
  1858. "title": title,
  1859. "url": f"{request.app.state.config.WEBUI_URL}/c/{metadata['chat_id']}",
  1860. },
  1861. )
  1862. await event_emitter(
  1863. {
  1864. "type": "chat:completion",
  1865. "data": data,
  1866. }
  1867. )
  1868. await background_tasks_handler()
  1869. except asyncio.CancelledError:
  1870. log.warning("Task was cancelled!")
  1871. await event_emitter({"type": "task-cancelled"})
  1872. if not ENABLE_REALTIME_CHAT_SAVE:
  1873. # Save message in the database
  1874. Chats.upsert_message_to_chat_by_id_and_message_id(
  1875. metadata["chat_id"],
  1876. metadata["message_id"],
  1877. {
  1878. "content": serialize_content_blocks(content_blocks),
  1879. },
  1880. )
  1881. if response.background is not None:
  1882. await response.background()
  1883. # background_tasks.add_task(post_response_handler, response, events)
  1884. task_id, _ = create_task(post_response_handler(response, events))
  1885. return {"status": True, "task_id": task_id}
  1886. else:
  1887. # Fallback to the original response
  1888. async def stream_wrapper(original_generator, events):
  1889. def wrap_item(item):
  1890. return f"data: {item}\n\n"
  1891. for event in events:
  1892. event, _ = await process_filter_functions(
  1893. request=request,
  1894. filter_functions=filter_functions,
  1895. filter_type="stream",
  1896. form_data=event,
  1897. extra_params=extra_params,
  1898. )
  1899. if event:
  1900. yield wrap_item(json.dumps(event))
  1901. async for data in original_generator:
  1902. data, _ = await process_filter_functions(
  1903. request=request,
  1904. filter_functions=filter_functions,
  1905. filter_type="stream",
  1906. form_data=data,
  1907. extra_params=extra_params,
  1908. )
  1909. if data:
  1910. yield data
  1911. return StreamingResponse(
  1912. stream_wrapper(response.body_iterator, events),
  1913. headers=dict(response.headers),
  1914. background=response.background,
  1915. )