open-notebook/tests/test_ask_token_budget.py at 7ca621f13a9ed000b43e6ff656587a68e9f4f894 · bferanmi806-sketch/open-notebook · GitHub

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
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
from types import SimpleNamespace
from typing import cast
from unittest.mock import AsyncMock

import pytest

import open_notebook.graphs.ask as ask


def test_ask_max_tokens_defaults_when_env_is_unset(monkeypatch):
    monkeypatch.delenv(ask.ASK_MAX_TOKENS_ENV_VAR, raising=False)

    assert ask.get_ask_max_tokens() == 8192


def test_ask_max_tokens_reads_positive_override(monkeypatch):
    monkeypatch.setenv(ask.ASK_MAX_TOKENS_ENV_VAR, "12000")

    assert ask.get_ask_max_tokens() == 12000


@pytest.mark.parametrize("value", ["not-a-number", "0", "-5"])
def test_ask_max_tokens_invalid_values_fall_back_to_default(monkeypatch, value):
    monkeypatch.setenv(ask.ASK_MAX_TOKENS_ENV_VAR, value)

    assert ask.get_ask_max_tokens() == 8192


@pytest.mark.asyncio
async def test_strategy_uses_fixed_token_budget(monkeypatch):
    monkeypatch.setenv(ask.ASK_MAX_TOKENS_ENV_VAR, "12000")
    model = SimpleNamespace(
        ainvoke=AsyncMock(
            return_value=SimpleNamespace(
                content='{"reasoning":"Need one search","searches":[]}'
            )
        )
    )
    provision = AsyncMock(return_value=model)
    monkeypatch.setattr(ask, "provision_langchain_model", provision)
    monkeypatch.setattr(
        ask.Prompter, "render", lambda self, **kwargs: "strategy prompt"
    )

    result = await ask.call_model_with_messages(
        cast(
            ask.ThreadState,
            {"question": "What is this?", "answers": [], "final_answer": ""},
        ),
        {"configurable": {"strategy_model": "strategy-model"}},
    )

    assert result == {
        "strategy": ask.Strategy(reasoning="Need one search", searches=[])
    }
    provision.assert_awaited_once_with(
        "strategy prompt",
        "strategy-model",
        "tools",
        max_tokens=2000,
        structured={"type": "json"},
    )
    model.ainvoke.assert_awaited_once_with("strategy prompt")


@pytest.mark.asyncio
async def test_provide_answer_uses_configured_token_budget(monkeypatch):
    monkeypatch.setenv(ask.ASK_MAX_TOKENS_ENV_VAR, "12000")
    vector_search = AsyncMock(return_value=[{"id": "source:1"}])
    model = SimpleNamespace(
        ainvoke=AsyncMock(
            return_value=SimpleNamespace(content="<think>internal</think>Answer")
        )
    )
    provision = AsyncMock(return_value=model)
    monkeypatch.setattr(ask, "vector_search", vector_search)
    monkeypatch.setattr(ask, "provision_langchain_model", provision)
    monkeypatch.setattr(ask.Prompter, "render", lambda self, **kwargs: "answer prompt")

    result = await ask.provide_answer(
        cast(
            ask.SubGraphState,
            {
                "question": "What is this?",
                "term": "this",
                "instructions": "Explain it",
            },
        ),
        {"configurable": {"answer_model": "answer-model"}},
    )

    assert result == {"answers": ["Answer"]}
    vector_search.assert_awaited_once_with("this", 10, True, True)
    provision.assert_awaited_once_with(
        "answer prompt",
        "answer-model",
        "tools",
        max_tokens=12000,
    )
    model.ainvoke.assert_awaited_once_with("answer prompt")


@pytest.mark.asyncio
async def test_write_final_answer_uses_configured_token_budget(monkeypatch):
    monkeypatch.setenv(ask.ASK_MAX_TOKENS_ENV_VAR, "12000")
    model = SimpleNamespace(
        ainvoke=AsyncMock(
            return_value=SimpleNamespace(content="<think>internal</think>Final answer")
        )
    )
    provision = AsyncMock(return_value=model)
    monkeypatch.setattr(ask, "provision_langchain_model", provision)
    monkeypatch.setattr(ask.Prompter, "render", lambda self, **kwargs: "final prompt")

    result = await ask.write_final_answer(
        {
            "question": "What is this?",
            "strategy": ask.Strategy(reasoning="", searches=[]),
            "answers": ["Answer"],
            "final_answer": "",
        },
        {"configurable": {"final_answer_model": "final-model"}},
    )

    assert result == {"final_answer": "Final answer"}
    provision.assert_awaited_once_with(
        "final prompt",
        "final-model",
        "tools",
        max_tokens=12000,
    )
    model.ainvoke.assert_awaited_once_with("final prompt")