Skip to content

Route MCP tool results through the existing citation pipeline #50

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
cableman wants to merge 1 commit into
base: upstream/v0.11.0
Choose a base branch
from
Open

Route MCP tool results through the existing citation pipeline #50

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
75 changes: 55 additions & 20 deletions backend/open_webui/utils/middleware.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -307,6 +307,14 @@ def get_citation_source_from_tool_result(
if isinstance(tool_result, dict) and 'error' in tool_result:
return []

# Unwrap {'results': [...]} wrapper added by process_tool_result for MCP tools
if (
isinstance(tool_result, dict)
and 'results' in tool_result
and isinstance(tool_result['results'], list)
):
tool_result = tool_result['results']

# Validate tool_result type based on what the branch expects
if tool_name in _EXPECTS_LIST and not isinstance(tool_result, list):
return []
Expand Down Expand Up @@ -1312,24 +1320,43 @@ async def tool_call_handler(tool_call):
tool = tools[tool_function_name]
tool_id = tool.get('tool_id', '')

tool_name = f'{tool_id}/{tool_function_name}' if tool_id else f'{tool_function_name}'

# Citation is enabled for this tool
sources.append(
{
'source': {
'name': (f'{tool_name}'),
},
'document': [str(tool_result)],
'metadata': [
{
'source': (f'{tool_name}'),
'parameters': tool_function_params,
}
],
'tool_result': True,
}
)
# Use original (un-prefixed) name for MCP tools so that
# e.g. "websearch_search_web" matches the "search_web" handler.
citation_tool_name = tool.get('original_name', tool_function_name)

if citation_tool_name in [
'search_web',
'fetch_url',
'view_file',
'view_knowledge_file',
'query_knowledge_files',
]:
try:
citation_sources = get_citation_source_from_tool_result(
tool_name=citation_tool_name,
tool_params=tool_function_params,
tool_result=tool_result,
tool_id=tool_id,
)
sources.extend(citation_sources)
Comment thread
cableman marked this conversation as resolved.
except Exception as e:
log.exception(f'Error extracting citation source: {e}')
else:
tool_name = f'{tool_id}/{tool_function_name}' if tool_id else f'{tool_function_name}'
sources.append(
{
'source': {
'name': (f'{tool_name}'),
},
'document': [str(tool_result)],
'metadata': [
{
'source': (f'{tool_name}'),
'parameters': tool_function_params,
}
],
}
)

if tools[tool_function_name].get('metadata', {}).get('file_handler', False):
skip_files = True
Expand Down Expand Up @@ -2776,6 +2803,7 @@ async def tool_function(**kwargs):
'type': 'mcp',
'client': client,
'direct': False,
'original_name': tool_spec['name'],
}
except Exception as e:
log.debug(e)
Expand Down Expand Up @@ -5056,9 +5084,16 @@ async def execute_tool_call(tool_call):
)

# Extract citation sources from tool results
# Use original (un-prefixed) name for MCP tools so that
# e.g. "websearch_search_web" matches the "search_web" handler.
citation_tool_name = (
tool.get('original_name', tool_function_name)
if tool
else tool_function_name
)
if (
citations_enabled
and tool_function_name
and citation_tool_name
in [
'search_web',
'fetch_url',
Expand All @@ -5071,7 +5106,7 @@ async def execute_tool_call(tool_call):
):
try:
citation_sources = get_citation_source_from_tool_result(
tool_name=tool_function_name,
tool_name=citation_tool_name,
tool_params=tool_function_params,
tool_result=tool_result,
tool_id=tool.get('tool_id', '') if tool else '',
Expand Down