# Copyright 2025 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Asserts for checking tool calls."""


import re


def check_tool_used_with_args_match(_: str, context: dict) -> dict:
    """Checks if one of the given tools was used with arguments matching all of
    a list of regexes.

    The assertion should be in the format:
    - type: >
        python:path/to/this/file.py:check_tool_used_with_args_match
      config:
        args_regexes: [<regex1>, <regex2>, ...]
        tool_names: [<tool_name_1>, <tool_name_2>, ...]
    """
    config = context.get('config', {})
    args_regexes = config.get('args_regexes')
    if not args_regexes:
        return {
            'pass': False,
            'reason': ('No regexes specified in the `args_regexes` field of '
                       'the assertion config.'),
            'score': 0,
        }
    if isinstance(args_regexes, str):
        args_regexes = [args_regexes]
    tool_names = config.get('tool_names')
    if not tool_names:
        return {
            'pass': False,
            'reason': ('No tools specified in the `tool_names` field of the '
                       'assertion config.'),
            'score': 0,
        }
    if isinstance(tool_names, str):
        tool_names = [tool_names]
    tool_names = set(tool_names)

    provider_response = context.get('providerResponse', {})
    metrics = provider_response.get('metrics', {})
    tool_calls = metrics.get('tool_calls', [])

    tool_used = None
    for call in tool_calls:
        function_name = call.get('function_name')
        if function_name in tool_names:
            tool_used = function_name
            args_str = call.get('function_args', '')
            if not args_str:
                continue

            if all(re.search(regex, args_str) for regex in args_regexes):
                return {
                    'pass': True,
                    'reason': (f'Tool "{function_name}" was used with '
                               'arguments matching all regexes.'),
                    'score': 1,
                }

    if not tool_used:
        used_tool_names = {c.get('function_name') for c in tool_calls}
        return {
            'pass': False,
            'reason': (f'None of the tools {sorted(tool_names)} were used. '
                       f'Used tools: {sorted(list(used_tool_names))}'),
            'score': 0,
        }

    # A tool was used, but no arguments matched the regex.
    return {
        'pass': False,
        'reason': (f'Tool "{tool_used}" was used, but its arguments did not '
                   f'match all of the regexes: "{args_regexes}".'),
        'score': 0,
    }