File size: 3,825 Bytes
d3cadd5 | 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 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | 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
|