middleware.py 103 KB

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