File size: 3,288 Bytes
aed88a2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | 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))
|