diff --git a/.env.example b/.env.example index 9e4b3b8f4e..c21f662995 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 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 + # 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..bcaa6b6b79 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 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 ### Added diff --git a/docs/5-CONFIGURATION/environment-reference.md b/docs/5-CONFIGURATION/environment-reference.md index e7a082b9e0..ab162bfb23 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 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. +`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..ee96e927c1 100644 --- a/open_notebook/graphs/ask.py +++ b/open_notebook/graphs/ask.py @@ -1,4 +1,6 @@ import operator +import os +from functools import cache from typing import Annotated, List from ai_prompter import Prompter @@ -6,6 +8,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 +19,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_MAX_TOKENS_ENV_VAR = "OPEN_NOTEBOOK_ASK_MAX_TOKENS" + + +@cache +def get_ask_max_tokens() -> int: + """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 + + 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 +93,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=get_ask_max_tokens(), structured=dict(type="json"), ) # model = model.bind_tools(tools) @@ -109,12 +142,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 +168,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..6d2789c81a --- /dev/null +++ b/tests/test_ask_token_budget.py @@ -0,0 +1,151 @@ +from types import SimpleNamespace +from typing import cast +from unittest.mock import AsyncMock + +import pytest + +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) + + 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 + + +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) + + assert ask.get_ask_max_tokens() == 8192 + + +@pytest.mark.asyncio +async def test_strategy_uses_configured_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=12000, + 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")