File size: 1,680 Bytes
fa599d5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""Tests for rewards.recursion_extract — REW-04 v2."""
from rewards.recursion_extract import count_llm_calls


def test_empty_completion():
    assert count_llm_calls("") == 0


def test_no_code_block():
    assert count_llm_calls("The answer is <answer>silver</answer>") == 0


def test_single_call():
    c = "```python\nresult = llm('find', ctx[:5000])\n```\n<answer>silver</answer>"
    assert count_llm_calls(c) == 1


def test_two_calls_in_one_block():
    c = "```python\na = llm('q1', ctx[:1000])\nb = llm('q2', ctx[1000:])\n```"
    assert count_llm_calls(c) == 2


def test_calls_across_two_blocks():
    c = "```python\nx=llm('q', ctx)\n```\nthen\n```python\ny=llm('q2', ctx)\n```"
    assert count_llm_calls(c) == 2


def test_call_in_comment_not_counted():
    c = "```python\n# would call llm(stuff) but skipping\nprint('done')\n```"
    assert count_llm_calls(c) == 0


def test_call_in_string_literal_not_counted():
    c = '```python\nnote = "earlier code did llm(...)"\nprint(note)\n```'
    assert count_llm_calls(c) == 0


def test_call_outside_code_block_not_counted():
    c = "Maybe I should call llm(question, chunk) but I won't actually."
    assert count_llm_calls(c) == 0


def test_call_in_loop_counts_literal_occurrence():
    c = "```python\nfor chunk in chunks:\n    r = llm('find', chunk)\n```"
    assert count_llm_calls(c) == 1


def test_invalid_python_falls_back_to_regex():
    c = "```python\nthis is not valid python !!!\nresult = llm('q', ctx)\n```"
    assert count_llm_calls(c) >= 1  # fallback regex finds it


def test_bare_python_fence():
    c = "```\nans = llm('q', ctx)\n```"  # no language tag
    assert count_llm_calls(c) == 1