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

Open WebUI 0.9.5 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 = """                        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
                        )
                        if force_final_response:
                            new_form_data.pop('tools', None)
                            new_form_data.pop('tool_ids', 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,
                            )

                        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)