From 7ca621f13a9ed000b43e6ff656587a68e9f4f894 Mon Sep 17 00:00:00 2001 From: Balogun Feranmi Date: Thu, 6 Aug 2026 11:54:05 +0100 Subject: [PATCH 1/3] fix: make Ask output token budget configurable --- .env.example | 6 + CHANGELOG.md | 3 + docs/5-CONFIGURATION/environment-reference.md | 5 + open_notebook/graphs/ask.py | 42 +++++- tests/test_ask_token_budget.py | 132 ++++++++++++++++++ 5 files changed, 184 insertions(+), 4 deletions(-) create mode 100644 tests/test_ask_token_budget.py diff --git a/.env.example b/.env.example index 9e4b3b8f4e..74048204ef 100644 --- a/.env.example +++ b/.env.example @@ -70,6 +70,12 @@ SURREAL_DATABASE=open_notebook # to match. # OPEN_NOTEBOOK_MAX_UPLOAD_SIZE_MB=100 +# Maximum output tokens for Ask/Q&A intermediate and final answers (default: 8192). +# The structured strategy-generation step remains fixed at 2000. Invalid or +# non-positive values fall back to 8192. Restart the API process/container after +# changing this value. Higher values can increase latency and token usage. +# OPEN_NOTEBOOK_ASK_MAX_TOKENS=8192 + # Security # Password to protect this Open Notebook instance. Supports Docker secrets # via OPEN_NOTEBOOK_PASSWORD_FILE. Auth is fully disabled if unset. diff --git a/CHANGELOG.md b/CHANGELOG.md index 6252b92e5b..d9be60826f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Community contribution intake now separates exploration from execution: feature requests, product/design/architecture ideas and contribution proposals start in GitHub Discussions, while Issues are reserved for reproducible bugs and maintainer-approved work items. The Issue chooser routes contributors accordingly, a structured Ideas Discussion form starts from user goals and outcomes, and the contributor/maintainer docs plus PR template now describe the Discussion → Issue → PR graduation path (#1204). - Release image gate gained a `probe` scenario (`make release-test` runs it as part of `all`): container-level checks that a Python test suite can't cover because they depend on the shipped image's process supervision — `OPEN_NOTEBOOK_WORKER_MAX_TASKS` reaching the in-image worker (the supervisord `sh -c` expansion), and the worker surviving startup with `HTTP_PROXY` set while a user's `NO_PROXY` value is preserved (the internal SurrealDB websocket not being tunneled). Both were manual probes during the v1.14.0 release; they now run automatically. Release-process docs gained the post-tag re-cut sequence and a note on never leaving the version bump uncommitted (v1.14.0 retro) +### Fixed +- Ask/Q&A intermediate and final answers now use a configurable 8192-token output budget via `OPEN_NOTEBOOK_ASK_MAX_TOKENS`, while structured strategy generation remains fixed at 2000; invalid or non-positive values safely fall back to the default (#1221) + ## [1.14.0] - 2026-07-20 ### Added diff --git a/docs/5-CONFIGURATION/environment-reference.md b/docs/5-CONFIGURATION/environment-reference.md index e7a082b9e0..21f3036d62 100644 --- a/docs/5-CONFIGURATION/environment-reference.md +++ b/docs/5-CONFIGURATION/environment-reference.md @@ -16,9 +16,14 @@ Comprehensive list of all environment variables available in Open Notebook. | `FRONTEND_BIND_HOST` | No | `0.0.0.0` (in Docker) | Network interface for Next.js to bind to. Default `0.0.0.0` ensures accessibility from reverse proxies. (Replaces `HOSTNAME`, which container runtimes such as Podman override with the container/pod hostname, causing Next.js to bind to the wrong address) | | `API_HOST` | No | `0.0.0.0` (in Docker) | Network interface for the API (uvicorn) to bind to. Set to `::` for IPv6 dual-stack environments (listens on IPv6 and, on Linux defaults, IPv4 too) | | `OPEN_NOTEBOOK_MAX_UPLOAD_SIZE_MB` | No | 100 | Maximum request body size (in MB) the API will accept, enforced before auth/routing. Raise this if you need to upload larger audio/video files. A fronting reverse proxy's own limit (e.g. nginx `client_max_body_size`) still applies and should be raised to match. | +| `OPEN_NOTEBOOK_ASK_MAX_TOKENS` | No | 8192 | Maximum output tokens for the intermediate-answer and final-answer model calls used by Ask/Q&A. It does not control the structured strategy-generation step, which remains fixed at 2000. Invalid or non-positive values fall back to 8192. | > **Important**: `OPEN_NOTEBOOK_ENCRYPTION_KEY` is required for storing AI provider credentials via the Settings UI. Without it, you cannot save credentials. If you change or lose this key, all stored credentials become unreadable. +`OPEN_NOTEBOOK_ASK_MAX_TOKENS` is process-level configuration, so changing it +normally requires restarting the API process or container. Higher values can +increase latency and token usage. + --- ## Database: SurrealDB diff --git a/open_notebook/graphs/ask.py b/open_notebook/graphs/ask.py index 405cc7c093..151d81d239 100644 --- a/open_notebook/graphs/ask.py +++ b/open_notebook/graphs/ask.py @@ -1,4 +1,5 @@ import operator +import os from typing import Annotated, List from ai_prompter import Prompter @@ -6,6 +7,7 @@ from langchain_core.runnables import RunnableConfig from langgraph.graph import END, START, StateGraph from langgraph.types import Send +from loguru import logger from pydantic import BaseModel, Field from typing_extensions import TypedDict @@ -16,6 +18,36 @@ from open_notebook.utils.error_classifier import classify_error from open_notebook.utils.text_utils import extract_text_content +DEFAULT_ASK_MAX_TOKENS = 8192 +ASK_STRATEGY_MAX_TOKENS = 2000 +ASK_MAX_TOKENS_ENV_VAR = "OPEN_NOTEBOOK_ASK_MAX_TOKENS" + + +def get_ask_max_tokens() -> int: + """Read the configured output budget for Ask's prose responses.""" + raw = os.environ.get(ASK_MAX_TOKENS_ENV_VAR) + if raw is None: + return DEFAULT_ASK_MAX_TOKENS + + raw = raw.strip() + try: + max_tokens = int(raw) + except ValueError: + logger.warning( + f"{ASK_MAX_TOKENS_ENV_VAR}={raw!r} is not a valid integer; " + f"using the default of {DEFAULT_ASK_MAX_TOKENS}" + ) + return DEFAULT_ASK_MAX_TOKENS + + if max_tokens <= 0: + logger.warning( + f"{ASK_MAX_TOKENS_ENV_VAR}={raw!r} is not positive; " + f"using the default of {DEFAULT_ASK_MAX_TOKENS}" + ) + return DEFAULT_ASK_MAX_TOKENS + + return max_tokens + class SubGraphState(TypedDict): question: str @@ -60,7 +92,7 @@ async def call_model_with_messages(state: ThreadState, config: RunnableConfig) - system_prompt, config.get("configurable", {}).get("strategy_model"), "tools", - max_tokens=2000, + max_tokens=ASK_STRATEGY_MAX_TOKENS, structured=dict(type="json"), ) # model = model.bind_tools(tools) @@ -109,12 +141,14 @@ async def provide_answer(state: SubGraphState, config: RunnableConfig) -> dict: payload["results"] = results ids = [r["id"] for r in results] payload["ids"] = ids - system_prompt = Prompter(prompt_template="ask/query_process").render(data=payload) # type: ignore[arg-type] + system_prompt = Prompter(prompt_template="ask/query_process").render( + data=payload, # type: ignore[arg-type] + ) model = await provision_langchain_model( system_prompt, config.get("configurable", {}).get("answer_model"), "tools", - max_tokens=2000, + max_tokens=get_ask_max_tokens(), ) ai_message = await model.ainvoke(system_prompt) ai_content = extract_text_content(ai_message.content) @@ -133,7 +167,7 @@ async def write_final_answer(state: ThreadState, config: RunnableConfig) -> dict system_prompt, config.get("configurable", {}).get("final_answer_model"), "tools", - max_tokens=2000, + max_tokens=get_ask_max_tokens(), ) ai_message = await model.ainvoke(system_prompt) final_content = extract_text_content(ai_message.content) diff --git a/tests/test_ask_token_budget.py b/tests/test_ask_token_budget.py new file mode 100644 index 0000000000..f077f3f97b --- /dev/null +++ b/tests/test_ask_token_budget.py @@ -0,0 +1,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="internalAnswer") + ) + ) + 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="internalFinal 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") From 6dcfc993743d74ed6359959107be718bcdf29cf8 Mon Sep 17 00:00:00 2001 From: Balogun Feranmi Date: Thu, 6 Aug 2026 12:37:22 +0100 Subject: [PATCH 2/3] fix: cache Ask output token budget --- open_notebook/graphs/ask.py | 2 ++ tests/test_ask_token_budget.py | 19 +++++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/open_notebook/graphs/ask.py b/open_notebook/graphs/ask.py index 151d81d239..9ca6c2ec49 100644 --- a/open_notebook/graphs/ask.py +++ b/open_notebook/graphs/ask.py @@ -1,5 +1,6 @@ import operator import os +from functools import cache from typing import Annotated, List from ai_prompter import Prompter @@ -23,6 +24,7 @@ ASK_MAX_TOKENS_ENV_VAR = "OPEN_NOTEBOOK_ASK_MAX_TOKENS" +@cache def get_ask_max_tokens() -> int: """Read the configured output budget for Ask's prose responses.""" raw = os.environ.get(ASK_MAX_TOKENS_ENV_VAR) diff --git a/tests/test_ask_token_budget.py b/tests/test_ask_token_budget.py index f077f3f97b..8e5e0cbfca 100644 --- a/tests/test_ask_token_budget.py +++ b/tests/test_ask_token_budget.py @@ -7,6 +7,13 @@ import open_notebook.graphs.ask as ask +@pytest.fixture(autouse=True) +def clear_ask_max_tokens_cache(): + ask.get_ask_max_tokens.cache_clear() + yield + ask.get_ask_max_tokens.cache_clear() + + def test_ask_max_tokens_defaults_when_env_is_unset(monkeypatch): monkeypatch.delenv(ask.ASK_MAX_TOKENS_ENV_VAR, raising=False) @@ -19,6 +26,18 @@ def test_ask_max_tokens_reads_positive_override(monkeypatch): assert ask.get_ask_max_tokens() == 12000 +def test_ask_max_tokens_caches_process_value(monkeypatch): + monkeypatch.setenv(ask.ASK_MAX_TOKENS_ENV_VAR, "12000") + + assert ask.get_ask_max_tokens() == 12000 + + monkeypatch.setenv(ask.ASK_MAX_TOKENS_ENV_VAR, "16000") + assert ask.get_ask_max_tokens() == 12000 + + ask.get_ask_max_tokens.cache_clear() + assert ask.get_ask_max_tokens() == 16000 + + @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) From cbb7daf082c9379e481cfbab3690a81de389fc15 Mon Sep 17 00:00:00 2001 From: Balogun Feranmi Date: Sat, 15 Aug 2026 20:56:58 +0100 Subject: [PATCH 3/3] fix: apply Ask token budget to strategy generation --- .env.example | 4 ++-- CHANGELOG.md | 2 +- docs/5-CONFIGURATION/environment-reference.md | 2 +- open_notebook/graphs/ask.py | 5 ++--- tests/test_ask_token_budget.py | 4 ++-- 5 files changed, 8 insertions(+), 9 deletions(-) diff --git a/.env.example b/.env.example index 74048204ef..c21f662995 100644 --- a/.env.example +++ b/.env.example @@ -70,8 +70,8 @@ SURREAL_DATABASE=open_notebook # to match. # OPEN_NOTEBOOK_MAX_UPLOAD_SIZE_MB=100 -# Maximum output tokens for Ask/Q&A intermediate and final answers (default: 8192). -# The structured strategy-generation step remains fixed at 2000. Invalid or +# Maximum output tokens for all Ask/Q&A model-generation stages: structured +# strategy, intermediate answers, and final answers (default: 8192). Invalid or # non-positive values fall back to 8192. Restart the API process/container after # changing this value. Higher values can increase latency and token usage. # OPEN_NOTEBOOK_ASK_MAX_TOKENS=8192 diff --git a/CHANGELOG.md b/CHANGELOG.md index d9be60826f..bcaa6b6b79 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,7 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Release image gate gained a `probe` scenario (`make release-test` runs it as part of `all`): container-level checks that a Python test suite can't cover because they depend on the shipped image's process supervision — `OPEN_NOTEBOOK_WORKER_MAX_TASKS` reaching the in-image worker (the supervisord `sh -c` expansion), and the worker surviving startup with `HTTP_PROXY` set while a user's `NO_PROXY` value is preserved (the internal SurrealDB websocket not being tunneled). Both were manual probes during the v1.14.0 release; they now run automatically. Release-process docs gained the post-tag re-cut sequence and a note on never leaving the version bump uncommitted (v1.14.0 retro) ### Fixed -- Ask/Q&A intermediate and final answers now use a configurable 8192-token output budget via `OPEN_NOTEBOOK_ASK_MAX_TOKENS`, while structured strategy generation remains fixed at 2000; invalid or non-positive values safely fall back to the default (#1221) +- Ask/Q&A structured strategy generation, intermediate answers, and final answers now use the configurable 8192-token output budget from `OPEN_NOTEBOOK_ASK_MAX_TOKENS`; invalid or non-positive values safely fall back to the default (#1221) ## [1.14.0] - 2026-07-20 diff --git a/docs/5-CONFIGURATION/environment-reference.md b/docs/5-CONFIGURATION/environment-reference.md index 21f3036d62..ab162bfb23 100644 --- a/docs/5-CONFIGURATION/environment-reference.md +++ b/docs/5-CONFIGURATION/environment-reference.md @@ -16,7 +16,7 @@ Comprehensive list of all environment variables available in Open Notebook. | `FRONTEND_BIND_HOST` | No | `0.0.0.0` (in Docker) | Network interface for Next.js to bind to. Default `0.0.0.0` ensures accessibility from reverse proxies. (Replaces `HOSTNAME`, which container runtimes such as Podman override with the container/pod hostname, causing Next.js to bind to the wrong address) | | `API_HOST` | No | `0.0.0.0` (in Docker) | Network interface for the API (uvicorn) to bind to. Set to `::` for IPv6 dual-stack environments (listens on IPv6 and, on Linux defaults, IPv4 too) | | `OPEN_NOTEBOOK_MAX_UPLOAD_SIZE_MB` | No | 100 | Maximum request body size (in MB) the API will accept, enforced before auth/routing. Raise this if you need to upload larger audio/video files. A fronting reverse proxy's own limit (e.g. nginx `client_max_body_size`) still applies and should be raised to match. | -| `OPEN_NOTEBOOK_ASK_MAX_TOKENS` | No | 8192 | Maximum output tokens for the intermediate-answer and final-answer model calls used by Ask/Q&A. It does not control the structured strategy-generation step, which remains fixed at 2000. Invalid or non-positive values fall back to 8192. | +| `OPEN_NOTEBOOK_ASK_MAX_TOKENS` | No | 8192 | Maximum output tokens for all Ask/Q&A model-generation calls: structured strategy generation, intermediate answers, and final answers. Invalid or non-positive values fall back to 8192. | > **Important**: `OPEN_NOTEBOOK_ENCRYPTION_KEY` is required for storing AI provider credentials via the Settings UI. Without it, you cannot save credentials. If you change or lose this key, all stored credentials become unreadable. diff --git a/open_notebook/graphs/ask.py b/open_notebook/graphs/ask.py index 9ca6c2ec49..ee96e927c1 100644 --- a/open_notebook/graphs/ask.py +++ b/open_notebook/graphs/ask.py @@ -20,13 +20,12 @@ from open_notebook.utils.text_utils import extract_text_content DEFAULT_ASK_MAX_TOKENS = 8192 -ASK_STRATEGY_MAX_TOKENS = 2000 ASK_MAX_TOKENS_ENV_VAR = "OPEN_NOTEBOOK_ASK_MAX_TOKENS" @cache def get_ask_max_tokens() -> int: - """Read the configured output budget for Ask's prose responses.""" + """Read the configured output budget for all Ask model-generation calls.""" raw = os.environ.get(ASK_MAX_TOKENS_ENV_VAR) if raw is None: return DEFAULT_ASK_MAX_TOKENS @@ -94,7 +93,7 @@ async def call_model_with_messages(state: ThreadState, config: RunnableConfig) - system_prompt, config.get("configurable", {}).get("strategy_model"), "tools", - max_tokens=ASK_STRATEGY_MAX_TOKENS, + max_tokens=get_ask_max_tokens(), structured=dict(type="json"), ) # model = model.bind_tools(tools) diff --git a/tests/test_ask_token_budget.py b/tests/test_ask_token_budget.py index 8e5e0cbfca..6d2789c81a 100644 --- a/tests/test_ask_token_budget.py +++ b/tests/test_ask_token_budget.py @@ -46,7 +46,7 @@ def test_ask_max_tokens_invalid_values_fall_back_to_default(monkeypatch, value): @pytest.mark.asyncio -async def test_strategy_uses_fixed_token_budget(monkeypatch): +async def test_strategy_uses_configured_token_budget(monkeypatch): monkeypatch.setenv(ask.ASK_MAX_TOKENS_ENV_VAR, "12000") model = SimpleNamespace( ainvoke=AsyncMock( @@ -76,7 +76,7 @@ async def test_strategy_uses_fixed_token_budget(monkeypatch): "strategy prompt", "strategy-model", "tools", - max_tokens=2000, + max_tokens=12000, structured={"type": "json"}, ) model.ainvoke.assert_awaited_once_with("strategy prompt")