#!/usr/bin/env python3
"""Make Open WebUI finish with an answer when its internal tool budget is used.

The pinned Open WebUI revision otherwise ends the request with only
``Tool-call limit reached``.  The patch is deliberately assertion-based: an
upstream source change makes the image build fail instead of silently applying
the modification at the wrong location.
"""

from pathlib import Path
import sys


if len(sys.argv) != 2:
    raise SystemExit("usage: patch_tool_finalization.py MIDDLEWARE_PY")

path = Path(sys.argv[1])
source = path.read_text()

patch_marker = "The tool budget for this turn is exhausted. Do not call"
if patch_marker in source:
    raise SystemExit("Open WebUI tool-finalization patch is already present")

needle = """                tool_call_iterations = 0
                max_tool_call_iterations = getattr(
"""
replacement = """                tool_call_iterations = 0
                # Open WebUI counts batches, while one model turn may request many
                # functions in parallel. Bound actual executions as well so a small
                # local model cannot expand eight rounds into dozens of API calls.
                tool_call_executions = 0
                max_tool_call_executions = 6
                max_tool_call_iterations = getattr(
"""
if source.count(needle) != 1:
    raise SystemExit(
        f"expected exactly one Open WebUI tool counter anchor, found {source.count(needle)}"
    )
source = source.replace(needle, replacement)

needle = """                    response_tool_calls = tool_calls.pop(0)

                    # Append function_call items for each tool call
"""
replacement = """                    response_tool_calls = tool_calls.pop(0)

                    remaining_tool_calls = max(
                        0, max_tool_call_executions - tool_call_executions
                    )
                    skipped_tool_calls = response_tool_calls[remaining_tool_calls:]
                    response_tool_calls = response_tool_calls[:remaining_tool_calls]
                    if skipped_tool_calls:
                        skipped_ids = {call.get('id', '') for call in skipped_tool_calls}
                        # Responses API streaming may already have exposed all calls
                        # in `output`. Remove deliberately skipped calls so the next
                        # completion never receives an orphan function call.
                        output[:] = [
                            item
                            for item in output
                            if not (
                                item.get('type') == 'function_call'
                                and item.get('call_id', '') in skipped_ids
                            )
                        ]
                    tool_call_executions += len(response_tool_calls)

                    # Append function_call items for each tool call
"""
if source.count(needle) != 1:
    raise SystemExit(
        f"expected exactly one Open WebUI tool batch anchor, found {source.count(needle)}"
    )
source = source.replace(needle, replacement)

needle = """                        res = await generate_chat_completion(
                            request,
                            new_form_data,
                            user,
                            bypass_system_prompt=True,
                        )
"""

replacement = """                        # The upstream loop otherwise stops with only
                        # `Tool-call limit reached` after executing this batch. Reserve
                        # the final completion for synthesis: keep every accumulated
                        # result, but expose no tools to the model and explicitly require
                        # a useful, evidence-bounded answer.
                        force_final_response = (
                            (
                                max_tool_call_iterations is not None
                                and tool_call_iterations >= max_tool_call_iterations
                            )
                            or tool_call_executions >= max_tool_call_executions
                        )
                        if force_final_response:
                            new_form_data.pop('tools', None)
                            new_form_data.pop('tool_ids', None)
                            new_form_data.pop('tool_choice', None)
                            final_metadata = dict(metadata)
                            final_metadata['tools'] = {}
                            final_metadata['tool_ids'] = []
                            final_metadata['tool_servers'] = []
                            new_form_data['metadata'] = final_metadata
                            new_form_data['messages'] = add_or_update_system_message(
                                'The tool budget for this turn is exhausted. Do not call '
                                'or imitate any more tools. Give the user a concise final '
                                'answer now using only the tool results already present. '
                                'Explicitly distinguish verified findings from inference, '
                                'and state what could not be verified. Never end without a '
                                'visible answer.',
                                new_form_data['messages'],
                                append=True,
                            )
                            # A final user-role instruction is deliberately stronger
                            # than another system suffix after a long tool-call pattern.
                            # It is request-local and is not added to the saved chat.
                            new_form_data['messages'].append(
                                {
                                    'role': 'user',
                                    'content': (
                                        'The research phase is finished. Answer my original '
                                        'question now in normal prose. Do not output XML, JSON, '
                                        'function names, tool_call blocks, or requests for more '
                                        'files. Use the available evidence, mention uncertainty, '
                                        'and provide a useful final conclusion.'
                                    ),
                                }
                            )

                        res = await generate_chat_completion(
                            request,
                            new_form_data,
                            user,
                            bypass_system_prompt=True,
                        )
"""

if source.count(needle) != 1:
    raise SystemExit(
        f"expected exactly one Open WebUI continuation anchor, found {source.count(needle)}"
    )
source = source.replace(needle, replacement)

needle = """                            await stream_body_handler(res, new_form_data)
                            output[:0] = prior_output
"""
replacement = """                            await stream_body_handler(res, new_form_data)
                            if force_final_response:
                                # A model may still print tool-shaped output after the
                                # schemas were removed. It must not trigger the upstream
                                # hard-error path or another execution attempt.
                                tool_calls.clear()
                            output[:0] = prior_output
"""
if source.count(needle) != 1:
    raise SystemExit(
        f"expected exactly one Open WebUI stream anchor, found {source.count(needle)}"
    )
source = source.replace(needle, replacement)

path.write_text(source)