1
0

middleware.py 119 KB

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