tests/run_injected_tools.py

import os, sys
sys.path.append(os.getcwd())
from app import ChatCompletionRequest
from utils import detect_tools_and_reasoning
from config import CONFIG
from pprint import pprint


def emulate_injection(req: ChatCompletionRequest):
    raw_prompt = req.prompt.strip() if req.prompt is not None else '\n\n'.join([m.content for m in req.messages]) if req.messages else ''
    detection = detect_tools_and_reasoning(raw_prompt)

    # compute web_search_enabled
    web_search_enabled = (
        True
        if (req.enable_web_search is not None and req.enable_web_search)
        else (
            req.web_search
            or (req.auto_web_search if req.auto_web_search is not None else CONFIG.AUTO_ENABLE_WEB_SEARCH and detection.get('need_web_search'))
        )
    )

    if req.enable_tools is not None:
        tools_enabled = bool(req.enable_tools)
    else:
        # Respect sampler override if present
        if req.sampler and isinstance(req.sampler, dict) and req.sampler.get('ALLOW_TOOLS') is not None:
            tools_enabled = bool(req.sampler.get('ALLOW_TOOLS'))
        else:
            auto_tools_flag = req.auto_tools if req.auto_tools is not None else CONFIG.AUTO_ENABLE_TOOLS
            tools_enabled = bool(req.tools) or CONFIG.ENABLE_TOOLS_BY_DEFAULT or (auto_tools_flag and (detection.get('need_calc') or detection.get('need_web_search')))

    if req.enable_reasoning is not None:
        reasoning_enabled = bool(req.enable_reasoning)
    else:
        reasoning_enabled = False

    # If tools_enabled and not provided, add detected tools
    if tools_enabled and not req.tools and detection.get('detected_tools'):
        req.tools = detection.get('detected_tools')

    # If web_search should be used, and not already set, set flag
    if web_search_enabled and not req.web_search:
        req.web_search = True

    return {
        'raw_prompt': raw_prompt,
        'detection': detection,
        'web_search_enabled': web_search_enabled,
        'tools_enabled': tools_enabled,
        'reasoning_enabled': reasoning_enabled,
        'req': req.model_dump(),
    }

# test cases
cases = [
    ChatCompletionRequest(model='rwkv-latest', prompt='Who is the current president of France?', stream=None, auto_web_search=True, auto_tools=None, auto_reasoning=None),
    ChatCompletionRequest(model='rwkv-latest', prompt='Calculate 2+3*4 for me', stream=None, auto_web_search=True, auto_tools=True, auto_reasoning=None),
    ChatCompletionRequest(model='rwkv-latest', prompt='Explain why the sky is blue', stream=None, auto_web_search=False, auto_tools=None, auto_reasoning=True),
    # Sampler override should disable web_search even though auto_web_search is True
    ChatCompletionRequest(model='rwkv-latest', prompt='Who is the current president of France?', stream=None, auto_web_search=True, auto_tools=None, auto_reasoning=None, sampler_allow_web_search=False),
    # Per-request sampler object also should disable tools
    ChatCompletionRequest(model='rwkv-latest', prompt='Calculate 2+3*4 for me', stream=None, auto_web_search=True, auto_tools=None, auto_reasoning=None, sampler= { 'ALLOW_TOOLS': False }),
]

for c in cases:
    print('---')
    pprint(emulate_injection(c))