github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2026-07-16 17:15:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix:budget bug

Browse Source
This commit is contained in:
aries_ckt
2026-05-15 18:15:56 +08:00
parent 3f22486586
commit 6d16b8775b
11 changed files with 74 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
docs/docs/agents/modules/context-management.md
View File

@@ -276,8 +276,8 @@ Agent context management can be configured in the application TOML file:
```toml
[service.web.agent_context]
# Set to 0 to auto-detect from the selected model metadata.
max_context_tokens = 0
# Non-positive values fall back to the default context budget.
max_context_tokens = 120000
reserved_tokens = 4096
warning_threshold = 0.70
error_threshold = 0.90
@@ -289,11 +289,8 @@ min_keep_tokens = 10000
max_compact_failures = 3
```
When `max_context_tokens` is `0`, DB-GPT tries to read the selected model's
`context_length` from `llm_client.get_model_metadata(model_name)`. If the model
metadata is unavailable, it falls back to the default budget.
For stable behavior, set `context_length` on each LLM deployment:
For stable behavior, set `context_length` on each LLM deployment when you want
the model metadata to reflect the real provider window:
```toml
[[models.llms]]

7
docs/i18n/zh-CN/docusaurus-plugin-content-docs/current/agents/modules/context-management.md
View File

@@ -243,8 +243,7 @@ Observation: <observation text>
```toml
[service.web.agent_context]
# 设置为 0 时，会从当前模型 metadata 自动读取上下文窗口。
max_context_tokens = 0
max_context_tokens = 120000
reserved_tokens = 4096
warning_threshold = 0.70
error_threshold = 0.90
@@ -256,9 +255,7 @@ min_keep_tokens = 10000
max_compact_failures = 3
```
当 `max_context_tokens` 为 `0` 时，DB-GPT 会尝试通过 `llm_client.get_model_metadata(model_name)` 读取所选模型的 `context_length`。如果模型 metadata 不可用，则回退到默认预算。
为了让行为稳定，建议在每个 LLM 部署配置中显式设置 `context_length`：
为了让行为稳定，建议在每个 LLM 部署配置中显式设置 `context_length`，这样模型 metadata 也会反映真实的 provider 上下文窗口：
```toml
[[models.llms]]

9
packages/dbgpt-app/src/dbgpt_app/config.py
View File

@@ -1,6 +1,7 @@
from dataclasses import dataclass, field
from typing import List, Optional
from dbgpt.agent.core.context.budget import DEFAULT_MAX_CONTEXT_TOKENS
from dbgpt.datasource.parameter import BaseDatasourceParameters
from dbgpt.model.parameter import (
ModelsDeployParameters,
@@ -210,11 +211,11 @@ class AgentContextParameters(BaseParameters):
__cfg_type__ = "service"
max_context_tokens: Optional[int] = field(
default=0,
default=DEFAULT_MAX_CONTEXT_TOKENS,
metadata={
"help": _(
"Maximum context-window tokens for agent calls. "
"Set to 0 to auto-detect from model metadata."
"Non-positive values fall back to the default context budget."
)
},
)
@@ -255,6 +256,10 @@ class AgentContextParameters(BaseParameters):
metadata={"help": _("Consecutive compaction failures before circuit break")},
)
def __post_init__(self):
if self.max_context_tokens is None or self.max_context_tokens <= 0:
self.max_context_tokens = DEFAULT_MAX_CONTEXT_TOKENS
@dataclass
class ServiceWebParameters(BaseParameters):

12
packages/dbgpt-app/src/dbgpt_app/openapi/api_v1/agentic_data_api.py
View File

@@ -81,12 +81,9 @@ async def _load_context_budget_config(
app_config = CFG.SYSTEM_APP.config.configs.get("app_config")
web_config = getattr(getattr(app_config, "service", None), "web", None)
agent_context = getattr(web_config, "agent_context", None)
max_context_tokens = _value(agent_context, "max_context_tokens", 0)
if max_context_tokens <= 0:
max_context_tokens = (
await _resolve_model_context_tokens(llm_client, model_name)
or defaults.max_context_tokens
)
max_context_tokens = _value(
agent_context, "max_context_tokens", defaults.max_context_tokens
)
return ContextBudgetConfig(
max_context_tokens=max_context_tokens,
warning_threshold=_value(
@@ -133,9 +130,6 @@ async def _load_context_budget_config(
logger.debug(
"Failed to load agent context config; using defaults", exc_info=True
)
metadata_tokens = await _resolve_model_context_tokens(llm_client, model_name)
if metadata_tokens:
return ContextBudgetConfig(max_context_tokens=metadata_tokens)
return defaults

13
packages/dbgpt-core/src/dbgpt/agent/core/context/budget.py
View File

@@ -10,6 +10,8 @@ from dbgpt.model.utils.token_utils import ProxyTokenizerWrapper
logger = logging.getLogger(__name__)
DEFAULT_MAX_CONTEXT_TOKENS = 120000
class TokenState(Enum):
"""Context budget state levels."""
@@ -42,7 +44,7 @@ class ContextBudgetConfig:
def __init__(
self,
max_context_tokens: int = 120000,
max_context_tokens: int = DEFAULT_MAX_CONTEXT_TOKENS,
warning_threshold: float = 0.70,
error_threshold: float = 0.90,
critical_threshold: float = 0.95,
@@ -53,7 +55,14 @@ class ContextBudgetConfig:
truncated_observation_max_chars: int = 200,
min_keep_tokens: int = 10000,
):
self.max_context_tokens = max_context_tokens
# `max_context_tokens <= 0` means "auto-detect from model metadata".
# If the caller could not resolve metadata, keep the context window usable
# by falling back to the system default budget instead of leaving it at 0.
self.max_context_tokens = (
max_context_tokens
if max_context_tokens and max_context_tokens > 0
else DEFAULT_MAX_CONTEXT_TOKENS
)
self.warning_threshold = warning_threshold
self.error_threshold = error_threshold
self.critical_threshold = critical_threshold

6
packages/dbgpt-core/src/dbgpt/agent/core/context/manager.py
View File

@@ -59,6 +59,12 @@ class ContextManager:
) -> None:
"""Push a context.status event to the registered callback (if any)."""
budget = self.tracker.config.effective_budget
if budget <= 0:
logger.debug(
"Skip context status emit because effective budget is non-positive: %d",
budget,
)
return
ratio = round(token_count / budget, 4) if budget > 0 else 1.0
logger.info(
"Context status: tokens=%d, budget=%d, ratio=%.4f, state=%s, layer=%s",

8
packages/dbgpt-core/src/dbgpt/agent/core/context/tests/test_budget.py
View File

@@ -1,6 +1,7 @@
"""Tests for context budget tracking."""
from dbgpt.agent.core.context.budget import (
DEFAULT_MAX_CONTEXT_TOKENS,
ContextBudgetConfig,
ContextBudgetTracker,
TokenState,
@@ -42,6 +43,13 @@ class TestContextBudgetConfig:
assert cfg.error_threshold == 0.80
assert cfg.critical_threshold == 0.90
def test_non_positive_max_context_tokens_falls_back_to_default(self):
cfg = ContextBudgetConfig(max_context_tokens=0)
assert cfg.max_context_tokens == DEFAULT_MAX_CONTEXT_TOKENS
cfg = ContextBudgetConfig(max_context_tokens=-1)
assert cfg.max_context_tokens == DEFAULT_MAX_CONTEXT_TOKENS
class _FakeMsg:
"""Lightweight message stub for testing."""

19
packages/dbgpt-core/src/dbgpt/agent/core/context/tests/test_manager.py
View File

@@ -94,12 +94,13 @@ class TestContextManagerCompaction:
class TestReactiveCompact:
def test_reactive_compact_reduces_messages(self):
@pytest.mark.asyncio
async def test_reactive_compact_reduces_messages(self):
cfg = ContextBudgetConfig()
mgr = ContextManager(config=cfg)
msgs = _build_long_conversation(num_rounds=10)
result = mgr.reactive_compact(msgs)
result = await mgr.reactive_compact(msgs)
# Should keep system + last 2 rounds only
assert len(result) == 1 + 2 * 3 # system + 2 rounds × 3 msgs
@@ -114,3 +115,17 @@ class TestContextManagerRecordsHistory:
await mgr.manage_context(msgs, current_round=0)
assert len(mgr.tracker.token_history) == 1
assert mgr.tracker.token_history[0] > 0
@pytest.mark.asyncio
async def test_skip_status_emit_when_effective_budget_invalid(self):
received = []
async def _on_status(status):
received.append(status)
cfg = ContextBudgetConfig(max_context_tokens=1, reserved_tokens=10)
mgr = ContextManager(config=cfg, on_status_event=_on_status)
msgs = [_sys("system"), _ai("T"), _human("A"), _obs("data")]
await mgr.manage_context(msgs, current_round=0)
assert received == []

7
web/hooks/use-chat.ts
View File

@@ -118,12 +118,17 @@ const useChat = ({ queryAgentURL = '/api/v1/chat/completions', app_code }: Props
// React-agent format: {"type": "context.status", "used": ..., "budget": ..., ...}
const cs = parsedData.context_status ?? (parsedData.type === 'context.status' ? parsedData : null);
if (cs) {
const budget = Number(cs.budget ?? 0);
if (!Number.isFinite(budget) || budget <= 0) {
setContextStatus(null);
return;
}
// Only show banner when Layer 3 (LLM compression) is active
if (cs.compact_layer === 'layer3') {
setContextStatus({
state: mapContextState(cs.state || 'normal'),
used_tokens: cs.used ?? 0,
max_tokens: cs.budget ?? 0,
max_tokens: budget,
usage_percent: (cs.ratio ?? 0) * 100,
layer: cs.compact_layer,
message: cs.message,

7
web/pages/index.tsx
View File

@@ -1562,6 +1562,11 @@ const Playground: NextPage = () => {
return;
}
if (payload.type === 'context.status') {
const budget = Number(payload.budget ?? 0);
if (!Number.isFinite(budget) || budget <= 0) {
setContextStatus(null);
return;
}
const stateMap: Record<string, 'OK' | 'WARNING' | 'ERROR'> = {
normal: 'OK',
warning: 'WARNING',
@@ -1572,7 +1577,7 @@ const Playground: NextPage = () => {
setContextStatus({
state: stateMap[payload.state] || 'OK',
used_tokens: payload.used ?? 0,
max_tokens: payload.budget ?? 0,
max_tokens: budget,
usage_percent: (payload.ratio ?? 0) * 100,
layer: payload.compact_layer ?? null,
});

4
web/utils/react-sse-parser.ts
View File

@@ -201,6 +201,10 @@ export class ReActSSEState {
}
private handleContextStatus(event: SSEContextStatusEvent): void {
if (!Number.isFinite(event.budget) || event.budget <= 0) {
this._contextStatus = null;
return;
}
// Map backend state values (lowercase: "normal", "warning", "error",
// "critical", "overflow") to frontend display states ("OK", "WARNING", "ERROR").
const stateMap: Record<string, 'OK' | 'WARNING' | 'ERROR'> = {