|
|
from pathlib import Path |
|
|
import sys |
|
|
|
|
|
import pytest |
|
|
|
|
|
sys.path.insert(0, str(Path(__file__).parent.parent)) |
|
|
|
|
|
from kiro_proxy.core.thinking import ( |
|
|
ThinkingConfig, |
|
|
build_thinking_prompt, |
|
|
build_user_prompt_with_thinking, |
|
|
extract_thinking_config_from_gemini_body, |
|
|
extract_thinking_config_from_openai_body, |
|
|
infer_thinking_from_anthropic_messages, |
|
|
infer_thinking_from_gemini_contents, |
|
|
infer_thinking_from_openai_messages, |
|
|
infer_thinking_from_openai_responses_input, |
|
|
normalize_thinking_config, |
|
|
) |
|
|
|
|
|
|
|
|
def test_normalize_thinking_config_defaults_to_disabled_unlimited(): |
|
|
cfg = normalize_thinking_config(None) |
|
|
assert cfg == ThinkingConfig(False, None) |
|
|
|
|
|
|
|
|
@pytest.mark.parametrize( |
|
|
"raw,expected", |
|
|
[ |
|
|
(True, ThinkingConfig(True, None)), |
|
|
("enabled", ThinkingConfig(True, None)), |
|
|
({"type": "enabled"}, ThinkingConfig(True, None)), |
|
|
({"thinking_type": "enabled", "budget_tokens": 20000}, ThinkingConfig(True, 20000)), |
|
|
({"enabled": True, "budget_tokens": 0}, ThinkingConfig(True, None)), |
|
|
({"includeThoughts": True, "thinkingBudget": 1234}, ThinkingConfig(True, 1234)), |
|
|
({"type": "disabled", "budget_tokens": 9999}, ThinkingConfig(False, 9999)), |
|
|
], |
|
|
) |
|
|
def test_normalize_thinking_config_variants(raw, expected): |
|
|
assert normalize_thinking_config(raw) == expected |
|
|
|
|
|
|
|
|
def test_extract_thinking_config_from_openai_body(): |
|
|
cfg, explicit = extract_thinking_config_from_openai_body({}) |
|
|
assert cfg == ThinkingConfig(False, None) |
|
|
assert explicit is False |
|
|
|
|
|
cfg, explicit = extract_thinking_config_from_openai_body({"thinking": {"type": "enabled"}}) |
|
|
assert cfg.enabled is True |
|
|
assert explicit is True |
|
|
|
|
|
cfg, explicit = extract_thinking_config_from_openai_body({"reasoning_effort": "high"}) |
|
|
assert cfg.enabled is True |
|
|
assert cfg.budget_tokens is None |
|
|
assert explicit is True |
|
|
|
|
|
cfg, explicit = extract_thinking_config_from_openai_body({"reasoning": {"effort": "medium"}}) |
|
|
assert cfg == ThinkingConfig(True, 20000) |
|
|
assert explicit is True |
|
|
|
|
|
|
|
|
def test_extract_thinking_config_from_gemini_body(): |
|
|
cfg, explicit = extract_thinking_config_from_gemini_body({}) |
|
|
assert cfg == ThinkingConfig(False, None) |
|
|
assert explicit is False |
|
|
|
|
|
cfg, explicit = extract_thinking_config_from_gemini_body( |
|
|
{"generationConfig": {"thinkingConfig": {"includeThoughts": True, "thinkingBudget": 1234}}} |
|
|
) |
|
|
assert cfg == ThinkingConfig(True, 1234) |
|
|
assert explicit is True |
|
|
|
|
|
|
|
|
def test_infer_thinking_from_payloads(): |
|
|
assert ( |
|
|
infer_thinking_from_anthropic_messages( |
|
|
[{"role": "assistant", "content": [{"type": "thinking", "thinking": "x"}]}] |
|
|
) |
|
|
is True |
|
|
) |
|
|
|
|
|
assert infer_thinking_from_openai_messages( |
|
|
[{"role": "assistant", "content": "<thinking>AAA</thinking>\nBBB"}] |
|
|
) |
|
|
|
|
|
assert infer_thinking_from_openai_responses_input( |
|
|
[ |
|
|
{ |
|
|
"type": "message", |
|
|
"role": "assistant", |
|
|
"content": [{"type": "output_text", "text": "<thinking>AAA</thinking>BBB"}], |
|
|
} |
|
|
] |
|
|
) |
|
|
|
|
|
assert infer_thinking_from_gemini_contents( |
|
|
[{"role": "model", "parts": [{"text": "<thinking>AAA</thinking>\nBBB"}]}] |
|
|
) |
|
|
|
|
|
|
|
|
def test_thinking_prompts_include_ultrathink_and_budget_hint(): |
|
|
p1 = build_thinking_prompt("hi", budget_tokens=None) |
|
|
assert "ULTRATHINK" in p1 |
|
|
assert "within" not in p1.lower() |
|
|
|
|
|
p2 = build_thinking_prompt("hi", budget_tokens=123) |
|
|
assert "ULTRATHINK" in p2 |
|
|
assert "123" in p2 |
|
|
|
|
|
|
|
|
def test_build_user_prompt_with_thinking_wraps_and_forbids_disclosure(): |
|
|
prompt = build_user_prompt_with_thinking("hello", "secret reasoning") |
|
|
assert "<thinking>" in prompt and "</thinking>" in prompt |
|
|
assert "Do NOT reveal" in prompt |
|
|
|
