-
Notifications
You must be signed in to change notification settings - Fork 659
Feat: split agent_context into submodules and add offload/reload for oversized step content #3243
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: develop
Are you sure you want to change the base?
Changes from all commits
790d3ea
94a287f
a628e76
16a176d
566d8c6
a5b522a
debefb3
37b24e9
0305481
d6d2709
8f6b3b1
5c38be4
6539899
cf2afc1
aab560f
63705b2
42b3675
a49dd10
9e65f2c
f393ca5
630aa89
f199c7a
776f4fd
1c041a2
731503b
46dd020
47f6b91
7cb5fa5
7b8cf05
dff5114
ed5c9fa
0c36c1d
0f3fa4c
e254c7b
5f8c993
eaf3dea
7202580
2b1aeb3
b0f98ba
d870fa8
f6730ff
8f7498d
83b0596
96f3189
f4638d7
2c0d4af
de7bbae
a708ea3
e292906
a705790
087672c
9a11ea4
018db44
6afbe7d
a59f0c9
c4511e9
5a49c33
5e869bd
620109f
1f8821f
8b97673
7683caf
16433a9
05815e6
9c04069
b7ff90f
a5729fe
4252520
cae5f9d
6cf6308
563e590
27855f7
23c44e0
0eb6f31
d515148
bad9b32
bdc3e43
fb3d544
b38b13b
ab8d8b0
aae8be2
d22a6c8
63be55d
e037c0e
b6f1279
a14d41c
275e80b
3ed0c1b
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -3,6 +3,7 @@ | |
| from .memory import * | ||
| from .storage import * | ||
| from .vector_database import * | ||
| from .container import * | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Wildcard import 风险:
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [代码规范] |
||
| from .skills import * | ||
|
|
||
|
|
||
|
|
||
This file was deleted.
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,29 @@ | ||
| """Agent context management for memory compression and summarization. | ||
|
|
||
| Provides ContextManager for token-aware memory compression of agent memory, | ||
| supporting incremental summarization with cache-based optimization. | ||
| """ | ||
|
|
||
| from .manager import ContextManager | ||
| from .offload_store import OffloadStore | ||
| from .summary_step import SummaryTaskStep | ||
| from .llm_summary import format_summary_output, _is_context_length_error | ||
| from .step_renderer import compress_history_offline | ||
|
|
||
| # Re-export types from sibling modules so that | ||
| # ``from agent_context import ContextManagerConfig`` still works. | ||
| from ..summary_config import ContextManagerConfig | ||
| from ..summary_cache import CompressionCallRecord, PreviousSummaryCache, CurrentSummaryCache | ||
|
|
||
| __all__ = [ | ||
| "ContextManager", | ||
| "OffloadStore", | ||
| "SummaryTaskStep", | ||
| "format_summary_output", | ||
| "_is_context_length_error", | ||
| "compress_history_offline", | ||
| "ContextManagerConfig", | ||
| "CompressionCallRecord", | ||
| "PreviousSummaryCache", | ||
| "CurrentSummaryCache", | ||
| ] |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,215 @@ | ||
| """Budget estimation, trimming, and pure data helpers for ContextManager.""" | ||
|
|
||
| import hashlib | ||
| import logging | ||
| from typing import Callable, List, Optional, Tuple | ||
|
|
||
| from smolagents.memory import ActionStep, MemoryStep, TaskStep | ||
|
|
||
| from ...utils.token_estimation import estimate_tokens_text | ||
| from ..summary_cache import PreviousSummaryCache, CurrentSummaryCache | ||
|
|
||
| logger = logging.getLogger("agent_context.budget") | ||
|
|
||
|
|
||
| # ============================================================ | ||
| # Pure data helpers (no dependencies beyond stdlib + types) | ||
| # ============================================================ | ||
|
|
||
| def extract_pairs(steps: List[MemoryStep]) -> List[tuple]: | ||
| """Extract (TaskStep, ActionStep) pairs from a step list.""" | ||
| pairs = [] | ||
| i = 0 | ||
| from .summary_step import SummaryTaskStep | ||
| while i < len(steps): | ||
| if isinstance(steps[i], TaskStep) and not isinstance(steps[i], SummaryTaskStep): | ||
| if i + 1 < len(steps) and isinstance(steps[i + 1], ActionStep): | ||
| pairs.append((steps[i], steps[i + 1])) | ||
| i += 2 | ||
| continue | ||
| i += 1 | ||
| return pairs | ||
|
|
||
|
|
||
| def action_content(action: ActionStep) -> str: | ||
| """Extract the output text from an ActionStep.""" | ||
| return action.action_output or getattr(action, "output", "") or "" | ||
|
|
||
|
|
||
| def pair_fingerprint(task_content: str, action_content: str) -> str: | ||
| """Compute a fingerprint hash for a (task, action) pair.""" | ||
| raw = (task_content[-200:] + action_content[-200:]) | ||
| return hashlib.md5(raw.encode()).hexdigest() | ||
|
|
||
|
|
||
| def action_fingerprint(action: ActionStep) -> str: | ||
| """Compute a fingerprint hash for an ActionStep.""" | ||
| raw = ( | ||
| str(action.step_number or "") | ||
| + (action.model_output or "")[-200:] | ||
| + ( | ||
| action.action_output if isinstance(action.action_output, str) | ||
| else str(action.action_output) if action.action_output else "" | ||
|
Check warning on line 52 in sdk/nexent/core/agents/agent_context/budget.py
|
||
| )[-200:] | ||
| ) | ||
| return hashlib.md5(raw.encode()).hexdigest() | ||
|
|
||
|
|
||
| def has_invoked_tools(action: ActionStep) -> bool: | ||
| """Check whether an ActionStep invokes any registered tool. | ||
|
|
||
| Unlike ``is_tool_call_step()`` which only checks for the generic | ||
| ``tool_calls is not None`` (always True for CodeAgent steps), this | ||
| function checks the ``invoked_tools`` list for actual tool names. | ||
|
|
||
| Returns True only when the step's code called at least one tool | ||
| that is registered in the agent's ``self.tools`` dict. | ||
| """ | ||
| invoked = getattr(action, "invoked_tools", None) | ||
| return bool(invoked) | ||
|
|
||
|
|
||
| def is_observation_step(action: ActionStep) -> bool: | ||
| """Check if an ActionStep is an observation step.""" | ||
| return action is not None and hasattr(action, 'observations') and action.observations is not None | ||
|
|
||
|
|
||
| def is_tool_call_step(action: ActionStep) -> bool: | ||
| """Check if an ActionStep is a tool call step.""" | ||
| return action is not None and hasattr(action, 'tool_calls') and action.tool_calls is not None | ||
|
|
||
|
|
||
| # ============================================================ | ||
| # Cache validation (depends on fingerprint functions, pure) | ||
| # ============================================================ | ||
|
|
||
| def is_prev_cache_valid( | ||
| prev_pairs: List[tuple], cache: Optional[PreviousSummaryCache], | ||
| ) -> Tuple[bool, int]: | ||
| """Check whether the previous cache covers a prefix of prev_pairs. | ||
|
|
||
| Returns (is_valid, covered_idx). When is_valid is True, | ||
| prev_pairs[0:covered_idx] can be replaced by cache.summary_text, | ||
| and prev_pairs[covered_idx:] represents the uncovered incremental portion. | ||
| """ | ||
| if cache is None or not prev_pairs: | ||
| return False, 0 | ||
| if cache.covered_pairs == 0 or cache.covered_pairs > len(prev_pairs): | ||
| return False, 0 | ||
| anchor_t, anchor_a = prev_pairs[cache.covered_pairs - 1] | ||
| fp = pair_fingerprint(anchor_t.task or "", action_content(anchor_a)) | ||
| if fp != cache.anchor_fingerprint: | ||
| return False, 0 | ||
| return True, cache.covered_pairs | ||
|
|
||
|
|
||
| def is_curr_cache_valid( | ||
| action_steps: List[ActionStep], cache: Optional[CurrentSummaryCache], | ||
| ) -> Tuple[bool, int]: | ||
| """Check whether the current cache covers a prefix of action_steps.""" | ||
| if cache is None or not action_steps: | ||
| return False, 0 | ||
| if cache.end_steps == 0 or cache.end_steps > len(action_steps): | ||
| return False, 0 | ||
| anchor = action_steps[cache.end_steps - 1] | ||
| if action_fingerprint(anchor) != cache.anchor_fingerprint: | ||
| return False, 0 | ||
| return True, cache.end_steps | ||
|
|
||
|
|
||
| # ============================================================ | ||
| # Budget trimming (depends on render_fn for text conversion) | ||
| # ============================================================ | ||
|
|
||
| def trim_pairs_to_budget( | ||
|
Check failure on line 124 in sdk/nexent/core/agents/agent_context/budget.py
|
||
| pairs: List[tuple], max_tokens: int, | ||
| render_fn: Callable[[List[tuple]], str], | ||
| keep_first: bool = True, | ||
| ) -> List[tuple]: | ||
| """Trim pairs to fit within a token budget. | ||
|
|
||
| Args: | ||
| pairs: List of (TaskStep, ActionStep) tuples. | ||
| max_tokens: Maximum token budget. | ||
| render_fn: Function to convert pairs to text (e.g. renderer.pairs_to_text). | ||
| keep_first: If True, always keep the first pair. | ||
| """ | ||
| if not pairs: | ||
| return [] | ||
| pair_tokens = [ | ||
| estimate_tokens_text(render_fn([p])) for p in pairs | ||
| ] | ||
| sep = estimate_tokens_text("\n\n") | ||
| total = sum(pair_tokens) + sep * max(0, len(pairs) - 1) | ||
| if total <= max_tokens: | ||
| return list(pairs) | ||
|
|
||
| if keep_first and len(pairs) > 1: | ||
| budget = max_tokens - pair_tokens[0] - sep | ||
| kept_tail = [] | ||
| for i in range(len(pairs) - 1, 0, -1): | ||
| cost = pair_tokens[i] + (sep if kept_tail else 0) | ||
| if cost > budget: | ||
| break | ||
| kept_tail.append(pairs[i]) | ||
| budget -= cost | ||
| return [pairs[0]] + list(reversed(kept_tail)) | ||
|
|
||
| budget = max_tokens | ||
| kept = [] | ||
| for i in range(len(pairs) - 1, -1, -1): | ||
| cost = pair_tokens[i] + (sep if kept else 0) | ||
| if cost > budget: | ||
| break | ||
| kept.append(pairs[i]) | ||
| budget -= cost | ||
| return list(reversed(kept)) if kept else [pairs[-1]] | ||
|
|
||
|
|
||
| def trim_actions_to_budget( | ||
| actions: List[ActionStep], task_text: str, max_tokens: int, | ||
| render_fn: Callable[[List[ActionStep]], str], | ||
| ) -> List[ActionStep]: | ||
| """Trim actions to fit within a token budget. | ||
|
|
||
| Args: | ||
| actions: List of ActionStep instances. | ||
| task_text: Task description text. | ||
| max_tokens: Maximum token budget. | ||
| render_fn: Function to convert actions to text (e.g. renderer.actions_to_text). | ||
| """ | ||
| if not actions: | ||
| return [] | ||
|
|
||
| def _total_tokens(acts): | ||
| return estimate_tokens_text(task_text + render_fn(acts)) | ||
|
|
||
| if _total_tokens(actions) <= max_tokens: | ||
| return list(actions) | ||
|
|
||
| for drop in range(1, len(actions) + 1): | ||
| remaining = actions[drop:] | ||
| if not remaining: | ||
| break | ||
| if is_observation_step(remaining[0]) and is_tool_call_step(actions[drop - 1]): | ||
| continue | ||
| if _total_tokens(remaining) <= max_tokens: | ||
| return list(remaining) | ||
|
|
||
| return _fallback_trim_actions(actions) | ||
|
|
||
|
|
||
| def _fallback_trim_actions(actions: List[ActionStep]) -> List[ActionStep]: | ||
| """Fallback trimming that preserves the last complete tool call pair.""" | ||
| if not actions: | ||
| return [] | ||
| last_action = actions[-1] | ||
| if len(actions) >= 2 and is_observation_step(last_action): | ||
| prev_action = actions[-2] | ||
| if is_tool_call_step(prev_action): | ||
| logger.warning( | ||
| "Fallback limit triggered: Retaining the last complete ToolCall + Observation pair intact. " | ||
| "This may exceed the token budget, and downstream truncation will be relied upon." | ||
| ) | ||
| return [prev_action, last_action] | ||
| return [last_action] | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
.gitignore中添加了sdk/nexent/core/agents/temp_scripts/,说明有调试脚本被提交到分支。即使被 gitignore,这些文件也不应存在于 PR 分支中。请在合并前清理。