middleware.py 130 KB

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