From 6a9d67b5bb4c93fd343b334bee3e37703dff59f6 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Sun, 22 Mar 2026 22:10:04 -0500 Subject: [PATCH] refac --- backend/open_webui/utils/middleware.py | 16 ++++++- backend/open_webui/utils/tools.py | 64 ++++++++++++++++++++++++-- src/lib/apis/index.ts | 28 ++++++++++- 3 files changed, 102 insertions(+), 6 deletions(-) diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index d721781ce4..044c2974b3 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -2557,7 +2557,7 @@ async def process_chat_payload(request, form_data, user, metadata, model): # so system terminals work even when no other tools are selected) if terminal_id: try: - terminal_tools = await get_terminal_tools( + terminal_tools, system_prompt = await get_terminal_tools( request, terminal_id, user, @@ -2565,11 +2565,25 @@ async def process_chat_payload(request, form_data, user, metadata, model): ) if terminal_tools: tools_dict = {**tools_dict, **terminal_tools} + if system_prompt: + form_data['messages'] = add_or_update_system_message( + system_prompt, + form_data['messages'], + append=True, + ) except Exception as e: log.exception(e) if direct_tool_servers: for tool_server in direct_tool_servers: + system_prompt = tool_server.pop('system_prompt', None) + if system_prompt: + form_data['messages'] = add_or_update_system_message( + system_prompt, + form_data['messages'], + append=True, + ) + tool_specs = tool_server.pop('specs', []) for tool in tool_specs: diff --git a/backend/open_webui/utils/tools.py b/backend/open_webui/utils/tools.py index 8febdfa31d..a098a83979 100644 --- a/backend/open_webui/utils/tools.py +++ b/backend/open_webui/utils/tools.py @@ -828,6 +828,43 @@ async def get_terminal_cwd( return None +async def get_terminal_system_prompt( + base_url: str, + headers: dict, + cookies: Optional[dict] = None, +) -> Optional[str]: + """Fetch the system prompt from a terminal server. + + Checks ``/api/config`` for the ``system`` feature flag first; + only fetches ``/system`` if the flag is present. Returns *None* + silently when the server doesn't support the endpoint. + """ + base = base_url.rstrip('/') + try: + async with aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout(total=3), + trust_env=True, + ) as session: + # 1. Check feature flag + async with session.get(f'{base}/api/config') as resp: + if resp.status != 200: + return None + config = await resp.json() + if not config.get('features', {}).get('system'): + return None + + # 2. Fetch system prompt + async with session.get( + f'{base}/system', headers=headers, cookies=cookies or {} + ) as resp: + if resp.status == 200: + data = await resp.json() + return data.get('prompt') + except Exception as e: + log.debug(f'Failed to fetch terminal system prompt: {e}') + return None + + async def set_terminal_servers(request: Request): """Load and cache OpenAPI specs from all TERMINAL_SERVER_CONNECTIONS.""" connections = request.app.state.config.TERMINAL_SERVER_CONNECTIONS or [] @@ -867,6 +904,25 @@ async def set_terminal_servers(request: Request): request.app.state.TERMINAL_SERVERS = await get_tool_servers_data(server_configs) + # Fetch system prompts concurrently (runs at cache time, not per-request) + connections_by_id = {c.get('id'): c for c in connections if c.get('id')} + + async def _fetch_system_prompt(server): + connection = connections_by_id.get(server.get('id')) + if not connection: + return + headers = {} + if connection.get('auth_type', 'bearer') == 'bearer': + headers['Authorization'] = f'Bearer {connection.get("key", "")}' + prompt = await get_terminal_system_prompt(server['url'], headers) + if prompt: + server['system_prompt'] = prompt + + await asyncio.gather( + *[_fetch_system_prompt(s) for s in request.app.state.TERMINAL_SERVERS], + return_exceptions=True, + ) + if request.app.state.redis is not None: await request.app.state.redis.set('terminal_servers', json.dumps(request.app.state.TERMINAL_SERVERS)) @@ -894,7 +950,7 @@ async def get_terminal_tools( terminal_id: str, user: UserModel, extra_params: dict, -) -> dict[str, dict]: +) -> tuple[dict[str, dict], Optional[str]]: """Resolve tools for a terminal server identified by terminal_id. - Finds the connection in TERMINAL_SERVER_CONNECTIONS @@ -941,14 +997,14 @@ async def get_terminal_tools( headers['Authorization'] = f'Bearer {oauth_token.get("access_token", "")}' # auth_type == "none": no Authorization header + system_prompt = server_data.get('system_prompt') terminal_cwd = await get_terminal_cwd(connection.get('url', ''), headers, cookies) tools_dict = {} for spec in specs: function_name = spec['name'] - - # Inject CWD into run_command description tool_spec = clean_openai_tool_schema(spec) + if function_name == 'run_command' and terminal_cwd: tool_spec['description'] = ( tool_spec.get('description', '') + f'\n\nThe current working directory is: {terminal_cwd}' @@ -977,7 +1033,7 @@ async def get_terminal_tools( 'type': 'terminal', } - return tools_dict + return tools_dict, system_prompt async def get_tool_server_data(url: str, headers: Optional[dict]) -> Dict[str, Any]: diff --git a/src/lib/apis/index.ts b/src/lib/apis/index.ts index f15acc8d08..020e75c758 100644 --- a/src/lib/apis/index.ts +++ b/src/lib/apis/index.ts @@ -392,12 +392,38 @@ export const getToolServersData = async (servers: object[]) => { specs: convertOpenApiToToolPayload(res) }; - return { + const result: Record = { url: server?.url, openapi: openapi, info: info, specs: specs }; + + // Fetch system prompt if the server supports it + try { + const baseUrl = (server?.url ?? '').replace(/\/$/, ''); + const configRes = await fetch(`${baseUrl}/api/config`); + if (configRes.ok) { + const config = await configRes.json(); + if (config?.features?.system) { + const headers: Record = {}; + if (toolServerToken) { + headers['Authorization'] = `Bearer ${toolServerToken}`; + } + const systemRes = await fetch(`${baseUrl}/system`, { headers }); + if (systemRes.ok) { + const systemData = await systemRes.json(); + if (systemData?.prompt) { + result.system_prompt = systemData.prompt; + } + } + } + } + } catch (e) { + // Server doesn't support /system — that's fine + } + + return result; } else if (error) { return { error,