diff --git a/.claude/skills/spec-coding/SKILL.md b/.claude/skills/spec-coding/SKILL.md new file mode 100644 index 000000000..b6ed2505d --- /dev/null +++ b/.claude/skills/spec-coding/SKILL.md @@ -0,0 +1,125 @@ +--- +name: spec-coding +description: Use for Nexent feature work, architecture changes, database/API changes, multi-file refactors, or any implementation that should be driven by SPEC documentation. Enforces documentation-first development through the Nexent Development SPECs Wiki: organize by implementation status, then feature scope, then lifecycle documents; update requirements, functional design, technical design, and development plan before coding. +license: Complete terms in LICENSE.txt +--- + +# Spec Coding + +Use this skill for Nexent coding work that changes product behavior, architecture, data models, APIs, persistence, runtime flows, or multiple modules. The goal is controlled implementation: document first, develop from the approved docs, and keep the Wiki as the source of truth. + +## Source of Truth + +Use the Feishu Wiki named `Nexent Development SPECs`. URL: https://dcnvjn24oieg.feishu.cn/wiki/KyU6wFj3siGJ1WkWlu8cTHgwnYb + +Top-level organization is by implementation status: + +```text +00 - Wiki Governance and Reading Guide +10 - Proposed Specs +20 - In Development Specs +30 - Implemented Specs +40 - Paused or Superseded Specs +90 - Templates and Standards +``` + +Inside each status section, organize by feature scope. Inside each feature scope, use lifecycle documents: + +```text + +├── 00 - Requirement Analysis +├── 01 - Functional Design +├── 02 - Technical Design +└── 03 - Development Plan +``` + +If a parent page has children, its body may include a `Quick Access` table, but only for direct children (depth 1). Every entry in `Quick Access` must be a clickable link to the child page. Do not maintain a global directory in page bodies; the Wiki UI already provides that. + +## Mandatory Workflow + +Before writing code: + +1. Identify the feature scope and current implementation status. +2. Locate or create the feature scope under the correct status section. +3. Ensure these lifecycle pages exist: + - `00 - Requirement Analysis` + - `01 - Functional Design` + - `02 - Technical Design` + - `03 - Development Plan` +4. Read the relevant lifecycle pages before editing code. +5. If the docs are missing or stale, update the Wiki first. +6. Only implement code that is traceable to the approved lifecycle docs. + +During coding: + +- Keep the implementation aligned with the `03 - Development Plan` PR/phase breakdown. +- If code discoveries invalidate the docs, stop broad implementation and update the relevant lifecycle page first. +- Keep acceptance criteria, tests, migrations, and compatibility requirements synchronized with the docs. + +After coding: + +- Update the relevant lifecycle page with implementation notes only when they change the plan, design, or acceptance criteria. +- Move the feature scope between status sections when its lifecycle status changes: + - `Proposed Specs` -> `In Development Specs` when implementation starts. + - `In Development Specs` -> `Implemented Specs` after implementation and acceptance. + - Any active status -> `Paused or Superseded Specs` when paused, abandoned, or replaced. + +## Lifecycle Page Responsibilities + +`00 - Requirement Analysis`: + +- Problem statement +- Goals and non-goals +- User or system impact +- Constraints +- Risks + +`01 - Functional Design`: + +- User-visible or system-visible behavior +- Capability boundaries +- Functional decomposition +- Error, empty, compatibility, and migration behavior where relevant + +`02 - Technical Design`: + +- Architecture +- Interfaces and contracts +- Data models and schema changes +- Runtime integration points +- Backward compatibility strategy + +`03 - Development Plan`: + +- Phase and PR breakdown +- File/module ownership +- Acceptance criteria +- Test plan +- Rollout and fallback notes + +## Nexent-Specific Checks + +For backend work, preserve the app/service/const layer boundaries described in `AGENTS.md`. + +For environment variables, keep `backend/consts/const.py` as the single source of truth. SDK code must not read environment variables directly. + +For DB schema changes, update all required locations: + +- versioned migration under `docker/sql/*.sql` or the project migration location in use +- Docker Compose fresh deploy init SQL +- K8s fresh deploy init SQL +- `APP_VERSION` if the project versioning rule requires it + +For tests, follow the project pytest conventions and add focused coverage matching the documented acceptance criteria. + +## When Documentation Can Be Lightweight + +Tiny mechanical fixes may use a short existing scope note instead of a full lifecycle set only when all are true: + +- The change is single-purpose and low risk. +- No API, DB, runtime contract, or user-visible behavior changes. +- No cross-module coordination is needed. +- The user explicitly wants a small fix. + +Even then, mention the relevant existing SPEC or explain why no SPEC update was needed. + diff --git a/.gitignore b/.gitignore index ebce75ae7..fc7a2044d 100644 --- a/.gitignore +++ b/.gitignore @@ -85,3 +85,7 @@ _doc/ /deploy/env/.env.bak agent_repository_frontend + +.tokensave + +.playwright-mcp/ diff --git a/AGENTS.md b/AGENTS.md index a09a54473..f33d7ee0f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -24,6 +24,12 @@ Usage notes: + +spec-coding +Use for Nexent feature work, architecture changes, database/API changes, multi-file refactors, or any implementation that should be driven by SPEC documentation. Enforces documentation-first development through the Nexent Development SPECs Wiki: organize by implementation status, then feature scope, then lifecycle documents; update requirements, functional design, technical design, and development plan before coding. +project + + prompts-writing Create, refine, and optimize high-quality YAML prompts for AI assistants. Use when working with prompt templates, system prompts, agent prompts, or any prompt engineering tasks. Provides structure guidelines, template patterns, and quality standards for YAML-based prompts. @@ -53,6 +59,31 @@ Nexent is a zero-code platform for auto-generating AI agents. Monorepo with: --- +## SPEC Coding Workflow (Mandatory) + +For any Nexent feature work, architecture change, database/API change, multi-file refactor, runtime behavior change, or other implementation that can affect product behavior, **invoke and follow the `spec-coding` skill before coding**. + +Development must be documentation-first: +- Use the Feishu Wiki `Nexent Development SPECs` as the source of truth. +- Organize SPEC documents by implementation status first, then feature Scope, then lifecycle document type. +- The expected lifecycle pages are `00 - Requirement Analysis`, `01 - Functional Design`, `02 - Technical Design`, and `03 - Development Plan`. +- Read the relevant lifecycle pages before editing code. +- If required SPEC pages are missing or stale, update the Wiki first, then implement. +- Code changes must trace back to the documented requirements, design, development plan, and acceptance criteria. +- If implementation discoveries invalidate the SPEC, update the relevant lifecycle page before continuing broad code changes. + +Only tiny mechanical fixes may skip a full SPEC update, and only when they do not change API, DB schema, runtime contracts, cross-module behavior, or user-visible behavior. In that case, state why no SPEC update was needed. + +Development must be test-verified against the documented acceptance criteria: +- Unit tests should relact the acceptance criteria and edge cases from the SPEC. +- Unit tests must reach 90% coverage for any new or modified module. +- Integration tests must verify cross-module behavior and runtime flows. +- For all frontend-affected changes, `playwright` must be used to verify user-visible behavior and acceptance criteria. +- For all backend-affected changes, `curl` or `wget` must be used to verify API behavior and acceptance criteria. +- For all SDK-affacted changes, when actual model calls are required to perform functional test, ask the user to provide one, and test with `LangFuse` to trace every step's input and output. + +--- + ## Developer Commands ### Backend (Python 3.11) @@ -165,4 +196,4 @@ Existing instruction files with detailed rules: - `CLAUDE.md` - Backend architecture, env var management, app/service layer rules - `.cursor/rules/environment_variable.mdc` - Env var centralization - `.cursor/rules/pytest_unit_test_rules.mdc` - Testing patterns -- `.cursor/rules/english_comments.mdc` - Comment language enforcement \ No newline at end of file +- `.cursor/rules/english_comments.mdc` - Comment language enforcement diff --git a/CONTEXT_MANAGEMENT_DEV_PLAN.md b/CONTEXT_MANAGEMENT_DEV_PLAN.md new file mode 100644 index 000000000..a4d80c519 --- /dev/null +++ b/CONTEXT_MANAGEMENT_DEV_PLAN.md @@ -0,0 +1,1030 @@ +# Agent 上下文管理优化 — 开发计划 + +> 基于 W8 (Progressive Component Reduction)、W12 (Release 1 History Projections)、W13 (Unified Context and Memory Policy) 三个需求文档，聚焦 SDK 层上下文管理能力的系统性提升。 + +--- + +## 1. 目标 + +**核心问题**：当前 agent 上下文管理存在三个结构性缺陷： + +1. **组件粒度粗**：`ContextComponent` 是整体选择/丢弃的最小单位，token 压力下只能整块移除（如整个 memory 组件），无法降级保留关键信息 +2. **选择逻辑分散**：memory level 过滤逻辑在 `store_memory_tool.py`、`search_memory_tool.py`、`create_agent_info.py` 三处重复；组件选择策略、预算分配、memory 决策没有统一策略引擎 +3. **历史投影缺失**：对话历史由前端每次请求回传，后端没有从持久化数据构建有界 `ContextItem` 候选集的能力 + +**目标状态**： + +- 每个上下文组件可在 `full → compressed → structured → pointer` 四级表示间渐进降级（W8） +- 统一的 `ContextPolicy` 引擎决定什么进入 prompt、什么被排除、以什么精度呈现（W13） +- 从执行历史中产出有界的、带源溯源的 `ContextItem` 候选集，替代当前的 ad-hoc 消息拼装（W12） + +--- + +## 2. 当前架构分析 + +### 2.1 上下文组装流程 + +```mermaid +flowchart TD + A[HTTP Request] --> B[agent_service.prepare_agent_run] + B --> C[create_agent_info.create_agent_config] + + C --> D[build_context_components
backend/utils/context_utils.py] + D --> D1[产出 15 个 ContextComponent 实例
system_prompt, memory, tools, skills, KB...] + + C --> E[build_memory_context
backend/services/memory_config_service.py] + E --> E1[搜索 mem0 四个 level 的 memory] + + C --> F[_resolve_input_budget
W2 安全预算] + + B --> G[ContextManager
sdk/nexent/core/agents/agent_context.py] + B --> H[ManagedContextRuntime
sdk/nexent/core/context_runtime/managed/runtime.py] + + H --> I[CoreAgent.run 循环] + + I --> J[context_runtime.prepare_run] + J --> J1[ContextManager.prepare_run_context] + J1 --> J2[build_context_messages
策略选择组件] + J1 --> J3[分离 stable system / dynamic user 消息] + J1 --> J4[写入 memory.system_prompt] + + I --> K[context_runtime.prepare_step] + K --> K1[ContextManager.assemble_final_context] + K1 --> K2[stable + dynamic + 压缩后的历史 → FinalContext] + K1 --> K3[compress_if_needed
两阶段压缩] + + I --> L[context_runtime.truncate_observation
单条观测截断] +``` + +### 2.2 关键模块清单 + +| 模块 | 文件 | 职责 | +|------|------|------| +| ContextManager | `sdk/nexent/core/agents/agent_context.py` (1787 行) | 压缩、组件注册/选择、managed 上下文组装 | +| ContextComponent | `sdk/nexent/core/agents/agent_model.py` (7 个子类) | 上下文组件类型系统 | +| ContextStrategy | `sdk/nexent/core/agents/agent_model.py` (4 个策略) | 组件选择算法 | +| ContextManagerConfig | `sdk/nexent/core/agents/summary_config.py` (123 行) | 全部配置：压缩、策略、预算、注入开关 | +| ContextRuntime | `sdk/nexent/core/context_runtime/contracts.py` (107 行) | CoreAgent 与上下文实现的协议接口 | +| ManagedContextRuntime | `sdk/nexent/core/context_runtime/managed/runtime.py` (105 行) | managed 路径适配器 | +| LegacyContextRuntime | `sdk/nexent/core/context_runtime/legacy/runtime.py` (118 行) | 旧路径回退 | +| StoreMemoryTool | `sdk/nexent/core/tools/store_memory_tool.py` | memory 写入 + level 过滤 | +| SearchMemoryTool | `sdk/nexent/core/tools/search_memory_tool.py` | memory 搜索 + level 过滤 | +| memory_service | `sdk/nexent/memory/memory_service.py` | mem0 异步 CRUD | +| build_context_components | `backend/utils/context_utils.py` | 后端侧 15 个组件拼装 | + +### 2.3 当前缺陷与改进方向 + +#### 2.3.1 上下文组装层缺陷 + +| 缺陷 | 当前状态 | W8/W12/W13 改进 | +|------|---------|-----------------| +| 组件只能整体保留或丢弃 | `ContextStrategy.select_components()` 返回完整的 `ContextComponent` 列表 | W8：引入多级表示（full/compressed/structured/pointer），组件可降级而非丢弃 | +| memory level 过滤逻辑重复 3 处 | `store_memory_tool._resolve_memory_levels()`、`search_memory_tool._resolve_memory_levels()`、`create_agent_info.py` | W13：统一为 `MemoryPolicy` 决策 | +| 组件预算是静态硬编码 | `component_budgets` 在 `ContextManagerConfig` 中固定 | W13：动态预算分配，基于策略引擎决策 | +| 没有 ContextItem 概念 | 最接近的是 `ContextComponent`（粗粒度） | W12：引入 `ContextItem` 作为有界的、带溯源的上下文候选单元 | +| 上下文选择没有统一策略 | 策略选择、memory 决策、预算分配各自独立 | W13：统一 `ContextPolicy` 引擎 | +| 压缩结果仅存内存 | `ContextManager` 的 summary cache 是进程内的 | W12：投影层为后续持久化做准备（但本期不实现 W5 event log） | + +#### 2.3.2 对话历史持久化层缺陷 + +当前 DB 存储了 ReAct 过程的大部分数据，但**缺乏结构化分组**，无法从 DB 重建完整的执行上下文。 + +**当前 DB 存储内容：** + +```mermaid +flowchart TD + A[conversation_message_t] --> B[message_role: user / assistant] + A --> C[message_content: 最终回答文本
仅 assistant] + A --> D[conversation_message_unit_t] + + D --> E[unit_type:
model_output_thinking / model_output_code /
tool / execution_logs / final_answer / ...] + D --> F[unit_content: 该单元的文本内容] + D --> G[unit_index: 前端展示排序号] +``` + +**ReAct 过程持久化缺陷：** + +| 缺陷 | 当前状态 | 影响 | +|------|---------|------| +| **无 run_id** | 同一次对话可能被多次运行，但 DB 无法区分"这次运行"和"上次运行" | 无法按 run 重建执行上下文 | +| **无 step_id** | ReAct 的每个 step（thinking → code → tool → observation）没有分组标识 | 无法将 tool 调用和对应的 execution_logs 配对 | +| **无 tool call/result 配对** | `tool` unit 存了工具名，`execution_logs` 存了结果，但没有显式关联 | 无法从 DB 重建"调用了什么工具、传了什么参数、得到了什么结果" | +| **无结构化 tool 参数** | 工具参数嵌在 `model_output_code` 的 Python 代码文本中 | 无法提取结构化的 tool call 信息 | +| **无事件时间戳** | `create_time` 是批量插入时间，不是实际事件发生时间 | 无法按时间顺序精确重建执行过程 | +| **unit_index 是展示排序** | 用于前端渲染，不是执行顺序 | 并发 run 时可能冲突 | +| **message_index 从请求历史计算** | `user_role_count * 2 + 1` | 并发 run 时会覆盖 | + +**核心问题**：当前 DB 存储的是**扁平的 UI 展示单元**，不是**结构化的执行日志**。上下文模块无法从 DB 读取历史来构建 ContextItem，因为 DB 中缺少 ReAct 的结构信息。 + +**改进方向**：在不引入 W5 event log 的前提下，增强现有 DB schema，使其能正确持久化 ReAct 全过程的结构化信息。然后基于增强后的 DB，建立 DB → ContextItem 的投影能力。 + +--- + +## 3. 目标架构 + +```mermaid +flowchart TD + subgraph W12["W12 投影层"] + P[projection] + end + + subgraph W13["ContextPolicy (W13)"] + RP[resolve_policy] + SC[select_context] + DM[decide_memory_operation] + end + + subgraph CM["ContextManager"] + ASM[组装 + 压缩] + end + + subgraph W8["ItemHandlerRegistry (W8)"] + RI["reduce_item(item, target)"] + HR[handler.reduce] + AV[AdmissibilityValidator] + end + + subgraph Handlers["Handler 实现"] + TH[ToolHandler.reduce
确定性，无 LLM] + MH[MemoryHandler.reduce
COMPRESSED 需 LLM] + SH[SkillHandler.reduce
确定性，无 LLM] + end + + P -->|"ContextItem[]"| ASM + W13 -->|SelectionDecision| ASM + ASM -->|FinalContext| LLM[LLM 调用] + ASM -->|请求降级表示| RI + RI --> HR + RI --> AV + HR --> TH + HR --> MH + HR --> SH +``` + +**数据流**： +1. W12 投影层从当前对话/agent 状态中产出 `ContextItem[]` 候选集 +2. W13 策略引擎对候选集做选择决策 → `SelectionDecision`（选中/排除/表示级别） +3. ContextManager 根据决策组装上下文，需要降级时通过 `ItemHandlerRegistry` 路由到对应 handler +4. 压缩后产出 `FinalContext` 进入模型 + +--- + +## 4. 共享接口定义 + +> 以下接口是三个 workstream 的协作契约，在 Phase 0 统一定义。 + +### 4.1 ContextItem（W12 产出，W13 消费，W8 降级） + +```python +# sdk/nexent/core/agents/context/context_item.py + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Dict, List, Optional + + +class ContextItemType(str, Enum): + """上下文项类型""" + SYSTEM_PROMPT = "system_prompt" + TOOL = "tool" + SKILL = "skill" + MEMORY = "memory" + KNOWLEDGE_BASE = "knowledge_base" + MANAGED_AGENT = "managed_agent" + EXTERNAL_AGENT = "external_agent" + HISTORY_TURN = "history_turn" # 一轮对话（user + assistant） + TOOL_CALL_RESULT = "tool_call_result" # 工具调用及结果 + WORKING_MEMORY = "working_memory" # 当前任务状态 + + +class RepresentationTier(str, Enum): + """表示精度层级（从高到低）""" + FULL = "full" # 完整内容 + COMPRESSED = "compressed" # 语义压缩（LLM 摘要） + STRUCTURED = "structured" # 最小结构化字段（名称 + 关键属性） + POINTER = "pointer" # 引用 + 元数据（决定是否加载） + + +class AuthorityTier(str, Enum): + """权威层级（高到低）""" + PLATFORM = "platform" # 平台安全策略 + TENANT = "tenant" # 租户配置 + USER = "user" # 用户指令 + WORKING_MEMORY = "working_memory" # 当前任务状态 + TOOL_RESULT = "tool_result" # 工具执行结果 + RETRIEVED_MEMORY = "retrieved_memory" # 检索的长期记忆 + SUMMARY = "summary" # 压缩摘要 + AGENT_INFERENCE = "agent_inference" # agent 推断 + + +@dataclass +class ContextItem: + """有界的、带溯源的上下文候选单元。 + + 替代当前 ContextComponent 的粗粒度模型，提供更细粒度的选择/降级能力。 + """ + item_id: str # 稳定标识符 + item_type: ContextItemType # 类型 + source_refs: List[str] = field(default_factory=list) # 源事件/数据引用 + authority_tier: AuthorityTier = AuthorityTier.AGENT_INFERENCE + minimum_fidelity: RepresentationTier = RepresentationTier.STRUCTURED + current_representation: RepresentationTier = RepresentationTier.FULL + content: Any = None # 当前表示的内容 + token_estimate: int = 0 # 当前表示的 token 估算 + metadata: Dict[str, Any] = field(default_factory=dict) # 扩展元数据 + lifecycle_status: str = "active" # active / stale / evicted + recompute_cost: Optional[int] = None # 重新计算的 token 成本 +``` + +### 4.2 SelectionDecision（W13 产出，ContextManager 消费） + +```python +# sdk/nexent/core/agents/context/policy_models.py + +@dataclass(frozen=True) +class SelectionDecision: + """上下文选择决策""" + selected_item_ids: List[str] # 选中的 item ID + excluded_item_ids: List[str] # 排除的 item ID + representation_requirements: Dict[str, RepresentationTier] # item_id → 要求的表示级别 + budget_allocations: Dict[str, int] # 类型 → 分配的 token 预算 + remaining_budget: int # 剩余预算 + conflicts: List[Dict[str, Any]] # 冲突决策记录 + reason_codes: List[str] # 决策原因码 + policy_version: str # 策略版本 + decision_fingerprint: str # 决策指纹 + + +@dataclass(frozen=True) +class MemoryDecision: + """Memory 操作决策""" + operation: str # retrieve / write / update / delete / no_write / confirm_required + allowed_scopes: List[str] # 允许的 memory level + excluded_candidates: List[str] # 排除的候选项 + conflict_decisions: List[Dict[str, Any]] # 冲突决策 + confirmation_required: Optional[Dict[str, Any]] # 需要的确认信息 + reason_codes: List[str] # 决策原因码 +``` + +### 4.3 ReductionResult（W8 产出，W13/ContextManager 消费） + +```python +# sdk/nexent/core/agents/context/reducer_models.py + +@dataclass(frozen=True) +class ReductionResult: + """组件降级结果""" + representation: RepresentationTier # 实际产出的表示级别 + source_fingerprint: str # 源内容指纹 + token_count: int # 降级后的 token 数 + generator: str # 生成器标识 + generator_version: str # 生成器版本 + admissible: bool # 是否通过准入检查 + loss_metadata: Dict[str, Any] # 丢失信息分类 + content: Any # 降级后的内容 +``` + +### 4.4 ContextItemHandler — 每种类型的选择与降级处理器 + +> **核心原则**：每种 `ContextItemType` 必须有一个 `ContextItemHandler`，负责该类型的选择评分和降级逻辑。 +> **初始实现**：所有 handler 在 PR-1 中以 **passthrough** 形式创建（选择返回全部、降级返回原内容）， +> 但必须用 `TODO` 注释标注建议算法。算法的具体实现在后续 PR 中逐步完善。 + +```python +# sdk/nexent/core/agents/context/item_handler.py + +from abc import ABC, abstractmethod + + +class ContextItemHandler(ABC): + """Per-type handler responsible for selection scoring and reduction. + + Every ContextItemType MUST have a registered handler. + Initial implementations are passthrough; suggested algorithms + are documented in TODO comments for incremental implementation. + """ + + @abstractmethod + def supported_types(self) -> List[ContextItemType]: + """Return the ContextItemType(s) this handler covers.""" + + def score( + self, + item: ContextItem, + query: str, + context: Dict[str, Any], + ) -> float: + """Return a selection score for this item (0.0 ~ 1.0). + + Higher score = more likely to be selected. + Default: passthrough, returns 1.0 (all items equally selected). + """ + return 1.0 + + def reduce( + self, + item: ContextItem, + target: RepresentationTier, + budget: int, + ) -> ReductionResult: + """Reduce item to target representation tier. + + Default: passthrough, returns original content unchanged. + """ + return ReductionResult( + representation=item.current_representation, + source_fingerprint="", + token_count=item.token_estimate, + generator="passthrough", + generator_version="0.1.0", + admissible=True, + loss_metadata={}, + content=item.content, + ) +``` + +#### 每种类型的 Handler 及建议算法（TODO） + +| Handler 类 | 覆盖类型 | 选择评分建议算法（TODO） | 降级建议算法（TODO） | +|---|---|---|---| +| `SystemPromptHandler` | `SYSTEM_PROMPT` | 不参与选择（mandatory，score=∞） | 不可降级（minimum=FULL），直接拒绝 | +| `ToolHandler` | `TOOL` | `score = priority × 0.4 + keyword_overlap(query, description) × 0.3 + usage_frequency × 0.3`；已被当前 run 调用过的 tool 额外加权 | `STRUCTURED`：模板裁剪，仅保留 name + 一句话 description + 参数名列表；`POINTER`：仅 name + 参数数量 | +| `SkillHandler` | `SKILL` | `score = keyword_overlap(query, description) × 0.6 + priority × 0.4` | `STRUCTURED`：name + description 首句截断；`POINTER`：仅 name | +| `MemoryHandler` | `MEMORY` | `score = mem0_relevance_score × 0.5 + recency × 0.2 + authority_weight × 0.3` | `COMPRESSED`：LLM 摘要（复用现有 compress prompt）；`STRUCTURED`：关键词/实体提取；`POINTER`：level + score + 前 50 字符预览 | +| `KnowledgeBaseHandler` | `KNOWLEDGE_BASE` | `score = relevance_score`（来自 KB 检索） | `COMPRESSED`：LLM 摘要；`STRUCTURED`：KB ID + 标题 + 相关度分数 | +| `ManagedAgentHandler` | `MANAGED_AGENT` | `score = keyword_overlap(query, description) × 0.5 + priority × 0.5` | `STRUCTURED`：name + routing metadata；`POINTER`：name + 能力标签 | +| `ExternalAgentHandler` | `EXTERNAL_AGENT` | 同 `ManagedAgentHandler` | 同 `ManagedAgentHandler` | +| `HistoryTurnHandler` | `HISTORY_TURN` | `score = recency × 0.5 + has_pending_action × 0.3 + keyword_overlap × 0.2` | `COMPRESSED`：LLM 摘要（复用现有压缩逻辑）；`STRUCTURED`：user query 摘要 + assistant 结论 | +| `ToolCallResultHandler` | `TOOL_CALL_RESULT` | `score = recency × 0.4 + is_active_tool × 0.4 + result_relevance × 0.2` | `STRUCTURED`：tool name + 结果摘要；`POINTER`：tool name + 状态 | +| `WorkingMemoryHandler` | `WORKING_MEMORY` | 不参与选择（mandatory，score=∞） | `STRUCTURED`：保留 active goals + constraints，丢弃 details | + +#### Handler 注册机制 + +```python +# sdk/nexent/core/agents/context/item_handler_registry.py + +class ItemHandlerRegistry: + """Registry mapping ContextItemType → ContextItemHandler. + + All handlers are registered at module load time. + Every ContextItemType MUST have exactly one handler. + """ + + _handlers: Dict[ContextItemType, ContextItemHandler] = {} + + @classmethod + def register(cls, handler: ContextItemHandler) -> None: ... + + @classmethod + def get(cls, item_type: ContextItemType) -> ContextItemHandler: ... + + @classmethod + def all_types_covered(cls) -> bool: + """Return True if every ContextItemType has a handler.""" +``` + +#### 开发约束 + +1. **PR-1 中必须创建所有 handler**（10 个类型对应 10 个 handler 类），即使大部分是 passthrough +2. 每个 handler 的 `score()` 和 `reduce()` 方法体中必须包含 `TODO` 注释，写明上表中的建议算法 +3. `ItemHandlerRegistry` 在初始化时校验所有 `ContextItemType` 都有对应 handler，缺失则报错 +4. 后续 PR（PR-4 选择引擎、PR-5/6 Reducer）通过调用 handler 的 `score()` / `reduce()` 来实现具体逻辑，而不是重新实现 +5. 每个 handler 的 TODO 注释格式： + ```python + def score(self, item, query, context): + # TODO(W13): Implement weighted scoring: + # score = priority * 0.4 + # + keyword_overlap(query, description) * 0.3 + # + usage_frequency * 0.3 + # Boost if tool was called in current run. + # Signals: item.metadata["priority"], item.metadata["usage_count"] + return 1.0 # passthrough + ``` + +--- + +## 5. 开发阶段与 PR 拆解 + +### 当前范围调整：Memory 相关工作暂缓 + +> **决策**：由于 Nexent 的 memory 模块近期将进行重构，本计划中所有依赖现有 memory 读写、level 过滤、scope 决策、memory reducer 的实现工作暂缓。 +> +> **当前阶段要求**：保持现有 memory 行为不变；只保留 ContextItem 类型覆盖和兼容性所需的最小占位接口。 + +具体调整： + +- `MEMORY` 类型和 `MemoryHandler` 继续保留，用于 handler registry 完整性校验 +- `MemoryHandler.score()` / `MemoryHandler.reduce()` 暂时保持 passthrough 或最小占位实现 +- `MemoryDecision` 模型如已存在则保留，但暂不接入运行时行为 +- 暂不新增或接入 `memory_policy.py` +- 暂不重构 `store_memory_tool.py`、`search_memory_tool.py`、`backend/agents/create_agent_info.py` 中现有 memory level 过滤逻辑 +- 暂不实现 `memory_decision_engine.py` +- 上下文选择与降级工作优先覆盖非 memory 类型：system prompt、tools、skills、knowledge base、managed/external agents、history turn、tool call result、working memory + +### 阶段依赖关系 + +```mermaid +flowchart TD + P0["Phase 0: 共享接口定义 (PR-0)"] --> P1A["Phase 1A: W12 组件投影 (PR-1)"] + P0 --> P2["Phase 2: W13 策略引擎 (PR-3, PR-4)"] + + P1A --> P1B["Phase 1B: W12 DB 历史投影 (PR-2)"] + P1A --> P2 + + P2 --> P3["Phase 3: W8 Handler 降级 (PR-5, PR-6)"] + + style P0 fill:#e1f5ff + style P1A fill:#fff4e1 + style P1B fill:#fff4e1 + style P2 fill:#e8f5e9 + style P3 fill:#fce4ec +``` + +### PR 合并策略 + +| 合并后 | 原始 PR | 合并理由 | 预估行数 | 主干影响 | +|--------|--------|---------|---------|---------| +| PR-1 | 原 PR-1 + PR-2 | 同属 W12 组件投影管线，PR-2 依赖 PR-1，顺序开发 | ~960 | 无：`use_context_items=False` 默认关闭 | +| PR-2 | 原 PR-3 + PR-4 + PR-5 | 同属 W12 DB 历史投影，PR-3→4→5 顺序依赖 | ~1480 | 无：新字段 DEFAULT NULL，旧路径不变 | +| PR-3 | 原 PR-6 + PR-7 | 同属 W13 策略基础，PR-6 是重构，PR-7 是新建 | ~770 | 无：PR-6 行为不变的重构 | +| PR-4 | 原 PR-8 + PR-9 | 同属 W13 策略引擎，都消费 PR-3 的 Policy 模型 | ~1060 | 无：新模块，旧路径不受影响 | +| PR-5 | 原 PR-10 + PR-11 | 同属 W8 确定性降级，PR-11 依赖 PR-10 | ~1000 | 无：仅影响 `use_context_items=True` 路径 | +| PR-6 | 原 PR-12 | W8 语义降级 + 集成，涉及 LLM 调用，独立交付 | ~700 | 无：仅影响 `use_context_items=True` 路径 | + +--- + +### Phase 0: 共享接口与模块基础 + +> **目标**：定义三个 workstream 的协作契约，创建新模块目录结构。 +> **预估**：1 个 PR + +#### PR-0: 共享接口定义 + +**开发内容**： + +1. 创建 `sdk/nexent/core/agents/context/` 模块目录： + ``` + sdk/nexent/core/agents/context/ + ├── __init__.py + ├── context_item.py # ContextItem, ContextItemType, RepresentationTier, AuthorityTier + ├── policy_models.py # SelectionDecision, MemoryDecision + ├── reducer_models.py # ReductionResult + ├── reason_codes.py # 稳定原因码注册表 + ├── item_handler.py # ContextItemHandler 抽象基类（见 4.4 节） + ├── item_handler_registry.py # ItemHandlerRegistry（见 4.4 节） + └── handlers/ # 每种类型的 handler 实现 + ├── __init__.py + ├── system_prompt_handler.py + ├── tool_handler.py + ├── skill_handler.py + ├── memory_handler.py + ├── knowledge_base_handler.py + ├── managed_agent_handler.py + ├── external_agent_handler.py + ├── history_turn_handler.py + ├── tool_call_result_handler.py + └── working_memory_handler.py + ``` + +2. 实现 `ContextItem` 数据模型（见 4.1 节） + +3. 实现 `SelectionDecision`、`MemoryDecision` 数据模型（见 4.2 节） + +4. 实现 `ReductionResult` 数据模型（见 4.3 节） + +5. 定义原因码注册表 `reason_codes.py`： + ```python + # 选择原因码 + SELECTED_MANDATORY_MINIMUM = "selected_mandatory_minimum" + SELECTED_BUDGET_UPGRADE = "selected_budget_upgrade" + EXCLUDED_BUDGET = "excluded_budget" + EXCLUDED_POLICY_DISABLED = "excluded_policy_disabled" + EXCLUDED_LOWER_AUTHORITY = "excluded_lower_authority" + + # Memory 原因码 + MEMORY_OPERATION_ALLOWED = "memory_operation_allowed" + MEMORY_OPERATION_DENIED = "memory_operation_denied" + CONFIRMATION_REQUIRED = "confirmation_required" + + # 降级原因码 + MINIMUM_FIDELITY_VIOLATION = "minimum_fidelity_violation" + REDUCER_FAILED = "reducer_failed" + REPRESENTATION_STALE = "representation_stale" + ``` + +6. 实现 `ContextItemHandler` 抽象基类（见 4.4 节），包含 `score()` 和 `reduce()` 的 passthrough 默认实现 + +7. 实现 `ItemHandlerRegistry`，初始化时校验所有 `ContextItemType` 都有对应 handler + +8. 创建所有 10 个 handler 类（passthrough 实现 + TODO 注释），每个 handler 的 `score()` 和 `reduce()` 方法中必须包含建议算法的 TODO 注释（具体算法见 4.4 节表格） + +9. 编写单元测试：数据模型序列化/反序列化、枚举值完整性、原因码不重复、handler 注册完整性校验 + +**验收标准**： +- [ ] `ContextItem` 可正确实例化，所有字段有类型注解和默认值 +- [ ] `SelectionDecision` 和 `MemoryDecision` 是 frozen dataclass（不可变） +- [ ] `ReductionResult` 是 frozen dataclass +- [ ] 原因码注册表无重复值 +- [ ] `ContextItemHandler` 抽象基类定义清晰，`score()` 和 `reduce()` 有 passthrough 默认实现 +- [ ] 所有 10 个 handler 类已创建，每个的 `score()` 和 `reduce()` 中有 TODO 注释写明建议算法 +- [ ] `ItemHandlerRegistry.all_types_covered()` 返回 True +- [ ] 所有新增代码有英文 docstring +- [ ] 单元测试全部通过 + +--- + +### Phase 1A: W12 组件投影基础设施 + +> **目标**：实现 ContextItem 投影管线，从 ContextComponent 桥接到 ContextItem，并集成到 ContextManager。 +> **预估**：1 个 PR（~960 行） +> **依赖**：PR-0 + +#### PR-1: ContextItem 投影器 + ContextManager 集成 + +**开发内容**： + +1. 创建 `sdk/nexent/core/agents/context/projector.py`： + - `ContextProjector` 类：将 `List[ContextComponent]` 转换为 `List[ContextItem]` + - 每个 `ContextComponent` 子类映射到一个或多个 `ContextItem` + - 映射规则： + - `SystemPromptComponent` → 1 个 `ContextItem(type=SYSTEM_PROMPT, authority=PLATFORM)` + - `ToolsComponent` → N 个 `ContextItem(type=TOOL)`，每个工具一个 item + - `SkillsComponent` → N 个 `ContextItem(type=SKILL)`，每个 skill 一个 item + - `MemoryComponent` → N 个 `ContextItem(type=MEMORY, authority=RETRIEVED_MEMORY)`，每条 memory 一个 item + - `KnowledgeBaseComponent` → 1 个 `ContextItem(type=KNOWLEDGE_BASE)` + - `ManagedAgentsComponent` → N 个 `ContextItem(type=MANAGED_AGENT)` + - `ExternalAgentsComponent` → N 个 `ContextItem(type=EXTERNAL_AGENT)` + +2. 为每个 `ContextItem` 计算 `token_estimate`（复用现有 `estimate_tokens()` 逻辑） + +3. 为每个 `ContextItem` 设置 `minimum_fidelity`： + - `SYSTEM_PROMPT` → `FULL`（不可降级） + - `TOOL` → `STRUCTURED`（至少保留 name + description） + - `SKILL` → `STRUCTURED`（至少保留 name + 简短描述） + - `MEMORY` → `STRUCTURED`（至少保留核心事实） + - `KNOWLEDGE_BASE` → `COMPRESSED` + - `MANAGED_AGENT` / `EXTERNAL_AGENT` → `STRUCTURED` + +4. 投影时通过 `ItemHandlerRegistry.get(item_type)` 验证每种产出的 ContextItem 都有对应 handler + +5. 在 `ContextManager` 中新增 `project_context_items()` 方法： + ```python + def project_context_items( + self, + components: Optional[Sequence[ContextComponent]] = None, + ) -> List[ContextItem]: + """Project registered components into ContextItem candidates.""" + ``` + +6. 在 `assemble_final_context()` 中增加 ContextItem 路径（可选启用）： + - 新增配置项 `ContextManagerConfig.use_context_items: bool = False` + - 当 `use_context_items=True` 时，走 ContextItem 投影 → 选择 → 组装路径 + - 当 `use_context_items=False` 时，保持现有 `build_context_messages()` 路径 + +7. 在 `FinalContext` 的 `ContextEvidence` 中增加 `context_items` 字段，记录投影决策 + +8. 在 `ManagedContextRuntime` 中适配新路径 + +9. 编写单元测试 + 集成测试：每种组件类型的投影正确性、token 估算、minimum_fidelity 设置、handler 覆盖校验、ContextItem 路径与旧路径的输出一致性 + +**验收标准**： +- [ ] `ContextProjector.project(components) → List[ContextItem]` 接口稳定 +- [ ] 每种 `ContextComponent` 子类都有对应的投影规则 +- [ ] 细粒度投影（tools/skills/memory 拆分为独立 item）正确工作 +- [ ] 每个 `ContextItem` 的 `token_estimate` > 0（空内容除外） +- [ ] 每种产出的 ContextItemType 在 `ItemHandlerRegistry` 中都有对应 handler +- [ ] `use_context_items=False` 时行为与当前完全一致（零回归） +- [ ] `use_context_items=True` 时，模型收到的消息内容与旧路径语义等价 +- [ ] `ContextEvidence` 中包含 `context_items` 投影信息 +- [ ] 单元测试覆盖所有 7 种组件类型 + 集成测试覆盖 managed 路径完整流程 +- [ ] 现有 `ContextManager.build_context_messages()` 行为不变（向后兼容） + +--- + +### Phase 1B: W12 对话历史持久化与投影（可与 Phase 2 并行） + +> **目标**：增强 DB schema 以正确持久化 ReAct 全过程，基于 DB 产出 `HISTORY_TURN` 和 `TOOL_CALL_RESULT` 类型的 ContextItem，并集成到 ContextManager。 +> **预估**：1 个 PR（~1480 行） +> **依赖**：PR-0 + +#### PR-2: DB Schema 增强 + 历史投影 + 集成 + +**开发内容**： + +**Part A: DB Schema 增强** + +1. 在 `conversation_message_unit_t` 表新增字段（通过 migration SQL）： + ```sql + ALTER TABLE nexent.conversation_message_unit_t ADD COLUMN run_id INTEGER DEFAULT NULL; + ALTER TABLE nexent.conversation_message_unit_t ADD COLUMN step_id INTEGER DEFAULT NULL; + ALTER TABLE nexent.conversation_message_unit_t ADD COLUMN tool_call_id VARCHAR(100) DEFAULT NULL; + ALTER TABLE nexent.conversation_message_unit_t ADD COLUMN event_time TIMESTAMP DEFAULT NULL; + ``` + + **新增字段详细说明：** + + | 字段 | 类型 | 含义 | 来源 | 写入时机 | 是否变化 | + |------|------|------|------|---------|---------| + | `run_id` | INTEGER | 该 unit 属于同一次对话中的第几次 agent 运行 | `_stream_agent_chunks()` 入口处生成，= 该 conversation 下已有 run 的最大值 + 1 | 每个 unit 创建时写入 | 写入后不再变化 | + | `step_id` | INTEGER | 该 unit 属于当前 run 中的第几个 ReAct 步骤 | `_stream_agent_chunks()` 中维护计数器，每次收到 `step_count` chunk 时递增 | 每个 unit 创建时写入 | 写入后不再变化 | + | `tool_call_id` | VARCHAR(100) | 将一次工具调用的"调用"和"结果"配对 | `_stream_agent_chunks()` 中，每次收到 `tool` chunk 时生成 UUID，同时附加到后续 `execution_logs` chunk | 仅在 `tool` 和 `execution_logs` 类型的 unit 上写入，其他类型为 NULL | 写入后不再变化 | + | `event_time` | TIMESTAMP | 该 unit 的实际产生时间（chunk 被处理的时间） | `_stream_agent_chunks()` 中每个 chunk 被处理时取 `datetime.now()` | 每个 unit 创建时写入 | 写入后不再变化 | + + **具体示例 — 一次 ReAct 执行过程：** + + 用户问"帮我查一下北京天气然后写个总结"，agent 的 ReAct 过程产生以下 chunk 流： + + ``` + run_id | step_id | tool_call_id | unit_type | unit_content | event_time + -------|---------|--------------|------------------------|-----------------------|------------------ + 1 | 1 | NULL | model_output_thinking | "让我查一下天气..." | 10:00:01.123 + 1 | 1 | NULL | model_output_code | "search_web('北京...')"| 10:00:01.456 + 1 | 1 | abc-123 | tool | "search_web" | 10:00:01.789 + 1 | 1 | abc-123 | execution_logs | "北京今天晴，25°C..." | 10:00:02.234 + 1 | 1 | NULL | step_count | "1" | 10:00:02.300 + 1 | 2 | NULL | model_output_thinking | "现在写总结..." | 10:00:03.100 + 1 | 2 | NULL | model_output_code | "write_file(...)" | 10:00:03.400 + 1 | 2 | def-456 | tool | "write_file" | 10:00:03.700 + 1 | 2 | def-456 | execution_logs | "文件已创建" | 10:00:04.100 + 1 | 2 | NULL | step_count | "2" | 10:00:04.200 + 1 | 3 | NULL | model_output_thinking | "整理结果..." | 10:00:05.000 + 1 | 3 | NULL | final_answer | "北京今天天气晴朗..." | 10:00:05.500 + ``` + + 用户第 2 次发消息时，`run_id` 变为 2，`step_id` 从 1 重新开始。 + + **增强后的 DB 视图使上下文模块能够：** + 1. 按 `run_id` 取出某次运行的全部 unit + 2. 按 `step_id` 分组为 ReAct 步骤 + 3. 按 `tool_call_id` 配对 tool call 和 result → 投影为 `TOOL_CALL_RESULT` ContextItem + 4. 按 `event_time` 精确排序 + +2. 在 `conversation_message_t` 表新增字段： + ```sql + ALTER TABLE nexent.conversation_message_t ADD COLUMN run_id INTEGER DEFAULT NULL; + ``` + +3. 修改 `backend/services/agent_service.py` 的 `_stream_agent_chunks()`： + - 维护当前 `run_id`（每次 agent run 递增）和 `step_id`（每次 `step_count` chunk 时递增） + - 为 `tool` 类型的 chunk 生成 `tool_call_id`（UUID），并将同一个 tool_call_id 传递给后续的 `execution_logs` chunk + - 为每个 unit 写入 `event_time = datetime.now()` + - 将 `run_id` 写入 assistant message 行 + +4. 修改 `save_message_unit()` 和 `create_message_unit()` 支持新字段 + +5. 同步更新 `docker/init.sql` 和 `k8s/helm/.../init.sql` + +**Part B: DB → ContextItem 历史投影** + +6. 创建 `sdk/nexent/core/agents/context/history_projector.py`： + - `HistoryProjector` 类：从 DB 查询对话历史并投影为 ContextItem + - 输入：`conversation_id` + `tenant_id` + `user_id` + - 投影三种 purpose： + - `chat_projection`：用户可见的对话历史（兼容当前 UI） + - `resume_projection`：重启恢复所需的活跃状态（目标、约束、待办、阻塞项） + - `model_context_projection`：有界的 ContextItem 候选集 + +7. `model_context_projection` 实现： + - 从 DB 查询 `conversation_message_unit_t`，按 `run_id` → `step_id` → `unit_index` 排序 + - 将 user message + assistant final_answer 配对投影为 `HISTORY_TURN` ContextItem + - 将 `tool` + `execution_logs`（通过 `tool_call_id` 配对）投影为 `TOOL_CALL_RESULT` ContextItem + - 每个 item 带 `source_refs`（指向 DB 中的 message_id + unit_id） + - 按 recency 排序（run_id + step_id），最近的步骤 authority 更高 + - `model_output_thinking` / `model_output_deep_thinking` 不投影（隐藏思维链） + +8. `resume_projection` 实现： + - 从最近一次 run 的 DB 记录中提取：最后一个 user query（活跃目标）、final_answer 中的结论、未完成的 tool call（有 tool unit 但无 execution_logs） + - 产出 `WORKING_MEMORY` 类型的 ContextItem + +9. `chat_projection` 实现： + - 从 DB 重建当前 UI 兼容的消息格式 + - 按 message_index 排序，unit 按 unit_index 排序 + - 输出格式与 `conversation_management_service.get_conversation_history_service()` 一致 + +**Part C: 集成** + +10. 将 `HistoryProjector` 集成到 `ContextManager.assemble_final_context()`： + - 当 `use_context_items=True` 时，合并组件投影（PR-1）和历史投影的 ContextItem + - 历史投影的 item 与组件投影的 item 统一进入选择流程 + +11. 添加投影决策记录： + - 每个 ContextItem 的 inclusion/exclusion 决策带 reason code + - 记录在 `ContextEvidence` 中 + +12. 编写单元测试 + 集成测试：DB → ContextItem 投影正确性、tool_call_id 配对、thinking 过滤、空历史处理、完整流程（组件 + DB 历史 → ContextItem → 选择 → FinalContext） + +**验收标准**： +- [ ] `conversation_message_unit_t` 新增 4 个字段，migration SQL 可执行 +- [ ] `conversation_message_t` 新增 `run_id` 字段 +- [ ] 每次 agent run 的 unit 都带有正确的 `run_id` 和 `step_id` +- [ ] `tool` unit 和对应的 `execution_logs` unit 共享相同的 `tool_call_id` +- [ ] `event_time` 记录实际事件时间 +- [ ] 现有 `unit_type` / `unit_content` / `unit_index` 行为不变（向后兼容） +- [ ] `docker/init.sql` 和 `k8s/helm/.../init.sql` 同步更新 +- [ ] `HistoryProjector.project(conversation_id, purpose) → List[ContextItem]` 接口稳定 +- [ ] `model_context_projection` 从 DB 正确产出 `HISTORY_TURN` 和 `TOOL_CALL_RESULT` ContextItem +- [ ] `tool` 和 `execution_logs` 通过 `tool_call_id` 正确配对 +- [ ] `model_output_thinking` / `model_output_deep_thinking` 不出现在 ContextItem 中 +- [ ] `resume_projection` 能从 DB 中提取活跃目标和未完成的 tool call +- [ ] `chat_projection` 输出与当前 `get_conversation_history_service()` 格式兼容 +- [ ] 组件投影和 DB 历史投影的 ContextItem 合并后正确进入选择流程 +- [ ] 每个 ContextItem 的决策有 reason code +- [ ] 端到端集成测试通过 +- [ ] 现有测试全部通过 + +--- + +### Phase 2: W13 统一上下文策略引擎 + +> **目标**：用统一的策略引擎替代分散的上下文选择和 memory 决策逻辑。 +> **预估**：2 个 PR（~770 + ~1060 行） +> **依赖**：PR-0；PR-1（ContextItem 模型） +> **可与 Phase 1B 并行** + +#### PR-3: Memory 统一 + Policy 模型 + +> **范围调整**：本 PR 中所有 memory 统一与 memory level 过滤重构工作暂缓。当前只实现与非 memory 上下文选择相关的 policy 模型；memory 行为保持现状。 + +**开发内容**： + +**Part A: Memory Level 过滤统一** + +> **暂缓**：以下 Part A 内容等待 memory 模块重构完成后再重新设计和实现。 + +1. 创建 `sdk/nexent/core/agents/context/memory_policy.py`： + - `resolve_memory_levels()` 函数：统一当前三处重复的 level 过滤逻辑 + - 输入：`MemoryUserConfig`（agent_share_option, disable_agent_ids, disable_user_agent_ids） + - 输出：允许的 memory level 列表 + - 区分 store 场景（候选 `["user_agent", "agent"]`）和 search 场景（候选 `["tenant", "user", "agent", "user_agent"]`） + +2. 重构 `store_memory_tool.py`：删除 `_resolve_memory_levels()` 方法，调用 `memory_policy.resolve_memory_levels()` + +3. 重构 `search_memory_tool.py`：删除 `_resolve_memory_levels()` 方法，调用 `memory_policy.resolve_memory_levels()` + +4. 重构 `backend/agents/create_agent_info.py`：被动 memory 注入路径使用 `memory_policy.resolve_memory_levels()` + +5. 编写单元测试：所有 agent_share_option 值的 level 解析、disable 列表过滤、store/search 场景差异 + +**Part B: ContextPolicy 与 MemoryPolicy 模型** + +6. 在 `sdk/nexent/core/agents/context/policy_models.py` 中扩展 `ContextPolicy` frozen dataclass + + > `MemoryPolicy` 如已在前置 PR 中存在则保留为兼容模型，但当前不要求完善字段、不要求接入运行时。 + +7. 实现 `resolve_policy()` 函数：合并 precedence（平台默认 → 租户配置 → agent 配置 → 用户请求覆盖），验证策略合法性 + +8. 实现 `validate_policy()` 函数：非法策略抛出 `PolicyInvalidError`（带 reason code） + +9. 编写单元测试：默认策略解析、合并 precedence、非法策略拒绝 + +**验收标准**： +- [ ] 当前阶段不实现 `resolve_memory_levels()`，现有 memory level 过滤逻辑保持不变 +- [ ] 当前阶段不改动 `store_memory_tool`、`search_memory_tool`、`create_agent_info.py` 的 memory 行为 +- [ ] 现有 memory 相关测试全部通过（行为不变） +- [ ] `ContextPolicy` 是 frozen dataclass +- [ ] `MemoryPolicy` 如已存在则保持兼容，但不作为当前验收重点 +- [ ] `resolve_policy()` 正确合并各层配置 +- [ ] 非法策略（预算超限、层级矛盾）在解析阶段被拒绝 +- [ ] 单元测试覆盖合并 precedence + 验证分支 + +#### PR-4: 选择引擎 + Memory 决策引擎 + +> **范围调整**：本 PR 只实现非 memory 的上下文选择引擎。Memory 操作策略路由暂缓，等待 memory 模块重构后再接入。 + +**开发内容**： + +**Part A: 上下文选择引擎** + +1. 创建 `sdk/nexent/core/agents/context/selection_engine.py`： + ```python + def select_context( + policy: ContextPolicy, + context_items: List[ContextItem], + safe_input_budget: int, + ) -> SelectionDecision: + """Two-phase selection: install mandatory, then spend remaining budget.""" + ``` + +2. 实现两阶段选择算法： + - **阶段 1**：安装所有 mandatory item 到 minimum_fidelity 表示；超出预算则返回 `mandatory_budget_impossible` + - **阶段 2**：调用 `ItemHandlerRegistry.get(item.item_type).score()` 获取评分，按 authority_tier + score 排序贪心升级（`STRUCTURED → COMPRESSED → FULL`） + - 当前 handler 的 `score()` 是 passthrough（返回 1.0），后续 PR 逐步实现具体算法 + +3. 实现冲突检测：同类型 item 内容矛盾时按 authority 层级决策，不可解冲突标记为 `authority_conflict_unresolved` + +4. 与现有 `ContextStrategy` 系统的桥接：当 `use_context_items=False` 时，现有策略系统继续工作 + +**Part B: Memory 操作策略路由** + +> **暂缓**：以下 Part B 内容等待 memory 模块重构完成后再重新设计和实现。 + +5. 创建 `sdk/nexent/core/agents/context/memory_decision_engine.py`： + ```python + def decide_memory_operation( + policy: ContextPolicy, + candidate_or_query: Any, + operation_type: str, # "retrieve" / "write" / "update" / "delete" + ) -> MemoryDecision: + """Route memory operations through unified policy engine.""" + ``` + +6. 重构 `SearchMemoryTool.forward()`：搜索前调用 `decide_memory_operation()`，根据 `allowed_scopes` 和 `excluded_candidates` 过滤 + +7. 重构 `StoreMemoryTool.forward()`：写入前调用 `decide_memory_operation()`，根据 `allowed_scopes` 和 `operation` 判断 + +8. 重构被动 memory 注入（`create_agent_info.py`）：pre-run memory search 通过 `decide_memory_operation()` 决策 + +9. 编写单元测试 + 集成测试：mandatory 安装、预算降级、authority 冲突、确定性测试、memory allow/deny/confirm-required + +**验收标准**： +- [ ] 两阶段选择算法正确实现 +- [ ] mandatory item 超出预算时返回 `mandatory_budget_impossible` +- [ ] authority 冲突按层级正确解决 +- [ ] 确定性测试：相同 `(policy, items, budget)` 产出相同的 `SelectionDecision` +- [ ] 与现有 `TokenBudgetStrategy` 行为兼容（不改变现有路径） +- [ ] 当前阶段不接入 `SearchMemoryTool`、`StoreMemoryTool`、被动 memory 注入 +- [ ] 默认策略下非 memory 上下文行为与重构前一致（零回归） +- [ ] 单元测试覆盖所有非 memory 选择分支 + +--- + +### Phase 3: W8 渐进式组件降级 + +> **目标**：在 token 压力下将组件降级到可接受的最低表示，而不是整块丢弃。 +> **预估**：2 个 PR（~1000 + ~700 行） +> **依赖**：PR-3（ContextPolicy）、PR-4（SelectionDecision 中的 representation_requirements） +> **实现方式**：实现 PR-0 中已创建的 Handler 的 `reduce()` 方法 + +#### PR-5: 准入验证 + 确定性 Handler 降级 + +**开发内容**： + +**Part A: 准入验证器与降级调度** + +1. 创建 `sdk/nexent/core/agents/context/admissibility_validator.py`：验证降级结果不低于 `minimum_fidelity`，mandatory item 不消失，源指纹匹配 + +2. 在 `ItemHandlerRegistry` 中新增 `reduce_item()` 调度方法：路由到 handler.reduce()，然后通过 AdmissibilityValidator 验证 + +**Part B: 确定性 Handler.reduce() 实现（无 LLM 调用）** + +3. 实现 `ToolHandler.reduce()`：FULL → 原内容；STRUCTURED → name + 一句话 description + 参数名列表；POINTER → name + 参数数量 + +4. 实现 `SkillHandler.reduce()`：FULL → 原内容；STRUCTURED → name + description 首句截断；POINTER → 仅 name + +5. 实现 `ManagedAgentHandler.reduce()` 和 `ExternalAgentHandler.reduce()`：FULL → 完整描述 + 工具列表；STRUCTURED → name + routing metadata；POINTER → name + 能力标签 + +6. 实现 `SystemPromptHandler.reduce()`：FULL → 原内容；其他级别 → `admissible=False`（不可降级） + +7. 实现 `WorkingMemoryHandler.reduce()`：FULL → 原内容；STRUCTURED → 保留 active goals + constraints，丢弃 details + +8. 编写单元测试：准入验证、路由正确、每个 handler 每个级别输出正确、token 估算、source_fingerprint + +**验收标准**： +- [ ] `AdmissibilityValidator` 正确拒绝非法降级 +- [ ] `ItemHandlerRegistry.reduce_item()` 正确路由到 handler 并执行验证 +- [ ] 5 个 handler 的 `reduce()` 从 passthrough 升级为确定性实现 +- [ ] 所有降级是确定性的（无 LLM 调用，相同输入 = 相同输出） +- [ ] token 估算误差 < 20% +- [ ] 低于 minimum_fidelity 的降级被拒绝 +- [ ] 单元测试覆盖所有验证分支 + 每个 handler 的每个表示级别 + +#### PR-6: 语义 Handler 降级 + ContextManager 集成 + +**开发内容**： + +1. 实现 `MemoryHandler.reduce()`：FULL → 原内容；COMPRESSED → LLM 摘要（复用现有 compress prompt）；STRUCTURED → 核心事实 key-value；POINTER → level + 时间戳 + 前 50 字符预览 + + > **暂缓**：等待 memory 模块重构完成后再实现。当前 `MemoryHandler.reduce()` 保持 passthrough 或最小占位行为。 + +2. 实现 `KnowledgeBaseHandler.reduce()`：FULL → 原内容；COMPRESSED → LLM 摘要；STRUCTURED → KB ID + 标题 + 相关度分数 + +3. 实现 `HistoryTurnHandler.reduce()`：FULL → 原内容；COMPRESSED → LLM 摘要（复用现有压缩逻辑）；STRUCTURED → user query 摘要 + assistant 结论 + +4. 实现 `ToolCallResultHandler.reduce()`：FULL → 原内容；STRUCTURED → tool name + 结果摘要；POINTER → tool name + 状态 + +5. 将降级流程集成到 `ContextManager.assemble_final_context()` 的 ContextItem 路径：根据 `SelectionDecision.representation_requirements` 确定目标表示级别，通过 `ItemHandlerRegistry.reduce_item()` 执行降级，降级失败时 fallback 到原始内容 + 记录 warning + +6. 编写集成测试：完整流程（投影 → 选择 → 降级 → FinalContext）、token 预算压力下的降级行为、降级 fallback + +**验收标准**： +- [ ] 非 memory 语义 handler 的 `reduce()` 从 passthrough 升级为语义实现 +- [ ] `MemoryHandler.reduce()` 当前保持 passthrough 或最小占位行为，不接入 LLM 摘要 +- [ ] LLM 调用复用现有压缩 prompt（不引入新 prompt） +- [ ] 集成到 `ContextManager` 的 ContextItem 路径 +- [ ] token 预算压力下组件正确降级而非丢弃 +- [ ] 降级失败时 fallback 到原始内容 +- [ ] 端到端集成测试通过 +- [ ] 现有测试无回归 + +--- + +## 6. PR 总览与并行策略 + +### PR 依赖图 + +```mermaid +flowchart TD + PR0["PR-0
共享接口 + Handler 骨架"] --> PR1["PR-1
W12 组件投影"] + PR0 --> PR3["PR-3
W13 Memory 统一 + Policy"] + + PR1 --> PR2["PR-2
W12 DB 历史投影"] + PR3 --> PR4["PR-4
W13 选择 + Memory 决策"] + + PR2 --> PR5["PR-5
W8 准入验证 + 确定性降级"] + PR4 --> PR5 + + PR5 --> PR6["PR-6
W8 语义降级 + 集成"] + + style PR0 fill:#e1f5ff + style PR1 fill:#fff4e1 + style PR2 fill:#fff4e1 + style PR3 fill:#e8f5e9 + style PR4 fill:#e8f5e9 + style PR5 fill:#fce4ec + style PR6 fill:#fce4ec +``` + +### 并行开发建议 + +| 阶段 | 可并行的 PR | 说明 | +|------|------------|------| +| Phase 0 完成后 | PR-1 ∥ PR-3 | W12 组件投影和 W13 非 memory Policy 基础无依赖 | +| Phase 1 中期 | PR-2 ∥ PR-4 | W12 DB 历史投影和 W13 非 memory 选择引擎可并行 | +| Phase 3 | PR-5 → PR-6 | 顺序执行，PR-6 依赖 PR-5 | + +### 预计 PR 数量 + +| Workstream | PR 数量 | 核心交付 | +|-----------|---------|---------| +| Phase 0 | 1 | 共享接口 + 10 个 Handler 骨架（passthrough + TODO） | +| W12 | 2 | 组件投影 + DB Schema 增强 + DB→ContextItem 历史投影 | +| W13 | 2 | Policy 模型 + 非 memory 选择引擎（memory 决策暂缓） | +| W8 | 2 | 准入验证 + 确定性降级 + 非 memory 语义降级 + 集成 | +| **合计** | **7** | | + +--- + +## 7. 全局验收标准 + +### 功能验收 + +- [ ] `use_context_items=False` 时，所有现有行为零回归 +- [ ] `use_context_items=True` 时，模型收到的上下文质量不低于旧路径 +- [ ] DB 正确持久化 ReAct 全过程（run_id、step_id、tool_call_id 配对、event_time） +- [ ] 上下文模块能从 DB 读取历史并构建 ContextItem（不依赖前端回传） +- [ ] 当前阶段不重构 memory level 过滤逻辑，现有 memory 行为保持不变 +- [ ] 每种 `ContextItemType` 都有对应的 `ContextItemHandler`，具备 `score()` 和 `reduce()` 方法 +- [ ] token 预算压力下，非 memory 组件通过 handler 降级而非整块丢弃 +- [ ] 每个上下文选择决策有 reason code 可追溯 + +### 测试验收 + +- [ ] 所有新增代码有对应单元测试 +- [ ] 所有重构代码的现有测试通过（零回归） +- [ ] 端到端集成测试覆盖完整流程 +- [ ] 确定性测试：相同输入产出相同决策 + +### 代码质量 + +- [ ] 所有注释和 docstring 使用英文 +- [ ] 新增模块遵循 SDK 层规范（不读取环境变量） +- [ ] 类型注解完整，无 `as any` / `@ts-ignore` 等价操作 +- [ ] `lsp_diagnostics` 无 error + +### 性能验收 + +- [ ] ContextItem 投影延迟 < 50ms（100 个组件） +- [ ] 策略引擎决策延迟 < 10ms +- [ ] 确定性 Handler.reduce() 延迟 < 5ms/个 +- [ ] 语义 Handler.reduce()（LLM 调用）延迟 < 2s/个 + +--- + +## 8. 风险与缓解 + +| 风险 | 影响 | 缓解措施 | +|------|------|---------| +| ContextItem 投影改变了消息顺序 | 模型行为变化 | PR-1 中做 A/B 对比测试，确保语义等价 | +| DB Schema 变更影响现有持久化逻辑 | 消息/单元写入异常 | PR-2 中新增字段全部 DEFAULT NULL，现有写入路径不受影响；全量回归测试 | +| Memory 模块即将重构 | 当前实现可能与新 memory 架构冲突 | 当前计划暂缓 memory policy、memory decision engine、memory reducer；只保留最小兼容占位 | +| LLM 降级调用增加延迟 | 响应变慢 | 确定性 Handler.reduce() 优先（PR-5）；语义 Handler.reduce() 仅在必要时调用（PR-6） | +| Handler passthrough 阶段功能不完整 | 降级无实际效果 | passthrough 是有意设计，保证框架可用；具体算法在后续 PR 逐步实现，每个 PR 独立可交付 | +| 策略引擎过于复杂 | 维护成本高 | 分阶段实现，每个 PR 独立可用 | +| 与 W5 event log 的后续集成 | 需要适配 | ContextItem 的 `source_refs` 字段预留了事件引用；DB 新增的 run_id/step_id 与 W5 的 agent_session/event_seq 概念对齐 | + +--- + +## 9. Memory 重构后的回补提醒 + +Memory 模块重构完成后，需要重新审视并回补本计划中暂缓的 memory 相关能力。回补时不要直接照搬旧设计，应先对齐新的 memory 架构、scope 模型、权限模型和数据结构。 + +待回补事项： + +- 重新设计 `MemoryPolicy`，明确 memory scope、authority、tenant/user/agent 边界 +- 实现统一的 memory level/scope 解析，替代当前分散在 `store_memory_tool.py`、`search_memory_tool.py`、`create_agent_info.py` 的逻辑 +- 实现 `memory_policy.py` 或等价模块，并迁移 memory 读写路径 +- 实现 `memory_decision_engine.py`，将 retrieve/write/update/delete 操作纳入统一策略决策 +- 完善 `MemoryHandler.score()`，结合新 memory relevance、recency、authority、scope 信号进行选择评分 +- 完善 `MemoryHandler.reduce()`，支持符合新 memory 数据结构的 compressed/structured/pointer 表示 +- 补充 memory policy、memory decision、memory reducer 的单元测试和端到端回归测试 +- 更新本计划的 PR 拆分和验收标准，移除“memory 暂缓”的临时约束 diff --git a/backend/agents/create_agent_info.py b/backend/agents/create_agent_info.py index 2158a5dec..686af9dab 100644 --- a/backend/agents/create_agent_info.py +++ b/backend/agents/create_agent_info.py @@ -667,6 +667,7 @@ async def create_agent_config( override_model_id: int | None = None, request_requested_output_tokens: int | None = None, tool_params: Optional[ToolParamsRequest | Dict[str, Any]] = None, + conversation_id: Optional[int] = None, ): normalized_tool_params = _normalize_tool_params_request(tool_params) agent_info = search_agent_info_by_agent_id( @@ -980,6 +981,7 @@ async def create_agent_config( capacity_snapshot=capacity_snapshot, safe_input_budget_snapshot=safe_input_budget_snapshot, verification_config=AgentVerificationConfig.model_validate(agent_info.get("verification_config") or {}), + conversation_id=conversation_id, ) return agent_config @@ -1392,6 +1394,7 @@ async def create_agent_run_info( override_version_no: int | None = None, override_model_id: int | None = None, requested_output_tokens: int | None = None, + conversation_id: Optional[int] = None, tool_params: Optional[ToolParamsRequest | Dict[str, Any]] = None, ): # Determine which version_no to use based on is_debug flag @@ -1421,6 +1424,7 @@ async def create_agent_run_info( "last_user_query": final_query, "allow_memory_search": allow_memory_search, "version_no": version_no, + "conversation_id": conversation_id, } if override_model_id is not None: create_config_kwargs["override_model_id"] = override_model_id diff --git a/backend/database/conversation_db.py b/backend/database/conversation_db.py index a82c00d24..53c6976a2 100644 --- a/backend/database/conversation_db.py +++ b/backend/database/conversation_db.py @@ -189,7 +189,8 @@ def create_message_units(message_units: List[Dict[str, Any]], message_id: int, c def create_message_unit(message_id: int, conversation_id: int, unit_index: int, unit_type: str, unit_content: str, user_id: Optional[str] = None, - unit_status: str = 'completed') -> int: + unit_status: str = 'completed', + step_index: Optional[int] = None) -> int: """ Insert a single ConversationMessageUnit row. @@ -201,6 +202,7 @@ def create_message_unit(message_id: int, conversation_id: int, unit_index: int, unit_content: Complete content of the unit user_id: Reserved parameter for created_by and updated_by fields unit_status: Lifecycle status (streaming / completed) + step_index: Optional ReAct step sequence number within this message Returns: int: Newly created unit ID (auto-increment ID) @@ -219,6 +221,8 @@ def create_message_unit(message_id: int, conversation_id: int, unit_index: int, "unit_status": unit_status, "delete_flag": 'N', } + if step_index is not None: + row_data["step_index"] = step_index if user_id: row_data["created_by"] = user_id row_data["updated_by"] = user_id @@ -418,6 +422,39 @@ def get_message_units(message_id: int) -> List[Dict[str, Any]]: return list(map(as_dict, records)) +def get_message_units_by_message(conversation_id: int, message_id: Optional[int] = None) -> List[Dict[str, Any]]: + """ + Get all units for a conversation, optionally filtered by message_id. + Used by HistoryProjector to reconstruct ReAct execution timeline. + + Args: + conversation_id: Conversation ID (integer) + message_id: Optional message ID to filter by. If None, returns all units. + + Returns: + List[Dict[str, Any]]: List of message units, ordered by message_id, step_index, unit_index + """ + with get_db_session() as session: + conversation_id = int(conversation_id) + + stmt = select(ConversationMessageUnit).where( + ConversationMessageUnit.conversation_id == conversation_id, + ConversationMessageUnit.delete_flag == 'N' + ) + + if message_id is not None: + stmt = stmt.where(ConversationMessageUnit.message_id == message_id) + + stmt = stmt.order_by( + asc(ConversationMessageUnit.message_id), + asc(ConversationMessageUnit.step_index), + asc(ConversationMessageUnit.unit_index) + ) + + records = session.scalars(stmt).all() + return list(map(as_dict, records)) + + def get_conversation_list(user_id: Optional[str] = None) -> List[Dict[str, Any]]: """ Get list of all undeleted conversations, sorted by creation time in descending order diff --git a/backend/database/db_models.py b/backend/database/db_models.py index b0d6d710f..d26f201ec 100644 --- a/backend/database/db_models.py +++ b/backend/database/db_models.py @@ -91,6 +91,8 @@ class ConversationMessageUnit(TableBase): unit_status = Column( String(30), default='completed', doc="Lifecycle status: streaming (still aggregating) or completed (fully persisted)") + step_index = Column( + Integer, doc="ReAct step sequence number within this message. Increments on step_count chunks") class ConversationSourceImage(TableBase): diff --git a/backend/services/agent_service.py b/backend/services/agent_service.py index 5ba8f548a..5066df472 100644 --- a/backend/services/agent_service.py +++ b/backend/services/agent_service.py @@ -895,8 +895,9 @@ async def _stream_agent_chunks( # so that units saved incrementally have a valid message_id to reference. streaming_message_id: Optional[int] = resume_message_id if not is_resume_mode and not agent_request.is_debug: + history_items = getattr(agent_request, "history", None) or [] user_role_count = sum( - 1 for item in getattr(agent_request, "history", []) + 1 for item in history_items if item.role == MESSAGE_ROLE["USER"] ) assistant_message_req = MessageRequest( @@ -917,6 +918,12 @@ async def _stream_agent_chunks( logger.error( "Failed to create streaming message row: %r", msg_exc, exc_info=True) + # History projection tracking state + current_step_id: int = 0 + pending_tool_content: Optional[str] = None # Buffered tool content for merge + pending_tool_step: Optional[int] = None # Step index of buffered tool + pending_tool_unit_index: Optional[int] = None # Unit index assigned to buffered tool + # Tracks the unit currently being accumulated in memory. Each entry is # a dict with keys: type, content, unit_id, unit_index, mergeable. current_unit: Optional[Dict[str, Any]] = None @@ -949,6 +956,9 @@ async def _stream_agent_chunks( chunk_type = data.get("type") chunk_content = data.get("content", "") or "" + if chunk_type == ProcessType.STEP_COUNT.value: + current_step_id += 1 + # Add unit_index to the chunk data for frontend resume skip logic. # This allows frontend to accurately skip chunks that were already persisted. # For mergeable types (continuing chunks), use the current unit's index. @@ -1095,6 +1105,7 @@ async def _stream_agent_chunks( unit_content='{"placeholder": true}', user_id=user_id, unit_status="completed", + step_index=current_step_id, ).result() try: search_results = json.loads(chunk_content) @@ -1140,6 +1151,47 @@ async def _stream_agent_chunks( yield f"data: {chunk}\n\n" continue + # Handle tool call merging: buffer TOOL, merge with EXECUTION_LOGS + if chunk_type == ProcessType.TOOL.value: + pending_tool_content = chunk_content + pending_tool_step = current_step_id + pending_tool_unit_index = next_unit_index + next_unit_index += 1 + await channel.publish(f"data: {chunk}\n\n") + yield f"data: {chunk}\n\n" + continue + + if chunk_type == ProcessType.EXECUTION_LOGS.value and pending_tool_content is not None: + merged_content = json.dumps({ + "tool_call": pending_tool_content, + "execution_result": chunk_content, + }) + if streaming_message_id is not None: + new_unit_id = submit( + save_message_unit, + message_id=streaming_message_id, + conversation_id=agent_request.conversation_id, + unit_index=pending_tool_unit_index, + unit_type="tool_call", + unit_content=merged_content, + user_id=user_id, + unit_status="completed", + step_index=pending_tool_step, + ).result() + current_unit = { + "type": "tool_call", + "content": merged_content, + "unit_id": new_unit_id, + "unit_index": pending_tool_unit_index, + "mergeable": False, + } + pending_tool_content = None + pending_tool_step = None + pending_tool_unit_index = None + await channel.publish(f"data: {chunk}\n\n") + yield f"data: {chunk}\n\n" + continue + # Default path: insert a new unit row with unit_status='streaming'. if streaming_message_id is not None and chunk_type not in ( "search_content_placeholder", @@ -1153,6 +1205,7 @@ async def _stream_agent_chunks( unit_content=chunk_content, user_id=user_id, unit_status="streaming", + step_index=current_step_id, ).result() current_unit = { "type": chunk_type, @@ -1166,6 +1219,21 @@ async def _stream_agent_chunks( await channel.publish(f"data: {chunk}\n\n") yield f"data: {chunk}\n\n" stream_completed_normally = True + + # Flush any buffered tool content that didn't receive execution_logs + if pending_tool_content is not None and streaming_message_id is not None: + submit( + save_message_unit, + message_id=streaming_message_id, + conversation_id=agent_request.conversation_id, + unit_index=pending_tool_unit_index, + unit_type="tool", + unit_content=pending_tool_content, + user_id=user_id, + unit_status="completed", + step_index=pending_tool_step, + ).result() + pending_tool_content = None except Exception as run_exc: logger.error("Agent run error: %r", run_exc, exc_info=True) await channel.publish(_safe_agent_stream_error_chunk()) @@ -2653,12 +2721,20 @@ async def prepare_agent_run( override_model_id=agent_request.model_id, requested_output_tokens=agent_request.requested_output_tokens, tool_params=agent_request.tool_params, + conversation_id=agent_request.conversation_id, ) # Mount conversation-level reusable ContextManager if enabled cm_config = getattr(agent_run_info.agent_config, 'context_manager_config', None) if cm_config and cm_config.enabled: + from nexent.core.agents.context.history_projector import HistoryProjector + from services.conversation_management_service import get_message_units_by_message_id + + cm_config.history_projector = HistoryProjector( + query_units_fn=get_message_units_by_message_id + ) + cm = agent_run_manager.get_or_create_context_manager( conversation_id=str(agent_request.conversation_id), config=cm_config, diff --git a/backend/services/conversation_management_service.py b/backend/services/conversation_management_service.py index 64482416c..36320061f 100644 --- a/backend/services/conversation_management_service.py +++ b/backend/services/conversation_management_service.py @@ -23,6 +23,7 @@ get_latest_assistant_message_id, get_last_unit_for_message, get_message_id_by_index, + get_message_units_by_message, get_source_images_by_conversation, get_source_images_by_message, get_source_searches_by_conversation, @@ -104,7 +105,8 @@ def save_message(request: MessageRequest, user_id: str, tenant_id: str, def save_message_unit(message_id: int, conversation_id: int, unit_index: int, unit_type: str, unit_content: str, user_id: Optional[str] = None, - unit_status: str = 'completed') -> int: + unit_status: str = 'completed', + step_index: Optional[int] = None) -> int: """ Insert exactly one ConversationMessageUnit row. @@ -116,6 +118,7 @@ def save_message_unit(message_id: int, conversation_id: int, unit_index: int, unit_content: Complete content of the unit user_id: Identifier of the user creating the unit unit_status: Lifecycle status (streaming / completed) + step_index: Optional ReAct step sequence number within this message Returns: int: Newly created unit_id @@ -128,9 +131,25 @@ def save_message_unit(message_id: int, conversation_id: int, unit_index: int, unit_content=unit_content, user_id=user_id, unit_status=unit_status, + step_index=step_index, ) +def get_message_units_by_message_id(conversation_id: int, message_id: Optional[int] = None) -> List[Dict[str, Any]]: + """Get all units for a conversation, optionally filtered by message_id. + + Used by HistoryProjector to reconstruct ReAct execution timeline. + + Args: + conversation_id: Conversation ID + message_id: Optional message ID to filter by + + Returns: + List of unit dictionaries ordered by message_id, step_index, unit_index + """ + return get_message_units_by_message(conversation_id, message_id) + + def update_message_status(message_id: int, status: str, user_id: str) -> None: """Update the lifecycle status of a conversation message.""" update_conversation_message_status(message_id, status, user_id=user_id) diff --git a/deploy/sql/init.sql b/deploy/sql/init.sql index 4dba737bf..a18be741d 100644 --- a/deploy/sql/init.sql +++ b/deploy/sql/init.sql @@ -45,6 +45,7 @@ CREATE TABLE IF NOT EXISTS "conversation_message_unit_t" ( "update_time" timestamp(0) DEFAULT CURRENT_TIMESTAMP, "updated_by" varchar(100) COLLATE "pg_catalog"."default", "created_by" varchar(100) COLLATE "pg_catalog"."default", + "step_index" int4, CONSTRAINT "conversation_message_unit_t_pk" PRIMARY KEY ("unit_id") ); ALTER TABLE "conversation_message_unit_t" OWNER TO "root"; @@ -58,6 +59,7 @@ COMMENT ON COLUMN "conversation_message_unit_t"."create_time" IS 'Creation time, COMMENT ON COLUMN "conversation_message_unit_t"."update_time" IS 'Update time, audit field'; COMMENT ON COLUMN "conversation_message_unit_t"."updated_by" IS 'Last updater ID, audit field'; COMMENT ON COLUMN "conversation_message_unit_t"."created_by" IS 'Creator ID, audit field'; +COMMENT ON COLUMN "conversation_message_unit_t"."step_index" IS 'ReAct step sequence number within this message. Increments on step_count chunks'; COMMENT ON TABLE "conversation_message_unit_t" IS 'Carries agent output content in each message'; CREATE TABLE IF NOT EXISTS "conversation_record_t" ( diff --git a/deploy/sql/migrations/v2.3.0_0703_history_projection_fields.sql b/deploy/sql/migrations/v2.3.0_0703_history_projection_fields.sql new file mode 100644 index 000000000..e15c38cbf --- /dev/null +++ b/deploy/sql/migrations/v2.3.0_0703_history_projection_fields.sql @@ -0,0 +1,37 @@ +-- Migration: Add step_index for ReAct step tracking +-- Date: 2026-07-03 (revised 2026-07-09) +-- Description: Add step_index column to conversation_message_unit_t. +-- Drops previously added run_id, tool_call_id, event_time columns +-- that are no longer needed after review. + +SET search_path TO nexent; +BEGIN; + +-- Add step_index (renamed from step_id) +ALTER TABLE nexent.conversation_message_unit_t + ADD COLUMN IF NOT EXISTS step_index INTEGER DEFAULT NULL; + +COMMENT ON COLUMN nexent.conversation_message_unit_t.step_index IS + 'ReAct step sequence number within this message. Increments on step_count chunks.'; + +-- Drop columns from previous revision (idempotent) +ALTER TABLE nexent.conversation_message_unit_t + DROP COLUMN IF EXISTS run_id; +ALTER TABLE nexent.conversation_message_unit_t + DROP COLUMN IF EXISTS step_id; +ALTER TABLE nexent.conversation_message_unit_t + DROP COLUMN IF EXISTS tool_call_id; +ALTER TABLE nexent.conversation_message_unit_t + DROP COLUMN IF EXISTS event_time; +ALTER TABLE nexent.conversation_message_t + DROP COLUMN IF EXISTS run_id; + +-- Drop obsolete indexes +DROP INDEX IF EXISTS nexent.idx_message_unit_conversation_run; +DROP INDEX IF EXISTS nexent.idx_message_unit_tool_call; + +-- New index for step-based queries +CREATE INDEX IF NOT EXISTS idx_message_unit_message_step + ON nexent.conversation_message_unit_t (message_id, step_index); + +COMMIT; diff --git a/frontend/app/[locale]/chat/streaming/chatStreamHandler.tsx b/frontend/app/[locale]/chat/streaming/chatStreamHandler.tsx index 244a1a239..f5ce5536b 100644 --- a/frontend/app/[locale]/chat/streaming/chatStreamHandler.tsx +++ b/frontend/app/[locale]/chat/streaming/chatStreamHandler.tsx @@ -213,6 +213,7 @@ const isSkippedUnitType = (unitType: string): boolean => { 'execution_logs', 'agent_new_run', 'tool', + 'tool_call', 'verification', 'memory_search', 'max_steps_reached', @@ -447,6 +448,7 @@ export const handleStreamResponse = async ( messageType === chatConfig.messageTypes.PARSE || messageType === chatConfig.messageTypes.EXECUTION_LOGS || messageType === chatConfig.messageTypes.TOOL || + messageType === chatConfig.messageTypes.TOOL_CALL || messageType === chatConfig.messageTypes.CARD || messageType === chatConfig.messageTypes.AGENT_NEW_RUN || messageType === chatConfig.messageTypes.VERIFICATION || @@ -860,6 +862,61 @@ export const handleStreamResponse = async ( // Execution result message, skip break; + case chatConfig.messageTypes.TOOL_CALL: { + if (!currentStep) { + currentStep = { + id: `step-tool-call-${Date.now()}-${Math.random() + .toString(36) + .substring(2, 9)}`, + title: "Tool Call", + content: "", + expanded: true, + contents: [], + metrics: null, + thinking: { content: "", expanded: true }, + code: { content: "", expanded: true }, + output: { content: "", expanded: true }, + }; + } + + try { + const toolCallData = JSON.parse(messageContent); + currentStep.contents.push({ + id: `executing-${Date.now()}-${Math.random() + .toString(36) + .substring(2, 7)}`, + type: chatConfig.messageTypes.EXECUTING, + content: toolCallData.tool_call || "", + expanded: true, + timestamp: Date.now(), + isLoading: true, + }); + currentStep.contents.push({ + id: `execution-${Date.now()}-${Math.random() + .toString(36) + .substring(2, 7)}`, + type: chatConfig.messageTypes.EXECUTION, + content: toolCallData.execution_result || "", + expanded: true, + timestamp: Date.now(), + }); + } catch { + currentStep.contents.push({ + id: `executing-${Date.now()}-${Math.random() + .toString(36) + .substring(2, 7)}`, + type: chatConfig.messageTypes.EXECUTING, + content: messageContent, + expanded: true, + timestamp: Date.now(), + isLoading: true, + }); + } + + lastContentType = chatConfig.contentTypes.EXECUTION; + break; + } + case chatConfig.messageTypes.AGENT_NEW_RUN: // If there's no currentStep, create one if (!currentStep) { diff --git a/frontend/const/chatConfig.ts b/frontend/const/chatConfig.ts index c206fa752..1180dd734 100644 --- a/frontend/const/chatConfig.ts +++ b/frontend/const/chatConfig.ts @@ -112,8 +112,9 @@ messageTypes: { FINAL_ANSWER: "final_answer" as const, PARSE: "parse" as const, TOOL: "tool" as const, - EXECUTION_LOGS: "execution_logs" as const, - ERROR: "error" as const, + EXECUTION_LOGS: "execution_logs" as const, + TOOL_CALL: "tool_call" as const, + ERROR: "error" as const, STEP_COUNT: "step_count" as const, TOKEN_COUNT: "token_count" as const, MAX_STEPS_REACHED: "max_steps_reached" as const, diff --git a/frontend/lib/chatMessageExtractor.ts b/frontend/lib/chatMessageExtractor.ts index 5a1d3afb3..d05fe7fb8 100644 --- a/frontend/lib/chatMessageExtractor.ts +++ b/frontend/lib/chatMessageExtractor.ts @@ -196,6 +196,50 @@ export function extractAssistantMsgFromResponse( break; } + case chatConfig.messageTypes.TOOL_CALL: { + // Merged tool_call: JSON with both tool call and execution result + const currentStep = getOrCreateCurrentStep(steps, "Tool Call"); + try { + const toolCallData = JSON.parse(msg.content); + // Add tool call content + const toolContentId = `tool-${Date.now()}-${Math.random() + .toString(36) + .substring(2, 7)}`; + currentStep.contents.push({ + id: toolContentId, + type: "executing", + content: toolCallData.tool_call || "", + expanded: true, + timestamp: Date.now(), + }); + // Add execution result + const execContentId = `execution-${Date.now()}-${Math.random() + .toString(36) + .substring(2, 7)}`; + currentStep.contents.push({ + id: execContentId, + type: "execution", + content: toolCallData.execution_result || "", + expanded: true, + timestamp: Date.now(), + }); + } catch { + // Fallback: treat as plain text + const contentId = `tool-${Date.now()}-${Math.random() + .toString(36) + .substring(2, 7)}`; + currentStep.contents.push({ + id: contentId, + type: "executing", + content: msg.content, + expanded: true, + timestamp: Date.now(), + }); + } + resetModelOutputTracking(); + break; + } + case chatConfig.messageTypes.ERROR: { const currentStep = getOrCreateCurrentStep(steps, "Error"); const contentId = `error-${Date.now()}-${Math.random() diff --git a/sdk/benchmark/acon_eval/dataset.py b/sdk/benchmark/acon_eval/dataset.py index ce3280381..b0e815233 100644 --- a/sdk/benchmark/acon_eval/dataset.py +++ b/sdk/benchmark/acon_eval/dataset.py @@ -1,79 +1,79 @@ - -"""Dataset loader for ACON's 8-objective QA benchmark (nq_multi_8). - -Adapted from ACON's experiments/smolagents/dataset.py. -Supports JSONL format with fields: id, question, answer. -""" -import json -from dataclasses import dataclass -from pathlib import Path -from typing import Any, Dict, Iterable, List, Optional - - -@dataclass -class QAExample: - id: str - question: str - answer: Any # str or list[list[str]] — each sub-answer is a list of acceptable variants - contexts: Optional[List[str]] = None - - -class QALoader: - def __init__(self, data_path: str): - self.path = Path(data_path) - if not self.path.exists(): - raise FileNotFoundError(f"Data file not found: {self.path}") - self.is_jsonl = self.path.suffix.lower() in {".jsonl", ".jl"} - - def count(self, limit: Optional[int] = None) -> int: - if self.is_jsonl: - total = 0 - with self.path.open("r", encoding="utf-8") as f: - for line in f: - if line.strip(): - total += 1 - else: - data = json.loads(self.path.read_text(encoding="utf-8")) - if isinstance(data, dict) and "data" in data: - data = data["data"] - total = len(data) - - if limit is not None: - total = min(total, limit) - return total - - def _iter_jsonl(self) -> Iterable[Dict[str, Any]]: - with self.path.open("r", encoding="utf-8") as f: - for line in f: - line = line.strip() - if not line: - continue - yield json.loads(line) - - def _iter_json(self) -> Iterable[Dict[str, Any]]: - data = json.loads(self.path.read_text(encoding="utf-8")) - if isinstance(data, dict) and "data" in data: - data = data["data"] - for item in data: - yield item - - def _normalize(self, raw: Dict[str, Any]) -> QAExample: - qid = str(raw.get("id") or raw.get("qid") or raw.get("question_id") or "") - question = raw.get("question") or raw.get("query") or "" - answer = raw.get("answer") - if answer is None: - answer = raw.get("answers") or raw.get("final_answer") or "" - contexts = raw.get("contexts") or raw.get("supporting_facts") or None - return QAExample(id=qid, question=question, answer=answer, contexts=contexts) - - def iter(self, limit: Optional[int] = None) -> Iterable[QAExample]: - it = self._iter_jsonl() if self.is_jsonl else self._iter_json() - count = 0 - for raw in it: - ex = self._normalize(raw) - if not ex.question: - continue - yield ex - count += 1 - if limit is not None and count >= limit: - break + +"""Dataset loader for ACON's 8-objective QA benchmark (nq_multi_8). + +Adapted from ACON's experiments/smolagents/dataset.py. +Supports JSONL format with fields: id, question, answer. +""" +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional + + +@dataclass +class QAExample: + id: str + question: str + answer: Any # str or list[list[str]] — each sub-answer is a list of acceptable variants + contexts: Optional[List[str]] = None + + +class QALoader: + def __init__(self, data_path: str): + self.path = Path(data_path) + if not self.path.exists(): + raise FileNotFoundError(f"Data file not found: {self.path}") + self.is_jsonl = self.path.suffix.lower() in {".jsonl", ".jl"} + + def count(self, limit: Optional[int] = None) -> int: + if self.is_jsonl: + total = 0 + with self.path.open("r", encoding="utf-8") as f: + for line in f: + if line.strip(): + total += 1 + else: + data = json.loads(self.path.read_text(encoding="utf-8")) + if isinstance(data, dict) and "data" in data: + data = data["data"] + total = len(data) + + if limit is not None: + total = min(total, limit) + return total + + def _iter_jsonl(self) -> Iterable[Dict[str, Any]]: + with self.path.open("r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + yield json.loads(line) + + def _iter_json(self) -> Iterable[Dict[str, Any]]: + data = json.loads(self.path.read_text(encoding="utf-8")) + if isinstance(data, dict) and "data" in data: + data = data["data"] + for item in data: + yield item + + def _normalize(self, raw: Dict[str, Any]) -> QAExample: + qid = str(raw.get("id") or raw.get("qid") or raw.get("question_id") or "") + question = raw.get("question") or raw.get("query") or "" + answer = raw.get("answer") + if answer is None: + answer = raw.get("answers") or raw.get("final_answer") or "" + contexts = raw.get("contexts") or raw.get("supporting_facts") or None + return QAExample(id=qid, question=question, answer=answer, contexts=contexts) + + def iter(self, limit: Optional[int] = None) -> Iterable[QAExample]: + it = self._iter_jsonl() if self.is_jsonl else self._iter_json() + count = 0 + for raw in it: + ex = self._normalize(raw) + if not ex.question: + continue + yield ex + count += 1 + if limit is not None and count >= limit: + break diff --git a/sdk/benchmark/acon_eval/eval_utils.py b/sdk/benchmark/acon_eval/eval_utils.py index 89d44b90f..b62123c5e 100644 --- a/sdk/benchmark/acon_eval/eval_utils.py +++ b/sdk/benchmark/acon_eval/eval_utils.py @@ -1,76 +1,76 @@ - -"""ACON-style evaluation utilities: exact match and F1 scoring. - -Adapted from ACON's experiments/smolagents/eval_utils.py for use with -the nexent agent evaluation pipeline. -""" -import re -import string -from typing import Any - - -def _normalize_answer(s: str) -> str: - """SQuAD-style answer normalization with plural handling.""" - def lower(text: str) -> str: - return text.lower() - - def remove_punc(text: str) -> str: - return text.translate(str.maketrans('', '', string.punctuation)) - - def remove_articles(text: str) -> str: - return re.sub(r"\b(a|an|the)\b", " ", text) - - def white_space_fix(text: str) -> str: - return " ".join(text.split()) - - def normalize_plurals(text: str) -> str: - """Strip trailing 's' from words longer than 3 chars to unify singular/plural.""" - return " ".join( - word[:-1] if len(word) > 3 and word.endswith("s") and not word.endswith("ss") else word - for word in text.split() - ) - - return normalize_plurals(white_space_fix(remove_articles(remove_punc(lower(s))))) - - -def _f1_score(prediction: str, ground_truth: str) -> float: - pred_tokens = _normalize_answer(prediction).split() - gold_tokens = _normalize_answer(ground_truth).split() - if len(pred_tokens) == 0 and len(gold_tokens) == 0: - return 1.0 - if len(pred_tokens) == 0 or len(gold_tokens) == 0: - return 0.0 - common: dict[str, int] = {} - for t in pred_tokens: - common[t] = common.get(t, 0) + 1 - overlap = 0 - for t in gold_tokens: - if common.get(t, 0) > 0: - overlap += 1 - common[t] -= 1 - if overlap == 0: - return 0.0 - precision = overlap / len(pred_tokens) - recall = overlap / len(gold_tokens) - return 2 * precision * recall / (precision + recall) - - -def exact_match(pred: Any, gold: Any) -> bool: - """SQuAD-style normalized exact match.""" - def norm_one(x: Any) -> str: - if isinstance(x, (list, tuple)): - x = x[0] if x else "" - return _normalize_answer(str(x)) - - p = norm_one(pred) - if isinstance(gold, (list, tuple)): - return max(p == norm_one(g) for g in gold) - return p == norm_one(gold) - - -def f1_max(pred: Any, gold: Any) -> float: - """Max F1 over gold answer variants.""" - p = str(pred) if pred is not None else "" - if isinstance(gold, (list, tuple)): - return max((_f1_score(p, str(g)) for g in gold), default=0.0) - return _f1_score(p, str(gold)) + +"""ACON-style evaluation utilities: exact match and F1 scoring. + +Adapted from ACON's experiments/smolagents/eval_utils.py for use with +the nexent agent evaluation pipeline. +""" +import re +import string +from typing import Any + + +def _normalize_answer(s: str) -> str: + """SQuAD-style answer normalization with plural handling.""" + def lower(text: str) -> str: + return text.lower() + + def remove_punc(text: str) -> str: + return text.translate(str.maketrans('', '', string.punctuation)) + + def remove_articles(text: str) -> str: + return re.sub(r"\b(a|an|the)\b", " ", text) + + def white_space_fix(text: str) -> str: + return " ".join(text.split()) + + def normalize_plurals(text: str) -> str: + """Strip trailing 's' from words longer than 3 chars to unify singular/plural.""" + return " ".join( + word[:-1] if len(word) > 3 and word.endswith("s") and not word.endswith("ss") else word + for word in text.split() + ) + + return normalize_plurals(white_space_fix(remove_articles(remove_punc(lower(s))))) + + +def _f1_score(prediction: str, ground_truth: str) -> float: + pred_tokens = _normalize_answer(prediction).split() + gold_tokens = _normalize_answer(ground_truth).split() + if len(pred_tokens) == 0 and len(gold_tokens) == 0: + return 1.0 + if len(pred_tokens) == 0 or len(gold_tokens) == 0: + return 0.0 + common: dict[str, int] = {} + for t in pred_tokens: + common[t] = common.get(t, 0) + 1 + overlap = 0 + for t in gold_tokens: + if common.get(t, 0) > 0: + overlap += 1 + common[t] -= 1 + if overlap == 0: + return 0.0 + precision = overlap / len(pred_tokens) + recall = overlap / len(gold_tokens) + return 2 * precision * recall / (precision + recall) + + +def exact_match(pred: Any, gold: Any) -> bool: + """SQuAD-style normalized exact match.""" + def norm_one(x: Any) -> str: + if isinstance(x, (list, tuple)): + x = x[0] if x else "" + return _normalize_answer(str(x)) + + p = norm_one(pred) + if isinstance(gold, (list, tuple)): + return max(p == norm_one(g) for g in gold) + return p == norm_one(gold) + + +def f1_max(pred: Any, gold: Any) -> float: + """Max F1 over gold answer variants.""" + p = str(pred) if pred is not None else "" + if isinstance(gold, (list, tuple)): + return max((_f1_score(p, str(g)) for g in gold), default=0.0) + return _f1_score(p, str(gold)) diff --git a/sdk/benchmark/acon_eval/retriever_sesrver.py b/sdk/benchmark/acon_eval/retriever_sesrver.py index 2703c4981..32196f95a 100644 --- a/sdk/benchmark/acon_eval/retriever_sesrver.py +++ b/sdk/benchmark/acon_eval/retriever_sesrver.py @@ -1,423 +1,423 @@ -import json -import os -import warnings -from typing import List, Dict, Optional -import argparse - -try: - import faiss -except: - print("faiss not found, try to install it via `pip install faiss-cpu` or `pip install faiss-gpu`") -import torch -import numpy as np -from transformers import AutoConfig, AutoTokenizer, AutoModel -from tqdm import tqdm -import datasets - -import uvicorn -from fastapi import FastAPI -from pydantic import BaseModel - - -parser = argparse.ArgumentParser(description="Launch the local faiss retriever.") -parser.add_argument("--index_path", type=str, default="search/database/wikipedia/bm25", help="Corpus indexing file.") -parser.add_argument("--corpus_path", type=str, default="search/database/wikipedia/wiki-18.jsonl", help="Local corpus file.") -parser.add_argument("--topk", type=int, default=3, help="Number of retrieved passages for one query.") -parser.add_argument("--retriever_model", type=str, default="intfloat/e5-base-v2", help="Name of the retriever model.") - -args = parser.parse_args() - -def load_corpus(corpus_path: str): - corpus = datasets.load_dataset( - 'json', - data_files=corpus_path, - split="train", - num_proc=4 - ) - return corpus - -def read_jsonl(file_path): - data = [] - with open(file_path, "r") as f: - for line in f: - data.append(json.loads(line)) - return data - -def load_docs(corpus, doc_idxs): - results = [corpus[int(idx)] for idx in doc_idxs] - return results - -def load_model(model_path: str, use_fp16: bool = False): - model_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) - model = AutoModel.from_pretrained(model_path, trust_remote_code=True) - model.eval() - model.cuda() - if use_fp16: - model = model.half() - tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True, trust_remote_code=True) - return model, tokenizer - -def pooling( - pooler_output, - last_hidden_state, - attention_mask = None, - pooling_method = "mean" -): - if pooling_method == "mean": - last_hidden = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0) - return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None] - elif pooling_method == "cls": - return last_hidden_state[:, 0] - elif pooling_method == "pooler": - return pooler_output - else: - raise NotImplementedError("Pooling method not implemented!") - -class Encoder: - def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16): - self.model_name = model_name - self.model_path = model_path - self.pooling_method = pooling_method - self.max_length = max_length - self.use_fp16 = use_fp16 - - self.model, self.tokenizer = load_model(model_path=model_path, use_fp16=use_fp16) - self.model.eval() - - @torch.no_grad() - def encode(self, query_list: List[str], is_query=True) -> np.ndarray: - # processing query for different encoders - if isinstance(query_list, str): - query_list = [query_list] - - if "e5" in self.model_name.lower(): - if is_query: - query_list = [f"query: {query}" for query in query_list] - else: - query_list = [f"passage: {query}" for query in query_list] - - if "bge" in self.model_name.lower(): - if is_query: - query_list = [f"Represent this sentence for searching relevant passages: {query}" for query in query_list] - - inputs = self.tokenizer(query_list, - max_length=self.max_length, - padding=True, - truncation=True, - return_tensors="pt" - ) - inputs = {k: v.cuda() for k, v in inputs.items()} - - if "T5" in type(self.model).__name__: - # T5-based retrieval model - decoder_input_ids = torch.zeros( - (inputs['input_ids'].shape[0], 1), dtype=torch.long - ).to(inputs['input_ids'].device) - output = self.model( - **inputs, decoder_input_ids=decoder_input_ids, return_dict=True - ) - query_emb = output.last_hidden_state[:, 0, :] - else: - output = self.model(**inputs, return_dict=True) - query_emb = pooling(output.pooler_output, - output.last_hidden_state, - inputs['attention_mask'], - self.pooling_method) - if "dpr" not in self.model_name.lower(): - query_emb = torch.nn.functional.normalize(query_emb, dim=-1) - - query_emb = query_emb.detach().cpu().numpy() - query_emb = query_emb.astype(np.float32, order="C") - - del inputs, output - torch.cuda.empty_cache() - - return query_emb - -class BaseRetriever: - def __init__(self, config): - self.config = config - self.retrieval_method = config.retrieval_method - self.topk = config.retrieval_topk - - self.index_path = config.index_path - self.corpus_path = config.corpus_path - - def _search(self, query: str, num: int, return_score: bool): - raise NotImplementedError - - def _batch_search(self, query_list: List[str], num: int, return_score: bool): - raise NotImplementedError - - def search(self, query: str, num: int = None, return_score: bool = False): - return self._search(query, num, return_score) - - def batch_search(self, query_list: List[str], num: int = None, return_score: bool = False): - return self._batch_search(query_list, num, return_score) -class BM25Retriever(BaseRetriever): - def __init__(self, config): - super().__init__(config) - from pyserini.search.lucene import LuceneSearcher - - if not os.path.exists(self.index_path): - raise FileNotFoundError(f"BM25 index path not found: {self.index_path}") - - self.searcher = LuceneSearcher(self.index_path) - self.contain_doc = self._check_contain_doc() - - if not self.contain_doc: - if not os.path.exists(self.corpus_path): - raise FileNotFoundError(f"Corpus file not found: {self.corpus_path}") - self.corpus = load_corpus(self.corpus_path) - - self.max_process_num = 8 - - def _check_contain_doc(self): - try: - doc = self.searcher.doc("0") or self.searcher.doc(0) - return doc is not None and doc.raw() is not None - except Exception: - return False - - def _search(self, query: str, num: int = None, return_score: bool = False): - if not query or not query.strip(): - return ([], []) if return_score else [] - - num = num or self.topk - hits = self.searcher.search(query, num) - - if not hits: - return ([], []) if return_score else [] - - scores = [hit.score for hit in hits] - - if len(hits) < num: - warnings.warn(f"Only retrieved {len(hits)} documents, fewer than requested topk={num}") - - if self.contain_doc: - results = [] - for hit in hits: - try: - raw = self.searcher.doc(hit.docid).raw() - obj = json.loads(raw) - content = obj.get("contents", "") - - lines = content.split("\n") - title = lines[0].strip("\"") if lines else "" - text = "\n".join(lines[1:]) if len(lines) > 1 else content - - results.append({ - "title": title, - "text": text, - "contents": content - }) - except Exception as e: - results.append({ - "title": "", - "text": "", - "contents": "", - "error": f"Failed to parse docid={hit.docid}: {str(e)}" - }) - else: - results = load_docs(self.corpus, [hit.docid for hit in hits]) - - return (results, scores) if return_score else results - - def _batch_search(self, query_list: List[str], num: int = None, return_score: bool = False): - if isinstance(query_list, str): - query_list = [query_list] - - results = [] - scores = [] - - for query in query_list: - item_result, item_score = self._search(query, num, True) - results.append(item_result) - scores.append(item_score) - - return (results, scores) if return_score else results - -class DenseRetriever(BaseRetriever): - def __init__(self, config): - super().__init__(config) - self.index = faiss.read_index(self.index_path) - if config.faiss_gpu: - co = faiss.GpuMultipleClonerOptions() - co.useFloat16 = True - co.shard = True - self.index = faiss.index_cpu_to_all_gpus(self.index, co=co) - - self.corpus = load_corpus(self.corpus_path) - self.encoder = Encoder( - model_name = self.retrieval_method, - model_path = config.retrieval_model_path, - pooling_method = config.retrieval_pooling_method, - max_length = config.retrieval_query_max_length, - use_fp16 = config.retrieval_use_fp16 - ) - self.topk = config.retrieval_topk - self.batch_size = config.retrieval_batch_size - - def _search(self, query: str, num: int = None, return_score: bool = False): - if num is None: - num = self.topk - query_emb = self.encoder.encode(query) - scores, idxs = self.index.search(query_emb, k=num) - idxs = idxs[0] - scores = scores[0] - results = load_docs(self.corpus, idxs) - if return_score: - return results, scores.tolist() - else: - return results - - def _batch_search(self, query_list: List[str], num: int = None, return_score: bool = False): - if isinstance(query_list, str): - query_list = [query_list] - if num is None: - num = self.topk - - results = [] - scores = [] - for start_idx in tqdm(range(0, len(query_list), self.batch_size), desc='Retrieval process: '): - query_batch = query_list[start_idx:start_idx + self.batch_size] - batch_emb = self.encoder.encode(query_batch) - batch_scores, batch_idxs = self.index.search(batch_emb, k=num) - batch_scores = batch_scores.tolist() - batch_idxs = batch_idxs.tolist() - - # load_docs is not vectorized, but is a python list approach - flat_idxs = sum(batch_idxs, []) - batch_results = load_docs(self.corpus, flat_idxs) - # chunk them back - batch_results = [batch_results[i*num : (i+1)*num] for i in range(len(batch_idxs))] - - results.extend(batch_results) - scores.extend(batch_scores) - - del batch_emb, batch_scores, batch_idxs, query_batch, flat_idxs, batch_results - torch.cuda.empty_cache() - - if return_score: - return results, scores - else: - return results - -def get_retriever(config): - if config.retrieval_method == "bm25": - return BM25Retriever(config) - else: - return DenseRetriever(config) - - -##################################### -# FastAPI server below -##################################### - -class Config: - """ - Minimal config class (simulating your argparse) - Replace this with your real arguments or load them dynamically. - """ - def __init__( - self, - retrieval_method: str = "bm25", - retrieval_topk: int = 10, - index_path: str = "./index/bm25", - corpus_path: str = "./data/corpus.jsonl", - dataset_path: str = "./data", - data_split: str = "train", - faiss_gpu: bool = True, - retrieval_model_path: str = "./model", - retrieval_pooling_method: str = "mean", - retrieval_query_max_length: int = 256, - retrieval_use_fp16: bool = False, - retrieval_batch_size: int = 128 - ): - self.retrieval_method = retrieval_method - self.retrieval_topk = retrieval_topk - self.index_path = index_path - self.corpus_path = corpus_path - self.dataset_path = dataset_path - self.data_split = data_split - self.faiss_gpu = faiss_gpu - self.retrieval_model_path = retrieval_model_path - self.retrieval_pooling_method = retrieval_pooling_method - self.retrieval_query_max_length = retrieval_query_max_length - self.retrieval_use_fp16 = retrieval_use_fp16 - self.retrieval_batch_size = retrieval_batch_size - - -class QueryRequest(BaseModel): - queries: List[str] - topk: Optional[int] = None - return_scores: bool = False - - -app = FastAPI() - -# 1) Build a config (could also parse from arguments). -# In real usage, you'd parse your CLI arguments or environment variables. -config = Config( - retrieval_method="bm25", - index_path=args.index_path, - corpus_path=args.corpus_path, - retrieval_topk=args.topk, - faiss_gpu=False, -) - -# 2) Instantiate a global retriever so it is loaded once and reused. -retriever = get_retriever(config) -@app.post("/retrieve") -def retrieve_endpoint(request: QueryRequest): - """ - Input: - { - "queries": ["What is Python?"], - "topk": 3, - "return_scores": true - } - """ - if not request.queries: - return {"result": [], "error": "queries cannot be empty"} - - topk = request.topk or config.retrieval_topk - - try: - if request.return_scores: - results, scores = retriever.batch_search( - query_list=request.queries, - num=topk, - return_score=True - ) - - resp = [] - for single_result, single_scores in zip(results, scores): - combined = [] - for doc, score in zip(single_result, single_scores): - combined.append({ - "document": doc, - "score": score - }) - resp.append(combined) - - return {"result": resp} - - else: - results = retriever.batch_search( - query_list=request.queries, - num=topk, - return_score=False - ) - return {"result": results} - - except Exception as e: - return { - "result": [], - "error": str(e) - } - -if __name__ == "__main__": - # 3) Launch the server. By default, it listens on http://127.0.0.1:8000 - uvicorn.run(app, host="0.0.0.0", port=8005) +import json +import os +import warnings +from typing import List, Dict, Optional +import argparse + +try: + import faiss +except: + print("faiss not found, try to install it via `pip install faiss-cpu` or `pip install faiss-gpu`") +import torch +import numpy as np +from transformers import AutoConfig, AutoTokenizer, AutoModel +from tqdm import tqdm +import datasets + +import uvicorn +from fastapi import FastAPI +from pydantic import BaseModel + + +parser = argparse.ArgumentParser(description="Launch the local faiss retriever.") +parser.add_argument("--index_path", type=str, default="search/database/wikipedia/bm25", help="Corpus indexing file.") +parser.add_argument("--corpus_path", type=str, default="search/database/wikipedia/wiki-18.jsonl", help="Local corpus file.") +parser.add_argument("--topk", type=int, default=3, help="Number of retrieved passages for one query.") +parser.add_argument("--retriever_model", type=str, default="intfloat/e5-base-v2", help="Name of the retriever model.") + +args = parser.parse_args() + +def load_corpus(corpus_path: str): + corpus = datasets.load_dataset( + 'json', + data_files=corpus_path, + split="train", + num_proc=4 + ) + return corpus + +def read_jsonl(file_path): + data = [] + with open(file_path, "r") as f: + for line in f: + data.append(json.loads(line)) + return data + +def load_docs(corpus, doc_idxs): + results = [corpus[int(idx)] for idx in doc_idxs] + return results + +def load_model(model_path: str, use_fp16: bool = False): + model_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) + model = AutoModel.from_pretrained(model_path, trust_remote_code=True) + model.eval() + model.cuda() + if use_fp16: + model = model.half() + tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True, trust_remote_code=True) + return model, tokenizer + +def pooling( + pooler_output, + last_hidden_state, + attention_mask = None, + pooling_method = "mean" +): + if pooling_method == "mean": + last_hidden = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0) + return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None] + elif pooling_method == "cls": + return last_hidden_state[:, 0] + elif pooling_method == "pooler": + return pooler_output + else: + raise NotImplementedError("Pooling method not implemented!") + +class Encoder: + def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16): + self.model_name = model_name + self.model_path = model_path + self.pooling_method = pooling_method + self.max_length = max_length + self.use_fp16 = use_fp16 + + self.model, self.tokenizer = load_model(model_path=model_path, use_fp16=use_fp16) + self.model.eval() + + @torch.no_grad() + def encode(self, query_list: List[str], is_query=True) -> np.ndarray: + # processing query for different encoders + if isinstance(query_list, str): + query_list = [query_list] + + if "e5" in self.model_name.lower(): + if is_query: + query_list = [f"query: {query}" for query in query_list] + else: + query_list = [f"passage: {query}" for query in query_list] + + if "bge" in self.model_name.lower(): + if is_query: + query_list = [f"Represent this sentence for searching relevant passages: {query}" for query in query_list] + + inputs = self.tokenizer(query_list, + max_length=self.max_length, + padding=True, + truncation=True, + return_tensors="pt" + ) + inputs = {k: v.cuda() for k, v in inputs.items()} + + if "T5" in type(self.model).__name__: + # T5-based retrieval model + decoder_input_ids = torch.zeros( + (inputs['input_ids'].shape[0], 1), dtype=torch.long + ).to(inputs['input_ids'].device) + output = self.model( + **inputs, decoder_input_ids=decoder_input_ids, return_dict=True + ) + query_emb = output.last_hidden_state[:, 0, :] + else: + output = self.model(**inputs, return_dict=True) + query_emb = pooling(output.pooler_output, + output.last_hidden_state, + inputs['attention_mask'], + self.pooling_method) + if "dpr" not in self.model_name.lower(): + query_emb = torch.nn.functional.normalize(query_emb, dim=-1) + + query_emb = query_emb.detach().cpu().numpy() + query_emb = query_emb.astype(np.float32, order="C") + + del inputs, output + torch.cuda.empty_cache() + + return query_emb + +class BaseRetriever: + def __init__(self, config): + self.config = config + self.retrieval_method = config.retrieval_method + self.topk = config.retrieval_topk + + self.index_path = config.index_path + self.corpus_path = config.corpus_path + + def _search(self, query: str, num: int, return_score: bool): + raise NotImplementedError + + def _batch_search(self, query_list: List[str], num: int, return_score: bool): + raise NotImplementedError + + def search(self, query: str, num: int = None, return_score: bool = False): + return self._search(query, num, return_score) + + def batch_search(self, query_list: List[str], num: int = None, return_score: bool = False): + return self._batch_search(query_list, num, return_score) +class BM25Retriever(BaseRetriever): + def __init__(self, config): + super().__init__(config) + from pyserini.search.lucene import LuceneSearcher + + if not os.path.exists(self.index_path): + raise FileNotFoundError(f"BM25 index path not found: {self.index_path}") + + self.searcher = LuceneSearcher(self.index_path) + self.contain_doc = self._check_contain_doc() + + if not self.contain_doc: + if not os.path.exists(self.corpus_path): + raise FileNotFoundError(f"Corpus file not found: {self.corpus_path}") + self.corpus = load_corpus(self.corpus_path) + + self.max_process_num = 8 + + def _check_contain_doc(self): + try: + doc = self.searcher.doc("0") or self.searcher.doc(0) + return doc is not None and doc.raw() is not None + except Exception: + return False + + def _search(self, query: str, num: int = None, return_score: bool = False): + if not query or not query.strip(): + return ([], []) if return_score else [] + + num = num or self.topk + hits = self.searcher.search(query, num) + + if not hits: + return ([], []) if return_score else [] + + scores = [hit.score for hit in hits] + + if len(hits) < num: + warnings.warn(f"Only retrieved {len(hits)} documents, fewer than requested topk={num}") + + if self.contain_doc: + results = [] + for hit in hits: + try: + raw = self.searcher.doc(hit.docid).raw() + obj = json.loads(raw) + content = obj.get("contents", "") + + lines = content.split("\n") + title = lines[0].strip("\"") if lines else "" + text = "\n".join(lines[1:]) if len(lines) > 1 else content + + results.append({ + "title": title, + "text": text, + "contents": content + }) + except Exception as e: + results.append({ + "title": "", + "text": "", + "contents": "", + "error": f"Failed to parse docid={hit.docid}: {str(e)}" + }) + else: + results = load_docs(self.corpus, [hit.docid for hit in hits]) + + return (results, scores) if return_score else results + + def _batch_search(self, query_list: List[str], num: int = None, return_score: bool = False): + if isinstance(query_list, str): + query_list = [query_list] + + results = [] + scores = [] + + for query in query_list: + item_result, item_score = self._search(query, num, True) + results.append(item_result) + scores.append(item_score) + + return (results, scores) if return_score else results + +class DenseRetriever(BaseRetriever): + def __init__(self, config): + super().__init__(config) + self.index = faiss.read_index(self.index_path) + if config.faiss_gpu: + co = faiss.GpuMultipleClonerOptions() + co.useFloat16 = True + co.shard = True + self.index = faiss.index_cpu_to_all_gpus(self.index, co=co) + + self.corpus = load_corpus(self.corpus_path) + self.encoder = Encoder( + model_name = self.retrieval_method, + model_path = config.retrieval_model_path, + pooling_method = config.retrieval_pooling_method, + max_length = config.retrieval_query_max_length, + use_fp16 = config.retrieval_use_fp16 + ) + self.topk = config.retrieval_topk + self.batch_size = config.retrieval_batch_size + + def _search(self, query: str, num: int = None, return_score: bool = False): + if num is None: + num = self.topk + query_emb = self.encoder.encode(query) + scores, idxs = self.index.search(query_emb, k=num) + idxs = idxs[0] + scores = scores[0] + results = load_docs(self.corpus, idxs) + if return_score: + return results, scores.tolist() + else: + return results + + def _batch_search(self, query_list: List[str], num: int = None, return_score: bool = False): + if isinstance(query_list, str): + query_list = [query_list] + if num is None: + num = self.topk + + results = [] + scores = [] + for start_idx in tqdm(range(0, len(query_list), self.batch_size), desc='Retrieval process: '): + query_batch = query_list[start_idx:start_idx + self.batch_size] + batch_emb = self.encoder.encode(query_batch) + batch_scores, batch_idxs = self.index.search(batch_emb, k=num) + batch_scores = batch_scores.tolist() + batch_idxs = batch_idxs.tolist() + + # load_docs is not vectorized, but is a python list approach + flat_idxs = sum(batch_idxs, []) + batch_results = load_docs(self.corpus, flat_idxs) + # chunk them back + batch_results = [batch_results[i*num : (i+1)*num] for i in range(len(batch_idxs))] + + results.extend(batch_results) + scores.extend(batch_scores) + + del batch_emb, batch_scores, batch_idxs, query_batch, flat_idxs, batch_results + torch.cuda.empty_cache() + + if return_score: + return results, scores + else: + return results + +def get_retriever(config): + if config.retrieval_method == "bm25": + return BM25Retriever(config) + else: + return DenseRetriever(config) + + +##################################### +# FastAPI server below +##################################### + +class Config: + """ + Minimal config class (simulating your argparse) + Replace this with your real arguments or load them dynamically. + """ + def __init__( + self, + retrieval_method: str = "bm25", + retrieval_topk: int = 10, + index_path: str = "./index/bm25", + corpus_path: str = "./data/corpus.jsonl", + dataset_path: str = "./data", + data_split: str = "train", + faiss_gpu: bool = True, + retrieval_model_path: str = "./model", + retrieval_pooling_method: str = "mean", + retrieval_query_max_length: int = 256, + retrieval_use_fp16: bool = False, + retrieval_batch_size: int = 128 + ): + self.retrieval_method = retrieval_method + self.retrieval_topk = retrieval_topk + self.index_path = index_path + self.corpus_path = corpus_path + self.dataset_path = dataset_path + self.data_split = data_split + self.faiss_gpu = faiss_gpu + self.retrieval_model_path = retrieval_model_path + self.retrieval_pooling_method = retrieval_pooling_method + self.retrieval_query_max_length = retrieval_query_max_length + self.retrieval_use_fp16 = retrieval_use_fp16 + self.retrieval_batch_size = retrieval_batch_size + + +class QueryRequest(BaseModel): + queries: List[str] + topk: Optional[int] = None + return_scores: bool = False + + +app = FastAPI() + +# 1) Build a config (could also parse from arguments). +# In real usage, you'd parse your CLI arguments or environment variables. +config = Config( + retrieval_method="bm25", + index_path=args.index_path, + corpus_path=args.corpus_path, + retrieval_topk=args.topk, + faiss_gpu=False, +) + +# 2) Instantiate a global retriever so it is loaded once and reused. +retriever = get_retriever(config) +@app.post("/retrieve") +def retrieve_endpoint(request: QueryRequest): + """ + Input: + { + "queries": ["What is Python?"], + "topk": 3, + "return_scores": true + } + """ + if not request.queries: + return {"result": [], "error": "queries cannot be empty"} + + topk = request.topk or config.retrieval_topk + + try: + if request.return_scores: + results, scores = retriever.batch_search( + query_list=request.queries, + num=topk, + return_score=True + ) + + resp = [] + for single_result, single_scores in zip(results, scores): + combined = [] + for doc, score in zip(single_result, single_scores): + combined.append({ + "document": doc, + "score": score + }) + resp.append(combined) + + return {"result": resp} + + else: + results = retriever.batch_search( + query_list=request.queries, + num=topk, + return_score=False + ) + return {"result": results} + + except Exception as e: + return { + "result": [], + "error": str(e) + } + +if __name__ == "__main__": + # 3) Launch the server. By default, it listens on http://127.0.0.1:8000 + uvicorn.run(app, host="0.0.0.0", port=8005) \ No newline at end of file diff --git a/sdk/benchmark/acon_eval/run_acon_qa.py b/sdk/benchmark/acon_eval/run_acon_qa.py index e59771e01..6f496beda 100644 --- a/sdk/benchmark/acon_eval/run_acon_qa.py +++ b/sdk/benchmark/acon_eval/run_acon_qa.py @@ -1,570 +1,570 @@ -#!/usr/bin/env python3 -"""Run ACON multi-objective QA benchmark with nexent agent. - -Loads ACON's nq_multi_8 data, builds a nexent CoreAgent with -wikipedia_search + final_answer tools, evaluates with EM/F1 scoring. - -Supports three modes: - baseline — no context compression - context_manager — nexent's built-in ContextManager - -Use --num_objectives to control how many sub-questions per sample -(e.g. --num_objectives 2 to use only the first 2 sub-questions). - -Usage: - # Start ACON retriever server first: - # cd acon/experiments/smolagents/search && python retriever_server.py - # (or download the corpus and start it per ACON README) - - python run_acon_qa.py \ - --data_folder data/nq_multi_8 \ - --split test \ - --mode baseline \ - --num_objectives 4 \ - --limit 5 - -Results saved to outputs///summary.json + predictions.jsonl -""" -import argparse -import asyncio -import json -import os -import sys -import threading -from datetime import datetime -from typing import Optional - -# ---- Path setup ---- -# Robust path resolution via paths.py (.git discovery) — works regardless of file location -# 1. Add benchmark/ to sys.path so paths.py can be found -# 2. import paths triggers setup_paths() which adds sdk/, backend/ to sys.path -# 3. Add this directory for local module imports (dataset, eval_utils, tools) -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -import paths # noqa: F401 — side-effect: adds sdk/, backend/ to sys.path -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - -# ---- Register ACON tools into nexent namespace before any agent creation ---- -from tools import register_acon_tools, get_acon_tool_configs -register_acon_tools() - -from dataset import QALoader -from eval_utils import exact_match, f1_max - -from agent_runner import ( - build_agent_run_info_with_custom_prompt, - run_agent_with_tracking, - AgentRunResult, - ContextManagerConfig, -) - -from nexent.core.agents.agent_model import AgentHistory -from nexent.core.agents.agent_context import ContextManager - - -# ---- QA-specific system prompt builder ---- - -def build_qa_system_prompt(num_objectives: int) -> str: - answer_slots = "; ".join(f"answer{i}" for i in range(1, num_objectives + 1)) - - return f"""You are a multi-hop QA agent. The input contains multiple sub-questions separated by "; ". -Answer them sequentially by actually calling `wikipedia_search`, then call `final_answer`. - -# Tools -- `wikipedia_search(query: str, n_results: int = 3)` — searches the local 2018 Wikipedia retriever. -- `final_answer(answer: str)` — submits the final answer. - -# Mandatory Tool-Use Protocol -For every search, you must use a real code block: - - -result = wikipedia_search(query="...", n_results=3) -print(result) - - -Only an actual Observation produced after a `` block counts as evidence. -Do not write fake Search/Result text. - -# Core Rules -For each sub-question, in order: -1. Run one `wikipedia_search` call. -2. Read the actual Observation. -3. If the Observation clearly answers the sub-question, register the canonical answer and move to the next sub-question. -4. Do not run confirmation searches after finding a clear answer. -5. Use at most 3 searches per sub-question. -6. If the first 2 searches fail, the 3rd query must be broader and centered on the main entity/topic. -7. If 3 searches are exhausted, commit to the best candidate from observed results and move on. - -# Anti-Loop & Exhaustion Rules (CRITICAL — overriding priority) -- Track the exact count of wikipedia_search calls for the current sub-question. -- When count reaches 3, STOP searching immediately. Output ANSWER_Q: and move to the next question. No exceptions, no additional searches. -- If the last 2 searches returned completely irrelevant results (no mention of the target entity), the query angle is wrong. Do NOT search a third time with minor wording tweaks of the same query. Instead, search the main entity broadly (e.g. "Formula One history" instead of "chain F1"), or if already at 3, infer the best answer from any indirect clues in the observations and output ANSWER_Q. -- Self-check: if you catch yourself writing "I'm not finding it", "Perhaps", "Let me search for" or similar frustration phrases, you have already done enough searching. Output ANSWER_Q with your best inference immediately. -- After 3 searches, you already have your answer. Do NOT write "However", "But", "I'm not sure", "I'm not entirely sure", "Let me try one more", "Let me check directly", or any similar hesitation phrase. These words mean you have a candidate answer but are delaying. Output that candidate as ANSWER_Q right now and move on. Uncertainty is expected and acceptable — your best guess IS the answer. -- If the conversation contains a user message starting with "Summary of earlier steps in this task:", that message is an authoritative checkpoint of your progress. Before each search, check its JSON fields: "status", "search_counts", "pending_q", "next_action". If pending_q is empty and next_action says to call final_answer, call final_answer immediately — do not search again. If a question is marked "exhausted" in the summary, do not search it further. - -# Query Rules -- Prefer entity-focused queries, e.g. "Asha Bhosle Guinness", not "most prolific singer ever". -- Each query must be meaningfully different. -- Use `n_results=3` by default. - -# Answer Rules -- Use concise canonical answers: Wikipedia-title-like names or one-line factual answers. -- Keep modifiers only when needed for correctness. -- Do not include explanations, citations, dates, chapter/verse references, or extra context. -- Final answers must be separated by "; " in the original sub-question order. - -# Answer Registration — mandatory -Before moving from one question to the next, output exactly one plain-text marker: - -ANSWER_Q: - - -JUST Examples: -ANSWER_Q1: Eva Lund -ANSWER_Q2: September 1980 - -Rules: -- The marker is plain text, not a code block. -- If an Observation clearly answers Q, output `ANSWER_Q: `. -- After 3 searches, if there is any usable candidate in the Observations, output `ANSWER_Q: `. -- Never move to the next question without an ANSWER_Q marker for the current question. -- Use the registered ANSWER_Q markers to construct the final answer. - -# Final Answer -Before calling `final_answer`, count your answers. -The final answer must contain exactly one answer per sub-question. -Never submit a partial answer. - -Use a real code block: - - -final_answer(answer="{answer_slots}") - - -Start answering the real questions, starting with obtaining ANSWER_Q1. -""" - -def _sanitize_for_path(name: str) -> str: - return ''.join(ch if ch.isalnum() or ch in ('-', '_', '.') else '-' for ch in name) - - -async def run_sample( - ex, - max_steps: int, - retriever_port: str, - mode: str, - cm_config: Optional[ContextManagerConfig], - debug: bool, - system_prompt: str, -) -> dict: - """Run a single QA example through the nexent agent.""" - tools = get_acon_tool_configs(port=retriever_port) - - agent_run_info = build_agent_run_info_with_custom_prompt( - query=ex.question, - system_prompt=system_prompt, - history=[], - tools=tools, - max_steps=max_steps, - agent_name="acon_qa_agent", - agent_description="ACON multi-objective QA agent", - language="en", - context_manager_config=cm_config, - temperature=0 - ) - - # Attach shared ContextManager if mode is context_manager - shared_cm = None - if mode == "context_manager" and cm_config and cm_config.enabled: - shared_cm = ContextManager(config=cm_config, max_steps=max_steps) - agent_run_info.context_manager = shared_cm - - result = await run_agent_with_tracking(agent_run_info, debug=debug) - pred_raw = result.final_answer or "" - - # Score: split prediction by semicolons, compare to gold answer list - pred_list = [p.strip() for p in pred_raw.split(";")] - - # Pad or truncate predictions to match number of gold sub-answers - n_sub = len(ex.answer) - while len(pred_list) < n_sub: - pred_list.append("") - pred_list = pred_list[:n_sub] - - em_list = [exact_match(p, a) for p, a in zip(pred_list, ex.answer)] - f1_list = [f1_max(p, a) for p, a in zip(pred_list, ex.answer)] - - em_score = sum(em_list) / n_sub if n_sub else 0.0 - f1_score = sum(f1_list) / n_sub if n_sub else 0.0 - - return { - "pred_raw": pred_raw, - "pred_list": pred_list, - "em_score": em_score, - "f1_score": f1_score, - "em_list": em_list, - "f1_list": f1_list, - "step_count": result.step_count, - "errors": result.errors, - "total_input_tokens": result.total_input_tokens, - "total_output_tokens": result.total_output_tokens, - "cm_stats": shared_cm.get_all_compression_stats() if shared_cm else None, - "cm_token_counts": shared_cm.get_token_counts() if shared_cm else None, - } - - -async def main( - data_folder: str, - split: str, - mode: str, - max_steps: int, - limit: Optional[int], - retriever_port: str, - token_threshold: int, - keep_recent_pairs: int, - keep_recent_steps: int, - max_observation_length: int, - debug: bool, - output_dir: Optional[str], - id_list_file: Optional[str], - num_objectives: int, -): - # Resolve data path - split_key = (split or "test").lower() - if split_key in {"dev", "validation", "val"}: - split_key = "test" - fname = "train.jsonl" if split_key == "train" else "test.jsonl" - data_path = os.path.join(data_folder, fname) - - if not os.path.exists(data_path): - print(f"ERROR: Data file not found: {data_path}") - print(f" Make sure to point --data_folder to ACON's nq_multi_8 directory,") - print(f" e.g., D:/path/to/acon/experiments/smolagents/data/nq_multi_8") - return - - loader = QALoader(data_path) - - # Optional ID filtering - filter_ids = None - if id_list_file and os.path.exists(id_list_file): - with open(id_list_file, "r", encoding="utf-8") as f: - filter_ids = {line.strip() for line in f if line.strip() and not line.strip().startswith("#")} - - # Build iterator - if filter_ids is not None: - materialized = [ex for ex in loader.iter(limit=None) if ex.id in filter_ids] - if limit is not None: - materialized = materialized[:limit] - iterator = materialized - total_count = len(materialized) - else: - iterator = list(loader.iter(limit=limit)) - total_count = len(iterator) - - # Truncate sub-questions if num_objectives < 8 - if num_objectives < 8: - for ex in iterator: - q_parts = [q.strip() for q in ex.question.split(";")] - ex.question = "; ".join(q_parts[:num_objectives]) - ex.answer = ex.answer[:num_objectives] - - # Build QA-specific system prompt with dynamic answer slots - qa_system_prompt = build_qa_system_prompt(num_objectives) - - # ContextManager config based on mode - cm_config = None - if mode == "context_manager": - # Custom summary JSON schema that emphasizes task progress tracking - custom_summary_schema = { - "n_questions": "Total number of sub-questions.", - "answers": ( - "Ordered list of final-answer candidates. Length must equal n_questions. " - "Each item is either an exact canonical answer string or 'Unknown'. " - ), - "status": ( - "Array of length n_questions. Each item must be one of: " - "'unstarted', 'searching', 'answered', 'exhausted'. " - "answered requires a non-null answer other than 'Unknown'. or null" - "exhausted requires answer that need to be inferred." - ), - "search_counts": ( - "Array of integers of length n_questions. " - "Count only actual wikipedia_search calls." - ), - "current_q": ( - "The 1-based index of the next question to solve. " - "Usually the first index whose status is not 'answered' or 'exhausted'." - ), - "pending_q": ( - "List of question numbers whose status is 'unstarted' or 'searching'. " - "Do not include answered or exhausted questions." - ), - "next_action": ( - "One direct mechanical next step. Example: " - "'Run wikipedia_search for Q5: Ash Wednesday ashes palm leaves'." - ), - } - # Custom summary system prompt that emphasizes multi-question task tracking - custom_incremental_summary_system_prompt = ( - "Update the compact QA checkpoint based on the latest agent action. " - "Output only strict JSON matching the schema. No markdown.\n\n" - "Treat ANSWER_Q: ... marker as authoritative; never replace with null or Unknown." - "INCREMENTAL UPDATE RULES:\n" - "- Preserve all answered values; never downgrade them to null or 'Unknown'.\n" - "- If the latest action executed wikipedia_search, increment only that question's search_counts entry.\n" - "- If the latest observation clearly answers the current question, write the canonical answer into answers and set status to 'answered'.\n" - "- ENFORCEMENT: If any search_counts reaches >=3, its status MUST be 'exhausted' (NEVER 'searching'). " - "Set its answer to the best observed candidate, or 'Unknown' if nothing was useful. " - "An exhausted question must be REMOVED from pending_q.\n" - "- current_q must advance past any exhausted question to the next unstarted/searching question.\n" - "- If ALL questions are answered or exhausted, set next_action to 'Call final_answer with the collected answers'.\n" - "- NEVER set next_action to search a question whose search_counts is already >=3.\n" - "- Otherwise, leave answer as null and status as 'searching'.\n" - "- pending_q must contain exactly the question numbers with status 'unstarted' or 'searching'.\n" - "- Overwrite the old state completely. Do not append logs, snippets, or history." - ) - - custom_summary_system_prompt = ( - "You are creating a compact execution checkpoint for a sequential multi-question QA agent. " - "Output only strict JSON matching the schema. No markdown, greetings, or backticks.\n\n" - "Treat ANSWER_Q: ... marker as authoritative; never replace an ANSWER_Q value with null or Unknown.\n" - "STATE RULES:\n" - "- Preserve exact canonical answer strings when explicitly available.\n" - "- answers, status, and search_counts must all have length n_questions.\n" - "- status must be consistent with answers and search_counts:\n" - " * unstarted => answer is null, search_counts is 0\n" - " * searching => answer is null, search_counts is 1 or 2\n" - " * answered => answer is non-null canonical string, search_counts is 1-3\n" - " * exhausted => search_counts is >=3, answer is best inference or 'Unknown'\n" - "- A question with search_counts >=3 must have status 'exhausted', never 'searching'.\n" - "- pending_q must contain exactly the question numbers with status 'unstarted' or 'searching'.\n" - "- current_q should be the first question in pending_q.\n" - "- If all questions are answered or exhausted, set next_action to 'Call final_answer'.\n\n" - - "COMPACTION RULES:\n" - "- Strip raw search logs, snippets, long reasons, file status, and failed query history.\n" - "- Count every wikipedia_search call visible in the trajectory for each question.\n" - "- Keep the checkpoint short and stable. Do not append history." - ) - cm_config = ContextManagerConfig( - enabled=True, - token_threshold=token_threshold, - keep_recent_pairs=keep_recent_pairs, - keep_recent_steps=keep_recent_steps, - max_observation_length=max_observation_length, - summary_json_schema=custom_summary_schema, - summary_system_prompt=custom_summary_system_prompt, - incremental_summary_system_prompt=custom_incremental_summary_system_prompt, - ) - else: - # baseline: no compression - cm_config = ContextManagerConfig(enabled=False, token_threshold=10**9) - - # Output directory - if output_dir is None: - acon_eval_dir = os.path.dirname(os.path.abspath(__file__)) - outputs_root = os.path.join(acon_eval_dir, "outputs") - else: - outputs_root = output_dir - - mode_part = _sanitize_for_path(mode) - split_part = _sanitize_for_path(split_key) - out_dir = os.path.join(outputs_root, f"{mode_part}", split_part) - os.makedirs(out_dir, exist_ok=True) - - print(f"\n{'='*60}") - obj_label = f"{num_objectives}-Objective" if num_objectives != 8 else "8-Objective" - print(f"ACON {obj_label} QA Evaluation (nexent agent)") - print(f"{'='*60}") - print(f" Data: {data_path}") - print(f" Split: {split_key}") - print(f" Mode: {mode}") - print(f" Num objectives: {num_objectives}") - print(f" Max steps: {max_steps}") - print(f" Limit: {limit or 'all'}") - print(f" Total: {total_count}") - print(f" Retriever: 127.0.0.1:{retriever_port}") - if mode == "context_manager": - print(f" CM config: threshold={token_threshold}, keep_recent_pairs={keep_recent_pairs}, " - f"keep_recent_steps={keep_recent_steps}, max_obs_len={max_observation_length}") - print(f" Output: {out_dir}") - print(f"{'='*60}\n") - - n = 0 - em_sum = 0.0 - f1_sum = 0.0 - all_rows = [] - - for ex in iterator: - print(f"[{n+1}/{total_count}] {ex.id[:40]}...", end=" ", flush=True) - - try: - sample_result = await run_sample( - ex=ex, - max_steps=max_steps, - retriever_port=retriever_port, - mode=mode, - cm_config=cm_config, - debug=debug, - system_prompt=qa_system_prompt, - ) - em_score = sample_result["em_score"] - f1_score = sample_result["f1_score"] - print(f"EM={em_score:.2f} F1={f1_score:.2f} steps={sample_result['step_count']}") - except Exception as e: - print(f"ERROR: {e}") - em_score = 0.0 - f1_score = 0.0 - sample_result = { - "pred_raw": "", - "pred_list": [], - "em_score": 0.0, - "f1_score": 0.0, - "em_list": [], - "f1_list": [], - "step_count": 0, - "errors": [str(e)], - "total_input_tokens": 0, - "total_output_tokens": 0, - "cm_stats": None, - "cm_token_counts": None, - } - - em_sum += em_score - f1_sum += f1_score - n += 1 - - all_rows.append({ - "id": ex.id, - "question": ex.question, - "answer": ex.answer, - "prediction": sample_result["pred_list"], - "pred_raw": sample_result["pred_raw"], - "em": em_score, - "f1": f1_score, - "em_list": sample_result["em_list"], - "f1_list": sample_result["f1_list"], - "step_count": sample_result["step_count"], - "errors": sample_result["errors"], - "total_input_tokens": sample_result["total_input_tokens"], - "total_output_tokens": sample_result["total_output_tokens"], - "cm_stats": sample_result.get("cm_stats"), - "cm_token_counts": sample_result.get("cm_token_counts"), - }) - - # Token aggregates - total_input_tokens = sum(row["total_input_tokens"] for row in all_rows) - total_output_tokens = sum(row["total_output_tokens"] for row in all_rows) - avg_input_tokens = (total_input_tokens / n) if n else 0.0 - avg_output_tokens = (total_output_tokens / n) if n else 0.0 - - # Compression cost aggregate (context_manager mode only) - total_compression_input_tokens = 0 - total_compression_output_tokens = 0 - for row in all_rows: - cm_stats = row.get("cm_stats") - if cm_stats: - total_compression_input_tokens += cm_stats.get("total_input_tokens", 0) - total_compression_output_tokens += cm_stats.get("total_output_tokens", 0) - avg_compression_input_tokens = (total_compression_input_tokens / n) if n else 0.0 - avg_compression_output_tokens = (total_compression_output_tokens / n) if n else 0.0 - - # Summary - summary = { - "total": n, - "avg_em": (em_sum / n) if n else 0.0, - "avg_f1": (f1_sum / n) if n else 0.0, - "mode": mode, - "split": split_key, - "num_objectives": num_objectives, - "data_path": data_path, - "max_steps": max_steps, - "token_threshold": token_threshold if mode == "context_manager" else None, - "keep_recent_pairs": keep_recent_pairs if mode == "context_manager" else None, - "keep_recent_steps": keep_recent_steps if mode == "context_manager" else None, - "avg_input_tokens": avg_input_tokens, - "avg_output_tokens": avg_output_tokens, - "total_input_tokens": total_input_tokens, - "total_output_tokens": total_output_tokens, - "total_compression_input_tokens": total_compression_input_tokens if mode == "context_manager" else None, - "total_compression_output_tokens": total_compression_output_tokens if mode == "context_manager" else None, - "avg_compression_input_tokens": avg_compression_input_tokens if mode == "context_manager" else None, - "avg_compression_output_tokens": avg_compression_output_tokens if mode == "context_manager" else None, - "timestamp": datetime.now().isoformat(), - } - - # Save results - with open(os.path.join(out_dir, "summary.json"), "w", encoding="utf-8") as f: - json.dump(summary, f, indent=2) - - with open(os.path.join(out_dir, "predictions.jsonl"), "w", encoding="utf-8") as f: - for row in all_rows: - f.write(json.dumps(row, ensure_ascii=False) + "\n") - - print(f"\n{'='*60}") - print(f"Results Summary") - print(f"{'='*60}") - print(f" Mode: {mode}") - print(f" Total: {n}") - print(f" Avg EM: {em_sum/n*100:.1f}% ({em_sum:.2f}/{n})" if n else " Avg EM: N/A") - print(f" Avg F1: {f1_sum/n:.3f}" if n else " Avg F1: N/A") - print(f" Avg Input Tokens: {avg_input_tokens:,.0f}") - print(f" Avg Output Tokens: {avg_output_tokens:,.0f}") - if mode == "context_manager": - print(f" Avg Compression Input Tokens: {avg_compression_input_tokens:,.0f}") - print(f" Avg Compression Output Tokens: {avg_compression_output_tokens:,.0f}") - print(f" Output: {out_dir}") - print(f"{'='*60}\n") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Run ACON multi-objective QA benchmark with nexent agent") - parser.add_argument( - "--data_folder", - type=str, - default="data/nq_multi_8", - help="Path to ACON nq_multi_8 data folder (containing train.jsonl and test.jsonl)", - ) - parser.add_argument("--split", type=str, default="test", help="Dataset split: train or test") - parser.add_argument( - "--mode", - type=str, - default="baseline", - choices=["baseline", "context_manager"], - help="Evaluation mode: baseline (no compression) or context_manager (nexent CM)", - ) - parser.add_argument("--max_steps", type=int, default=30, help="Max agent steps per question") - parser.add_argument("--limit", type=int, default=None, help="Limit number of examples") - parser.add_argument("--retriever_port", type=str, default="8005", help="ACON retriever server port") - parser.add_argument("--token_threshold", type=int, default=7200, help="ContextManager token threshold (for context_manager mode)") - parser.add_argument("--keep_recent_pairs", type=int, default=1, help="ContextManager keep_recent_pairs (for context_manager mode)") - parser.add_argument("--keep_recent_steps", type=int, default=4, help="ContextManager keep_recent_steps (for context_manager mode)") - parser.add_argument("--max_observation_length", type=int, default=20000, help="Max observation length in chars (for context_manager mode)") - parser.add_argument("--debug", action="store_true", help="Enable debug output") - parser.add_argument("--output_dir", type=str, default=None, help="Override output directory") - parser.add_argument("--id_list_file", type=str, default=None, help="File with example IDs to filter (one per line)") - parser.add_argument( - "--num_objectives", - type=int, - default=8, - help="Number of sub-questions to use per sample (1-8, default: 8)", - ) - - args = parser.parse_args() - - asyncio.run(main( - data_folder=args.data_folder, - split=args.split, - mode=args.mode, - max_steps=args.max_steps, - limit=args.limit, - retriever_port=args.retriever_port, - token_threshold=args.token_threshold, - keep_recent_pairs=args.keep_recent_pairs, - keep_recent_steps=args.keep_recent_steps, - max_observation_length=args.max_observation_length, - debug=args.debug, - output_dir=args.output_dir, - id_list_file=args.id_list_file, - num_objectives=args.num_objectives, - )) +#!/usr/bin/env python3 +"""Run ACON multi-objective QA benchmark with nexent agent. + +Loads ACON's nq_multi_8 data, builds a nexent CoreAgent with +wikipedia_search + final_answer tools, evaluates with EM/F1 scoring. + +Supports three modes: + baseline — no context compression + context_manager — nexent's built-in ContextManager + +Use --num_objectives to control how many sub-questions per sample +(e.g. --num_objectives 2 to use only the first 2 sub-questions). + +Usage: + # Start ACON retriever server first: + # cd acon/experiments/smolagents/search && python retriever_server.py + # (or download the corpus and start it per ACON README) + + python run_acon_qa.py \ + --data_folder data/nq_multi_8 \ + --split test \ + --mode baseline \ + --num_objectives 4 \ + --limit 5 + +Results saved to outputs///summary.json + predictions.jsonl +""" +import argparse +import asyncio +import json +import os +import sys +import threading +from datetime import datetime +from typing import Optional + +# ---- Path setup ---- +# Robust path resolution via paths.py (.git discovery) — works regardless of file location +# 1. Add benchmark/ to sys.path so paths.py can be found +# 2. import paths triggers setup_paths() which adds sdk/, backend/ to sys.path +# 3. Add this directory for local module imports (dataset, eval_utils, tools) +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import paths # noqa: F401 — side-effect: adds sdk/, backend/ to sys.path +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +# ---- Register ACON tools into nexent namespace before any agent creation ---- +from tools import register_acon_tools, get_acon_tool_configs +register_acon_tools() + +from dataset import QALoader +from eval_utils import exact_match, f1_max + +from agent_runner import ( + build_agent_run_info_with_custom_prompt, + run_agent_with_tracking, + AgentRunResult, + ContextManagerConfig, +) + +from nexent.core.agents.agent_model import AgentHistory +from nexent.core.agents.agent_context import ContextManager + + +# ---- QA-specific system prompt builder ---- + +def build_qa_system_prompt(num_objectives: int) -> str: + answer_slots = "; ".join(f"answer{i}" for i in range(1, num_objectives + 1)) + + return f"""You are a multi-hop QA agent. The input contains multiple sub-questions separated by "; ". +Answer them sequentially by actually calling `wikipedia_search`, then call `final_answer`. + +# Tools +- `wikipedia_search(query: str, n_results: int = 3)` — searches the local 2018 Wikipedia retriever. +- `final_answer(answer: str)` — submits the final answer. + +# Mandatory Tool-Use Protocol +For every search, you must use a real code block: + + +result = wikipedia_search(query="...", n_results=3) +print(result) + + +Only an actual Observation produced after a `` block counts as evidence. +Do not write fake Search/Result text. + +# Core Rules +For each sub-question, in order: +1. Run one `wikipedia_search` call. +2. Read the actual Observation. +3. If the Observation clearly answers the sub-question, register the canonical answer and move to the next sub-question. +4. Do not run confirmation searches after finding a clear answer. +5. Use at most 3 searches per sub-question. +6. If the first 2 searches fail, the 3rd query must be broader and centered on the main entity/topic. +7. If 3 searches are exhausted, commit to the best candidate from observed results and move on. + +# Anti-Loop & Exhaustion Rules (CRITICAL — overriding priority) +- Track the exact count of wikipedia_search calls for the current sub-question. +- When count reaches 3, STOP searching immediately. Output ANSWER_Q: and move to the next question. No exceptions, no additional searches. +- If the last 2 searches returned completely irrelevant results (no mention of the target entity), the query angle is wrong. Do NOT search a third time with minor wording tweaks of the same query. Instead, search the main entity broadly (e.g. "Formula One history" instead of "chain F1"), or if already at 3, infer the best answer from any indirect clues in the observations and output ANSWER_Q. +- Self-check: if you catch yourself writing "I'm not finding it", "Perhaps", "Let me search for" or similar frustration phrases, you have already done enough searching. Output ANSWER_Q with your best inference immediately. +- After 3 searches, you already have your answer. Do NOT write "However", "But", "I'm not sure", "I'm not entirely sure", "Let me try one more", "Let me check directly", or any similar hesitation phrase. These words mean you have a candidate answer but are delaying. Output that candidate as ANSWER_Q right now and move on. Uncertainty is expected and acceptable — your best guess IS the answer. +- If the conversation contains a user message starting with "Summary of earlier steps in this task:", that message is an authoritative checkpoint of your progress. Before each search, check its JSON fields: "status", "search_counts", "pending_q", "next_action". If pending_q is empty and next_action says to call final_answer, call final_answer immediately — do not search again. If a question is marked "exhausted" in the summary, do not search it further. + +# Query Rules +- Prefer entity-focused queries, e.g. "Asha Bhosle Guinness", not "most prolific singer ever". +- Each query must be meaningfully different. +- Use `n_results=3` by default. + +# Answer Rules +- Use concise canonical answers: Wikipedia-title-like names or one-line factual answers. +- Keep modifiers only when needed for correctness. +- Do not include explanations, citations, dates, chapter/verse references, or extra context. +- Final answers must be separated by "; " in the original sub-question order. + +# Answer Registration — mandatory +Before moving from one question to the next, output exactly one plain-text marker: + +ANSWER_Q: + + +JUST Examples: +ANSWER_Q1: Eva Lund +ANSWER_Q2: September 1980 + +Rules: +- The marker is plain text, not a code block. +- If an Observation clearly answers Q, output `ANSWER_Q: `. +- After 3 searches, if there is any usable candidate in the Observations, output `ANSWER_Q: `. +- Never move to the next question without an ANSWER_Q marker for the current question. +- Use the registered ANSWER_Q markers to construct the final answer. + +# Final Answer +Before calling `final_answer`, count your answers. +The final answer must contain exactly one answer per sub-question. +Never submit a partial answer. + +Use a real code block: + + +final_answer(answer="{answer_slots}") + + +Start answering the real questions, starting with obtaining ANSWER_Q1. +""" + +def _sanitize_for_path(name: str) -> str: + return ''.join(ch if ch.isalnum() or ch in ('-', '_', '.') else '-' for ch in name) + + +async def run_sample( + ex, + max_steps: int, + retriever_port: str, + mode: str, + cm_config: Optional[ContextManagerConfig], + debug: bool, + system_prompt: str, +) -> dict: + """Run a single QA example through the nexent agent.""" + tools = get_acon_tool_configs(port=retriever_port) + + agent_run_info = build_agent_run_info_with_custom_prompt( + query=ex.question, + system_prompt=system_prompt, + history=[], + tools=tools, + max_steps=max_steps, + agent_name="acon_qa_agent", + agent_description="ACON multi-objective QA agent", + language="en", + context_manager_config=cm_config, + temperature=0 + ) + + # Attach shared ContextManager if mode is context_manager + shared_cm = None + if mode == "context_manager" and cm_config and cm_config.enabled: + shared_cm = ContextManager(config=cm_config, max_steps=max_steps) + agent_run_info.context_manager = shared_cm + + result = await run_agent_with_tracking(agent_run_info, debug=debug) + pred_raw = result.final_answer or "" + + # Score: split prediction by semicolons, compare to gold answer list + pred_list = [p.strip() for p in pred_raw.split(";")] + + # Pad or truncate predictions to match number of gold sub-answers + n_sub = len(ex.answer) + while len(pred_list) < n_sub: + pred_list.append("") + pred_list = pred_list[:n_sub] + + em_list = [exact_match(p, a) for p, a in zip(pred_list, ex.answer)] + f1_list = [f1_max(p, a) for p, a in zip(pred_list, ex.answer)] + + em_score = sum(em_list) / n_sub if n_sub else 0.0 + f1_score = sum(f1_list) / n_sub if n_sub else 0.0 + + return { + "pred_raw": pred_raw, + "pred_list": pred_list, + "em_score": em_score, + "f1_score": f1_score, + "em_list": em_list, + "f1_list": f1_list, + "step_count": result.step_count, + "errors": result.errors, + "total_input_tokens": result.total_input_tokens, + "total_output_tokens": result.total_output_tokens, + "cm_stats": shared_cm.get_all_compression_stats() if shared_cm else None, + "cm_token_counts": shared_cm.get_token_counts() if shared_cm else None, + } + + +async def main( + data_folder: str, + split: str, + mode: str, + max_steps: int, + limit: Optional[int], + retriever_port: str, + token_threshold: int, + keep_recent_pairs: int, + keep_recent_steps: int, + max_observation_length: int, + debug: bool, + output_dir: Optional[str], + id_list_file: Optional[str], + num_objectives: int, +): + # Resolve data path + split_key = (split or "test").lower() + if split_key in {"dev", "validation", "val"}: + split_key = "test" + fname = "train.jsonl" if split_key == "train" else "test.jsonl" + data_path = os.path.join(data_folder, fname) + + if not os.path.exists(data_path): + print(f"ERROR: Data file not found: {data_path}") + print(f" Make sure to point --data_folder to ACON's nq_multi_8 directory,") + print(f" e.g., D:/path/to/acon/experiments/smolagents/data/nq_multi_8") + return + + loader = QALoader(data_path) + + # Optional ID filtering + filter_ids = None + if id_list_file and os.path.exists(id_list_file): + with open(id_list_file, "r", encoding="utf-8") as f: + filter_ids = {line.strip() for line in f if line.strip() and not line.strip().startswith("#")} + + # Build iterator + if filter_ids is not None: + materialized = [ex for ex in loader.iter(limit=None) if ex.id in filter_ids] + if limit is not None: + materialized = materialized[:limit] + iterator = materialized + total_count = len(materialized) + else: + iterator = list(loader.iter(limit=limit)) + total_count = len(iterator) + + # Truncate sub-questions if num_objectives < 8 + if num_objectives < 8: + for ex in iterator: + q_parts = [q.strip() for q in ex.question.split(";")] + ex.question = "; ".join(q_parts[:num_objectives]) + ex.answer = ex.answer[:num_objectives] + + # Build QA-specific system prompt with dynamic answer slots + qa_system_prompt = build_qa_system_prompt(num_objectives) + + # ContextManager config based on mode + cm_config = None + if mode == "context_manager": + # Custom summary JSON schema that emphasizes task progress tracking + custom_summary_schema = { + "n_questions": "Total number of sub-questions.", + "answers": ( + "Ordered list of final-answer candidates. Length must equal n_questions. " + "Each item is either an exact canonical answer string or 'Unknown'. " + ), + "status": ( + "Array of length n_questions. Each item must be one of: " + "'unstarted', 'searching', 'answered', 'exhausted'. " + "answered requires a non-null answer other than 'Unknown'. or null" + "exhausted requires answer that need to be inferred." + ), + "search_counts": ( + "Array of integers of length n_questions. " + "Count only actual wikipedia_search calls." + ), + "current_q": ( + "The 1-based index of the next question to solve. " + "Usually the first index whose status is not 'answered' or 'exhausted'." + ), + "pending_q": ( + "List of question numbers whose status is 'unstarted' or 'searching'. " + "Do not include answered or exhausted questions." + ), + "next_action": ( + "One direct mechanical next step. Example: " + "'Run wikipedia_search for Q5: Ash Wednesday ashes palm leaves'." + ), + } + # Custom summary system prompt that emphasizes multi-question task tracking + custom_incremental_summary_system_prompt = ( + "Update the compact QA checkpoint based on the latest agent action. " + "Output only strict JSON matching the schema. No markdown.\n\n" + "Treat ANSWER_Q: ... marker as authoritative; never replace with null or Unknown." + "INCREMENTAL UPDATE RULES:\n" + "- Preserve all answered values; never downgrade them to null or 'Unknown'.\n" + "- If the latest action executed wikipedia_search, increment only that question's search_counts entry.\n" + "- If the latest observation clearly answers the current question, write the canonical answer into answers and set status to 'answered'.\n" + "- ENFORCEMENT: If any search_counts reaches >=3, its status MUST be 'exhausted' (NEVER 'searching'). " + "Set its answer to the best observed candidate, or 'Unknown' if nothing was useful. " + "An exhausted question must be REMOVED from pending_q.\n" + "- current_q must advance past any exhausted question to the next unstarted/searching question.\n" + "- If ALL questions are answered or exhausted, set next_action to 'Call final_answer with the collected answers'.\n" + "- NEVER set next_action to search a question whose search_counts is already >=3.\n" + "- Otherwise, leave answer as null and status as 'searching'.\n" + "- pending_q must contain exactly the question numbers with status 'unstarted' or 'searching'.\n" + "- Overwrite the old state completely. Do not append logs, snippets, or history." + ) + + custom_summary_system_prompt = ( + "You are creating a compact execution checkpoint for a sequential multi-question QA agent. " + "Output only strict JSON matching the schema. No markdown, greetings, or backticks.\n\n" + "Treat ANSWER_Q: ... marker as authoritative; never replace an ANSWER_Q value with null or Unknown.\n" + "STATE RULES:\n" + "- Preserve exact canonical answer strings when explicitly available.\n" + "- answers, status, and search_counts must all have length n_questions.\n" + "- status must be consistent with answers and search_counts:\n" + " * unstarted => answer is null, search_counts is 0\n" + " * searching => answer is null, search_counts is 1 or 2\n" + " * answered => answer is non-null canonical string, search_counts is 1-3\n" + " * exhausted => search_counts is >=3, answer is best inference or 'Unknown'\n" + "- A question with search_counts >=3 must have status 'exhausted', never 'searching'.\n" + "- pending_q must contain exactly the question numbers with status 'unstarted' or 'searching'.\n" + "- current_q should be the first question in pending_q.\n" + "- If all questions are answered or exhausted, set next_action to 'Call final_answer'.\n\n" + + "COMPACTION RULES:\n" + "- Strip raw search logs, snippets, long reasons, file status, and failed query history.\n" + "- Count every wikipedia_search call visible in the trajectory for each question.\n" + "- Keep the checkpoint short and stable. Do not append history." + ) + cm_config = ContextManagerConfig( + enabled=True, + token_threshold=token_threshold, + keep_recent_pairs=keep_recent_pairs, + keep_recent_steps=keep_recent_steps, + max_observation_length=max_observation_length, + summary_json_schema=custom_summary_schema, + summary_system_prompt=custom_summary_system_prompt, + incremental_summary_system_prompt=custom_incremental_summary_system_prompt, + ) + else: + # baseline: no compression + cm_config = ContextManagerConfig(enabled=False, token_threshold=10**9) + + # Output directory + if output_dir is None: + acon_eval_dir = os.path.dirname(os.path.abspath(__file__)) + outputs_root = os.path.join(acon_eval_dir, "outputs") + else: + outputs_root = output_dir + + mode_part = _sanitize_for_path(mode) + split_part = _sanitize_for_path(split_key) + out_dir = os.path.join(outputs_root, f"{mode_part}", split_part) + os.makedirs(out_dir, exist_ok=True) + + print(f"\n{'='*60}") + obj_label = f"{num_objectives}-Objective" if num_objectives != 8 else "8-Objective" + print(f"ACON {obj_label} QA Evaluation (nexent agent)") + print(f"{'='*60}") + print(f" Data: {data_path}") + print(f" Split: {split_key}") + print(f" Mode: {mode}") + print(f" Num objectives: {num_objectives}") + print(f" Max steps: {max_steps}") + print(f" Limit: {limit or 'all'}") + print(f" Total: {total_count}") + print(f" Retriever: 127.0.0.1:{retriever_port}") + if mode == "context_manager": + print(f" CM config: threshold={token_threshold}, keep_recent_pairs={keep_recent_pairs}, " + f"keep_recent_steps={keep_recent_steps}, max_obs_len={max_observation_length}") + print(f" Output: {out_dir}") + print(f"{'='*60}\n") + + n = 0 + em_sum = 0.0 + f1_sum = 0.0 + all_rows = [] + + for ex in iterator: + print(f"[{n+1}/{total_count}] {ex.id[:40]}...", end=" ", flush=True) + + try: + sample_result = await run_sample( + ex=ex, + max_steps=max_steps, + retriever_port=retriever_port, + mode=mode, + cm_config=cm_config, + debug=debug, + system_prompt=qa_system_prompt, + ) + em_score = sample_result["em_score"] + f1_score = sample_result["f1_score"] + print(f"EM={em_score:.2f} F1={f1_score:.2f} steps={sample_result['step_count']}") + except Exception as e: + print(f"ERROR: {e}") + em_score = 0.0 + f1_score = 0.0 + sample_result = { + "pred_raw": "", + "pred_list": [], + "em_score": 0.0, + "f1_score": 0.0, + "em_list": [], + "f1_list": [], + "step_count": 0, + "errors": [str(e)], + "total_input_tokens": 0, + "total_output_tokens": 0, + "cm_stats": None, + "cm_token_counts": None, + } + + em_sum += em_score + f1_sum += f1_score + n += 1 + + all_rows.append({ + "id": ex.id, + "question": ex.question, + "answer": ex.answer, + "prediction": sample_result["pred_list"], + "pred_raw": sample_result["pred_raw"], + "em": em_score, + "f1": f1_score, + "em_list": sample_result["em_list"], + "f1_list": sample_result["f1_list"], + "step_count": sample_result["step_count"], + "errors": sample_result["errors"], + "total_input_tokens": sample_result["total_input_tokens"], + "total_output_tokens": sample_result["total_output_tokens"], + "cm_stats": sample_result.get("cm_stats"), + "cm_token_counts": sample_result.get("cm_token_counts"), + }) + + # Token aggregates + total_input_tokens = sum(row["total_input_tokens"] for row in all_rows) + total_output_tokens = sum(row["total_output_tokens"] for row in all_rows) + avg_input_tokens = (total_input_tokens / n) if n else 0.0 + avg_output_tokens = (total_output_tokens / n) if n else 0.0 + + # Compression cost aggregate (context_manager mode only) + total_compression_input_tokens = 0 + total_compression_output_tokens = 0 + for row in all_rows: + cm_stats = row.get("cm_stats") + if cm_stats: + total_compression_input_tokens += cm_stats.get("total_input_tokens", 0) + total_compression_output_tokens += cm_stats.get("total_output_tokens", 0) + avg_compression_input_tokens = (total_compression_input_tokens / n) if n else 0.0 + avg_compression_output_tokens = (total_compression_output_tokens / n) if n else 0.0 + + # Summary + summary = { + "total": n, + "avg_em": (em_sum / n) if n else 0.0, + "avg_f1": (f1_sum / n) if n else 0.0, + "mode": mode, + "split": split_key, + "num_objectives": num_objectives, + "data_path": data_path, + "max_steps": max_steps, + "token_threshold": token_threshold if mode == "context_manager" else None, + "keep_recent_pairs": keep_recent_pairs if mode == "context_manager" else None, + "keep_recent_steps": keep_recent_steps if mode == "context_manager" else None, + "avg_input_tokens": avg_input_tokens, + "avg_output_tokens": avg_output_tokens, + "total_input_tokens": total_input_tokens, + "total_output_tokens": total_output_tokens, + "total_compression_input_tokens": total_compression_input_tokens if mode == "context_manager" else None, + "total_compression_output_tokens": total_compression_output_tokens if mode == "context_manager" else None, + "avg_compression_input_tokens": avg_compression_input_tokens if mode == "context_manager" else None, + "avg_compression_output_tokens": avg_compression_output_tokens if mode == "context_manager" else None, + "timestamp": datetime.now().isoformat(), + } + + # Save results + with open(os.path.join(out_dir, "summary.json"), "w", encoding="utf-8") as f: + json.dump(summary, f, indent=2) + + with open(os.path.join(out_dir, "predictions.jsonl"), "w", encoding="utf-8") as f: + for row in all_rows: + f.write(json.dumps(row, ensure_ascii=False) + "\n") + + print(f"\n{'='*60}") + print(f"Results Summary") + print(f"{'='*60}") + print(f" Mode: {mode}") + print(f" Total: {n}") + print(f" Avg EM: {em_sum/n*100:.1f}% ({em_sum:.2f}/{n})" if n else " Avg EM: N/A") + print(f" Avg F1: {f1_sum/n:.3f}" if n else " Avg F1: N/A") + print(f" Avg Input Tokens: {avg_input_tokens:,.0f}") + print(f" Avg Output Tokens: {avg_output_tokens:,.0f}") + if mode == "context_manager": + print(f" Avg Compression Input Tokens: {avg_compression_input_tokens:,.0f}") + print(f" Avg Compression Output Tokens: {avg_compression_output_tokens:,.0f}") + print(f" Output: {out_dir}") + print(f"{'='*60}\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Run ACON multi-objective QA benchmark with nexent agent") + parser.add_argument( + "--data_folder", + type=str, + default="data/nq_multi_8", + help="Path to ACON nq_multi_8 data folder (containing train.jsonl and test.jsonl)", + ) + parser.add_argument("--split", type=str, default="test", help="Dataset split: train or test") + parser.add_argument( + "--mode", + type=str, + default="baseline", + choices=["baseline", "context_manager"], + help="Evaluation mode: baseline (no compression) or context_manager (nexent CM)", + ) + parser.add_argument("--max_steps", type=int, default=30, help="Max agent steps per question") + parser.add_argument("--limit", type=int, default=None, help="Limit number of examples") + parser.add_argument("--retriever_port", type=str, default="8005", help="ACON retriever server port") + parser.add_argument("--token_threshold", type=int, default=7200, help="ContextManager token threshold (for context_manager mode)") + parser.add_argument("--keep_recent_pairs", type=int, default=1, help="ContextManager keep_recent_pairs (for context_manager mode)") + parser.add_argument("--keep_recent_steps", type=int, default=4, help="ContextManager keep_recent_steps (for context_manager mode)") + parser.add_argument("--max_observation_length", type=int, default=20000, help="Max observation length in chars (for context_manager mode)") + parser.add_argument("--debug", action="store_true", help="Enable debug output") + parser.add_argument("--output_dir", type=str, default=None, help="Override output directory") + parser.add_argument("--id_list_file", type=str, default=None, help="File with example IDs to filter (one per line)") + parser.add_argument( + "--num_objectives", + type=int, + default=8, + help="Number of sub-questions to use per sample (1-8, default: 8)", + ) + + args = parser.parse_args() + + asyncio.run(main( + data_folder=args.data_folder, + split=args.split, + mode=args.mode, + max_steps=args.max_steps, + limit=args.limit, + retriever_port=args.retriever_port, + token_threshold=args.token_threshold, + keep_recent_pairs=args.keep_recent_pairs, + keep_recent_steps=args.keep_recent_steps, + max_observation_length=args.max_observation_length, + debug=args.debug, + output_dir=args.output_dir, + id_list_file=args.id_list_file, + num_objectives=args.num_objectives, + )) diff --git a/sdk/benchmark/acon_eval/tools.py b/sdk/benchmark/acon_eval/tools.py index 828f05e63..680303269 100644 --- a/sdk/benchmark/acon_eval/tools.py +++ b/sdk/benchmark/acon_eval/tools.py @@ -1,131 +1,131 @@ -"""ACON QA benchmark tools for nexent agent. - -Provides WikipediaSearchTool and FinalAnswerTool as smolagents.Tool -subclasses, plus a helper to register them in nexent's tool namespace -so that NexentAgent.create_local_tool() can find them via globals(). -""" -from typing import Any - -import requests -from smolagents.tools import Tool - -from nexent.core.agents.agent_model import ToolConfig - - -class WikipediaSearchTool(Tool): - name = "wikipedia_search" - description = ( - "Uses semantic search to retrieve the parts of 2018 wikipedia " - "that could be most relevant to answer your query." - ) - inputs = { - "query": { - "type": "string", - "description": ( - "The query to perform. This should be semantically close to " - "your target documents. Use the affirmative form rather than " - "a question." - ), - }, - "n_results": { - "type": "integer", - "nullable": True, - "description": "The number of results to return. Minimum is 3. Maximum is 10.", - }, - } - output_type = "string" - - def __init__(self, port: str = "8005", **kwargs): - super().__init__() - self.port = port - self.url = f"http://127.0.0.1:{self.port}/retrieve" - - def forward(self, query: str, n_results: int = 3) -> str: - if n_results < 3: - n_results = 3 - if n_results > 10: - n_results = 10 - - assert isinstance(query, str), "Your search query must be a string" - payload = { - "queries": [query], - "topk": n_results, - "return_scores": True, - } - - response = requests.post(self.url, json=payload) - response.raise_for_status() - - retrieved_data = response.json() - docs = retrieved_data["result"][0] - - return "Retrieved documents:" + "".join( - f"\n\n[Document {str(i)}]\n" + doc["document"]["contents"] - for i, doc in enumerate(docs) - ) - - -class FinalAnswerTool(Tool): - name = "final_answer" - description = "Provides a final answer to the given problem." - inputs = { - "answer": { - "type": "any", - "description": "The final answer to the problem", - }, - } - output_type = "any" - - def forward(self, answer: Any) -> Any: - return answer - - -# --------------------------------------------------------------------------- -# Tool registration and ToolConfig builders -# --------------------------------------------------------------------------- - -def register_acon_tools(): - """Inject ACON tool classes into nexent.core.tools AND nexent_agent namespaces. - - NexentAgent.create_local_tool() looks up tool classes via globals(), - which is populated by `from ..tools import *` at import time. - Since `setattr` on the tools module does NOT update nexent_agent's - already-executed `globals()`, we must inject into BOTH modules. - """ - import nexent.core.tools as _tools_mod - import nexent.core.agents.nexent_agent as _agent_mod - for cls in (WikipediaSearchTool, FinalAnswerTool): - setattr(_tools_mod, cls.__name__, cls) - setattr(_agent_mod, cls.__name__, cls) - - -def build_wikipedia_search_tool_config(port: str = "8005") -> ToolConfig: - return ToolConfig( - class_name="WikipediaSearchTool", - name="wikipedia_search", - description=WikipediaSearchTool.description, - inputs=str(WikipediaSearchTool.inputs), - output_type=WikipediaSearchTool.output_type, - params={"port": port}, - source="local", - ) - - -def build_final_answer_tool_config() -> ToolConfig: - return ToolConfig( - class_name="FinalAnswerTool", - name="final_answer", - description=FinalAnswerTool.description, - inputs=str(FinalAnswerTool.inputs), - output_type=FinalAnswerTool.output_type, - params={}, - source="local", - ) - - -def get_acon_tool_configs(port: str = "8005") -> list[ToolConfig]: - """Return the standard ACON QA tool config list.""" - return [ - build_wikipedia_search_tool_config(port=port), - build_final_answer_tool_config(), +"""ACON QA benchmark tools for nexent agent. + +Provides WikipediaSearchTool and FinalAnswerTool as smolagents.Tool +subclasses, plus a helper to register them in nexent's tool namespace +so that NexentAgent.create_local_tool() can find them via globals(). +""" +from typing import Any + +import requests +from smolagents.tools import Tool + +from nexent.core.agents.agent_model import ToolConfig + + +class WikipediaSearchTool(Tool): + name = "wikipedia_search" + description = ( + "Uses semantic search to retrieve the parts of 2018 wikipedia " + "that could be most relevant to answer your query." + ) + inputs = { + "query": { + "type": "string", + "description": ( + "The query to perform. This should be semantically close to " + "your target documents. Use the affirmative form rather than " + "a question." + ), + }, + "n_results": { + "type": "integer", + "nullable": True, + "description": "The number of results to return. Minimum is 3. Maximum is 10.", + }, + } + output_type = "string" + + def __init__(self, port: str = "8005", **kwargs): + super().__init__() + self.port = port + self.url = f"http://127.0.0.1:{self.port}/retrieve" + + def forward(self, query: str, n_results: int = 3) -> str: + if n_results < 3: + n_results = 3 + if n_results > 10: + n_results = 10 + + assert isinstance(query, str), "Your search query must be a string" + payload = { + "queries": [query], + "topk": n_results, + "return_scores": True, + } + + response = requests.post(self.url, json=payload) + response.raise_for_status() + + retrieved_data = response.json() + docs = retrieved_data["result"][0] + + return "Retrieved documents:" + "".join( + f"\n\n[Document {str(i)}]\n" + doc["document"]["contents"] + for i, doc in enumerate(docs) + ) + + +class FinalAnswerTool(Tool): + name = "final_answer" + description = "Provides a final answer to the given problem." + inputs = { + "answer": { + "type": "any", + "description": "The final answer to the problem", + }, + } + output_type = "any" + + def forward(self, answer: Any) -> Any: + return answer + + +# --------------------------------------------------------------------------- +# Tool registration and ToolConfig builders +# --------------------------------------------------------------------------- + +def register_acon_tools(): + """Inject ACON tool classes into nexent.core.tools AND nexent_agent namespaces. + + NexentAgent.create_local_tool() looks up tool classes via globals(), + which is populated by `from ..tools import *` at import time. + Since `setattr` on the tools module does NOT update nexent_agent's + already-executed `globals()`, we must inject into BOTH modules. + """ + import nexent.core.tools as _tools_mod + import nexent.core.agents.nexent_agent as _agent_mod + for cls in (WikipediaSearchTool, FinalAnswerTool): + setattr(_tools_mod, cls.__name__, cls) + setattr(_agent_mod, cls.__name__, cls) + + +def build_wikipedia_search_tool_config(port: str = "8005") -> ToolConfig: + return ToolConfig( + class_name="WikipediaSearchTool", + name="wikipedia_search", + description=WikipediaSearchTool.description, + inputs=str(WikipediaSearchTool.inputs), + output_type=WikipediaSearchTool.output_type, + params={"port": port}, + source="local", + ) + + +def build_final_answer_tool_config() -> ToolConfig: + return ToolConfig( + class_name="FinalAnswerTool", + name="final_answer", + description=FinalAnswerTool.description, + inputs=str(FinalAnswerTool.inputs), + output_type=FinalAnswerTool.output_type, + params={}, + source="local", + ) + + +def get_acon_tool_configs(port: str = "8005") -> list[ToolConfig]: + """Return the standard ACON QA tool config list.""" + return [ + build_wikipedia_search_tool_config(port=port), + build_final_answer_tool_config(), ] \ No newline at end of file diff --git a/sdk/benchmark/manual_cases/calibrate_thresholds.py b/sdk/benchmark/manual_cases/calibrate_thresholds.py index 701372c6e..2458a7e09 100644 --- a/sdk/benchmark/manual_cases/calibrate_thresholds.py +++ b/sdk/benchmark/manual_cases/calibrate_thresholds.py @@ -1,148 +1,148 @@ -"""Calibrate token_threshold in case.json files based on actual history token counts. - -For each case under ./cases/, computes the token count of history.json PLUS -the system prompt tokens (using the same estimate_tokens_text() used at runtime), -then writes that value into case.json's compressed_config.token_threshold so -compression triggers precisely when the full context reaches this size. - -The threshold must account for system_prompt + history, because the ContextManager -checks token count against the full message list (which includes system prompt). - -Uses the same BENCHMARK_SYSTEM_PROMPT as test_benchmark.py for consistency. - -Usage: - python calibrate_thresholds.py [--cases-root ./cases] [--system-prompt ] [--dry-run] -""" - -import glob -import json -import os -import sys - -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) -import paths # noqa: F401 — side-effect: adds sdk/ to sys.path - -from nexent.core.utils.token_estimation import estimate_tokens_text - -# Same lean benchmark prompt as test_benchmark.py — kept in sync. -BENCHMARK_SYSTEM_PROMPT = """You are a helpful assistant. Answer the user's questions based on the conversation history and your knowledge. - -- Be precise and concise. -- When the answer depends on information from earlier conversation, refer to it accurately. -- Do not fabricate information you do not know. -- Use final_answer to submit your response. - -Now start!""" - - -def calibrate_thresholds( - cases_root: str = "./cases", - system_prompt: str = None, - dry_run: bool = False, -) -> list[dict]: - """Calibrate token_threshold in every case.json under cases_root. - - token_threshold = system_prompt_tokens + history_tokens - - Args: - cases_root: Directory containing case subdirectories. - system_prompt: System prompt string. Defaults to BENCHMARK_SYSTEM_PROMPT - (matching test_benchmark.py runtime). - dry_run: If True, compute but do not write files. - - Returns: - List of dicts with calibration details for each case. - """ - sp = system_prompt if system_prompt is not None else BENCHMARK_SYSTEM_PROMPT - sp_tokens = estimate_tokens_text(sp) - - results = [] - case_paths = sorted(glob.glob(os.path.join(cases_root, "*/case.json"))) - - if not case_paths: - print(f"No cases found under {cases_root}") - return results - - print(f"System prompt tokens: {sp_tokens}") - - for case_path in case_paths: - case_dir = os.path.dirname(case_path) - case_name = os.path.basename(case_dir) - - with open(case_path, "r", encoding="utf-8") as f: - case = json.load(f) - - history_relpath = case.get("history_file", "history.json") - history_abspath = os.path.join(case_dir, history_relpath) - - if not os.path.exists(history_abspath): - print(f" SKIP {case_name}: {history_relpath} not found") - continue - - with open(history_abspath, "r", encoding="utf-8") as f: - history = json.load(f) - - history_text = "".join(msg["content"] for msg in history) - history_tokens = estimate_tokens_text(history_text) - total_tokens = sp_tokens + history_tokens - - old_threshold = case.get("compressed_config", {}).get("token_threshold") - changed = old_threshold != total_tokens - - results.append({ - "case": case_name, - "old_threshold": old_threshold, - "new_threshold": total_tokens, - "history_tokens": history_tokens, - "system_prompt_tokens": sp_tokens, - "changed": changed, - }) - - if changed: - case.setdefault("compressed_config", {})["token_threshold"] = total_tokens - if not dry_run: - with open(case_path, "w", encoding="utf-8") as f: - json.dump(case, f, ensure_ascii=False, indent=2) - print( - f" {case_name}: token_threshold {old_threshold} -> {total_tokens} " - f"(sp={sp_tokens} + history={history_tokens})" - ) - else: - print( - f" {case_name}: token_threshold {old_threshold} -> {total_tokens} " - f"(sp={sp_tokens} + history={history_tokens}) [dry-run]" - ) - else: - print(f" {case_name}: token_threshold already {total_tokens}, no change") - - changed_count = sum(1 for r in results if r["changed"]) - if dry_run: - print(f"\nDry-run: {changed_count} case(s) would be calibrated (no files written).") - else: - print(f"\nCalibrated {changed_count} case(s).") - - return results - - -if __name__ == "__main__": - import argparse - - parser = argparse.ArgumentParser(description="Calibrate token_threshold in case.json files") - parser.add_argument( - "--cases-root", default="./cases", - help="Root directory containing case subdirectories (default: ./cases)", - ) - parser.add_argument( - "--system-prompt", default=None, - help="Custom system prompt string (default: build from agent_runner template)", - ) - parser.add_argument( - "--dry-run", action="store_true", - help="Compute new thresholds but do not write to case.json files", - ) - args = parser.parse_args() - calibrate_thresholds( - cases_root=args.cases_root, - system_prompt=args.system_prompt, - dry_run=args.dry_run, +"""Calibrate token_threshold in case.json files based on actual history token counts. + +For each case under ./cases/, computes the token count of history.json PLUS +the system prompt tokens (using the same estimate_tokens_text() used at runtime), +then writes that value into case.json's compressed_config.token_threshold so +compression triggers precisely when the full context reaches this size. + +The threshold must account for system_prompt + history, because the ContextManager +checks token count against the full message list (which includes system prompt). + +Uses the same BENCHMARK_SYSTEM_PROMPT as test_benchmark.py for consistency. + +Usage: + python calibrate_thresholds.py [--cases-root ./cases] [--system-prompt ] [--dry-run] +""" + +import glob +import json +import os +import sys + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) +import paths # noqa: F401 — side-effect: adds sdk/ to sys.path + +from nexent.core.utils.token_estimation import estimate_tokens_text + +# Same lean benchmark prompt as test_benchmark.py — kept in sync. +BENCHMARK_SYSTEM_PROMPT = """You are a helpful assistant. Answer the user's questions based on the conversation history and your knowledge. + +- Be precise and concise. +- When the answer depends on information from earlier conversation, refer to it accurately. +- Do not fabricate information you do not know. +- Use final_answer to submit your response. + +Now start!""" + + +def calibrate_thresholds( + cases_root: str = "./cases", + system_prompt: str = None, + dry_run: bool = False, +) -> list[dict]: + """Calibrate token_threshold in every case.json under cases_root. + + token_threshold = system_prompt_tokens + history_tokens + + Args: + cases_root: Directory containing case subdirectories. + system_prompt: System prompt string. Defaults to BENCHMARK_SYSTEM_PROMPT + (matching test_benchmark.py runtime). + dry_run: If True, compute but do not write files. + + Returns: + List of dicts with calibration details for each case. + """ + sp = system_prompt if system_prompt is not None else BENCHMARK_SYSTEM_PROMPT + sp_tokens = estimate_tokens_text(sp) + + results = [] + case_paths = sorted(glob.glob(os.path.join(cases_root, "*/case.json"))) + + if not case_paths: + print(f"No cases found under {cases_root}") + return results + + print(f"System prompt tokens: {sp_tokens}") + + for case_path in case_paths: + case_dir = os.path.dirname(case_path) + case_name = os.path.basename(case_dir) + + with open(case_path, "r", encoding="utf-8") as f: + case = json.load(f) + + history_relpath = case.get("history_file", "history.json") + history_abspath = os.path.join(case_dir, history_relpath) + + if not os.path.exists(history_abspath): + print(f" SKIP {case_name}: {history_relpath} not found") + continue + + with open(history_abspath, "r", encoding="utf-8") as f: + history = json.load(f) + + history_text = "".join(msg["content"] for msg in history) + history_tokens = estimate_tokens_text(history_text) + total_tokens = sp_tokens + history_tokens + + old_threshold = case.get("compressed_config", {}).get("token_threshold") + changed = old_threshold != total_tokens + + results.append({ + "case": case_name, + "old_threshold": old_threshold, + "new_threshold": total_tokens, + "history_tokens": history_tokens, + "system_prompt_tokens": sp_tokens, + "changed": changed, + }) + + if changed: + case.setdefault("compressed_config", {})["token_threshold"] = total_tokens + if not dry_run: + with open(case_path, "w", encoding="utf-8") as f: + json.dump(case, f, ensure_ascii=False, indent=2) + print( + f" {case_name}: token_threshold {old_threshold} -> {total_tokens} " + f"(sp={sp_tokens} + history={history_tokens})" + ) + else: + print( + f" {case_name}: token_threshold {old_threshold} -> {total_tokens} " + f"(sp={sp_tokens} + history={history_tokens}) [dry-run]" + ) + else: + print(f" {case_name}: token_threshold already {total_tokens}, no change") + + changed_count = sum(1 for r in results if r["changed"]) + if dry_run: + print(f"\nDry-run: {changed_count} case(s) would be calibrated (no files written).") + else: + print(f"\nCalibrated {changed_count} case(s).") + + return results + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Calibrate token_threshold in case.json files") + parser.add_argument( + "--cases-root", default="./cases", + help="Root directory containing case subdirectories (default: ./cases)", + ) + parser.add_argument( + "--system-prompt", default=None, + help="Custom system prompt string (default: build from agent_runner template)", + ) + parser.add_argument( + "--dry-run", action="store_true", + help="Compute new thresholds but do not write to case.json files", + ) + args = parser.parse_args() + calibrate_thresholds( + cases_root=args.cases_root, + system_prompt=args.system_prompt, + dry_run=args.dry_run, ) \ No newline at end of file diff --git a/sdk/benchmark/manual_cases/eval_utils.py b/sdk/benchmark/manual_cases/eval_utils.py index 36dfde485..2f074df55 100644 --- a/sdk/benchmark/manual_cases/eval_utils.py +++ b/sdk/benchmark/manual_cases/eval_utils.py @@ -1,77 +1,77 @@ -from dataclasses import dataclass - - -@dataclass -class EvalResult: - passed: bool - score: float - details: dict - - -def contains_all(text: str, keywords: list[str]) -> bool: - text = text.lower() - return all(k.lower() in text for k in keywords) - - -def contains_any(text: str, keywords: list[str]) -> bool: - text = text.lower() - return any(k.lower() in text for k in keywords) - - -def count_matches(text: str, keywords: list[str]) -> int: - """Count how many keywords are present in the text (case-insensitive).""" - text = text.lower() - return sum(1 for k in keywords if k.lower() in text) - - -def eval_text(text: str, check: dict) -> EvalResult: - """Evaluate text against keyword checks with partial scoring. - - Scoring rules: - - must_contain: score = matched_count / total_keywords - (1.0 if all present, 0.6 if 3/5 present, etc.) - - must_contain_any: score = 1.0 if any present, 0.0 otherwise - - When both are present, score is the average of both sub-scores. - - passed is True only when all checks fully pass (backward compatible). - """ - passed = True - details = {} - scores = [] - - if "must_contain" in check: - keywords = check["must_contain"] - matched = count_matches(text, keywords) - ok = matched == len(keywords) - details["must_contain"] = { - "matched": matched, - "total": len(keywords), - "ok": ok, - } - scores.append(matched / len(keywords) if keywords else 1.0) - passed = passed and ok - - if "must_contain_any" in check: - keywords = check["must_contain_any"] - ok = contains_any(text, keywords) - matched = count_matches(text, keywords) - details["must_contain_any"] = { - "matched": matched, - "total": len(keywords), - "ok": ok, - } - scores.append(1.0 if ok else 0.0) - passed = passed and ok - - score = sum(scores) / len(scores) if scores else (1.0 if passed else 0.0) - - return EvalResult( - passed=passed, - score=score, - details=details, - ) - - -def average_score(results: list[EvalResult]) -> float: - if not results: - return 0.0 +from dataclasses import dataclass + + +@dataclass +class EvalResult: + passed: bool + score: float + details: dict + + +def contains_all(text: str, keywords: list[str]) -> bool: + text = text.lower() + return all(k.lower() in text for k in keywords) + + +def contains_any(text: str, keywords: list[str]) -> bool: + text = text.lower() + return any(k.lower() in text for k in keywords) + + +def count_matches(text: str, keywords: list[str]) -> int: + """Count how many keywords are present in the text (case-insensitive).""" + text = text.lower() + return sum(1 for k in keywords if k.lower() in text) + + +def eval_text(text: str, check: dict) -> EvalResult: + """Evaluate text against keyword checks with partial scoring. + + Scoring rules: + - must_contain: score = matched_count / total_keywords + (1.0 if all present, 0.6 if 3/5 present, etc.) + - must_contain_any: score = 1.0 if any present, 0.0 otherwise + - When both are present, score is the average of both sub-scores. + - passed is True only when all checks fully pass (backward compatible). + """ + passed = True + details = {} + scores = [] + + if "must_contain" in check: + keywords = check["must_contain"] + matched = count_matches(text, keywords) + ok = matched == len(keywords) + details["must_contain"] = { + "matched": matched, + "total": len(keywords), + "ok": ok, + } + scores.append(matched / len(keywords) if keywords else 1.0) + passed = passed and ok + + if "must_contain_any" in check: + keywords = check["must_contain_any"] + ok = contains_any(text, keywords) + matched = count_matches(text, keywords) + details["must_contain_any"] = { + "matched": matched, + "total": len(keywords), + "ok": ok, + } + scores.append(1.0 if ok else 0.0) + passed = passed and ok + + score = sum(scores) / len(scores) if scores else (1.0 if passed else 0.0) + + return EvalResult( + passed=passed, + score=score, + details=details, + ) + + +def average_score(results: list[EvalResult]) -> float: + if not results: + return 0.0 return sum(r.score for r in results) / len(results) \ No newline at end of file diff --git a/sdk/benchmark/manual_cases/summary_inspector.py b/sdk/benchmark/manual_cases/summary_inspector.py index 4dc459af9..248c4b288 100644 --- a/sdk/benchmark/manual_cases/summary_inspector.py +++ b/sdk/benchmark/manual_cases/summary_inspector.py @@ -1,330 +1,330 @@ -# -*- coding: utf-8 -*- -""" -Standalone Summary Inspector — quick evaluation of compression prompt/schema quality. - -Completely independent from test_benchmark.py and the cases/ directory. -Uses compress_history_offline to compress history and checks whether the -resulting summary retains key information. No agent runs needed — just -one LLM call per inspection + text-based checks. - -Use case: - - Iterate on summary prompt / schema in summary_config.py - - Verify that key facts survive compression without running full agent loops - - Compare different ContextManagerConfig settings side-by-side - -Directory layout (independent from cases/ and reports/): - - inspections/ - └── / - ├── history.json # [{"role": "user|assistant", "content": "..."}] - └── checks.json # [{"description": "...", "must_contain": [...]}] - -Result is written to inspections//_result.json (co-located with input). - -Usage: - python summary_inspector.py # all inspections - python summary_inspector.py -n example_infra # single inspection - python summary_inspector.py --config my_config.json # custom config overrides - python summary_inspector.py --save-summary # also save raw summary .txt -""" - -import argparse -import json -import os -import sys -import glob - -# ============ Path Setup ============ -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -import paths # noqa: F401 — side-effect: adds sdk/, backend/ to sys.path - -from dotenv import load_dotenv -load_dotenv() - -from nexent.core.agents.agent_context import compress_history_offline, ContextManagerConfig -from nexent.core.agents.agent_model import ModelConfig -from nexent.core.models.openai_llm import OpenAIModel - -from eval_utils import eval_text - - -# ============ Config ============ -LLM_API_KEY = os.getenv("LLM_API_KEY") -LLM_MODEL_NAME = os.getenv("LLM_MODEL_NAME") -LLM_API_URL = os.getenv("LLM_API_URL") - -INSPECTIONS_DIR = "./inspections" - - -def create_model(temperature: float = 0.1): - """Create an LLM model for offline compression.""" - from nexent.core.utils.observer import MessageObserver - - model_config = ModelConfig( - cite_name="inspector_model", - api_key=LLM_API_KEY, - model_name=LLM_MODEL_NAME, - url=LLM_API_URL, - temperature=temperature, - ssl_verify=False, - ) - return OpenAIModel( - observer=MessageObserver(), - model_id=model_config.model_name, - api_key=model_config.api_key, - api_base=model_config.url, - temperature=model_config.temperature, - top_p=model_config.top_p, - ssl_verify=model_config.ssl_verify, - ) - - -def history_to_pairs(history: list) -> list[tuple[str, str]]: - """Convert [{role, content}] to [(user_text, assistant_text)] pairs. - - Consecutive user messages are merged; same for assistant messages, - so the output is a clean alternating sequence of pairs. - """ - pairs = [] - current_user = [] - current_assistant = [] - - for entry in history: - role = entry["role"] - content = entry["content"] - if role == "user": - if current_assistant: - pairs.append(( - "\n".join(current_user).strip(), - "\n".join(current_assistant).strip(), - )) - current_user = [] - current_assistant = [] - current_user.append(content) - elif role == "assistant": - current_assistant.append(content) - - if current_user and current_assistant: - pairs.append(( - "\n".join(current_user).strip(), - "\n".join(current_assistant).strip(), - )) - - return pairs - - -def build_config(overrides: dict = None) -> ContextManagerConfig: - """Build ContextManagerConfig with optional field overrides.""" - config = ContextManagerConfig() - if not overrides: - return config - - for key, value in overrides.items(): - if hasattr(config, key): - setattr(config, key, value) - else: - print(f"WARNING: unknown config field '{key}', ignoring") - - return config - - -def run_inspection( - inspection_dir: str, - model, - config: ContextManagerConfig, -) -> dict: - """Run summary inspection for a single inspection set. - - Reads: - - /history.json - - /checks.json - - Writes: - - /_result.json - - /_summary.txt (optional, if --save-summary) - - Returns: - dict with name, summary, checks, score, and compression metadata. - """ - name = os.path.basename(inspection_dir) - - # Load history - history_path = os.path.join(inspection_dir, "history.json") - if not os.path.exists(history_path): - print(f" SKIP: history.json not found in {inspection_dir}") - return {"name": name, "skipped": True, "reason": "no history.json"} - - with open(history_path, "r", encoding="utf-8") as f: - history = json.load(f) - - # Load checks - checks_path = os.path.join(inspection_dir, "checks.json") - if not os.path.exists(checks_path): - print(f" SKIP: checks.json not found in {inspection_dir}") - return {"name": name, "skipped": True, "reason": "no checks.json"} - - with open(checks_path, "r", encoding="utf-8") as f: - checks = json.load(f) - - if not checks: - print(f" SKIP: checks.json is empty for {name}") - return {"name": name, "skipped": True, "reason": "empty checks"} - - # Convert history to pairs - pairs = history_to_pairs(history) - print(f" History: {len(history)} messages -> {len(pairs)} pairs") - - # Compress - result = compress_history_offline(pairs=pairs, model=model, config=config) - summary = result.get("summary") or "" - is_fallback = result.get("is_fallback", False) - is_incremental = result.get("is_incremental", False) - input_chars = result.get("input_chars", 0) - - if not summary: - print(f" FAILED: compression returned no summary (fallback={is_fallback})") - report = { - "name": name, - "summary": None, - "is_fallback": is_fallback, - "input_chars": input_chars, - "checks": [], - "score": 0.0, - } - _write_result(inspection_dir, report) - return report - - print(f" Summary: {len(summary)} chars, fallback={is_fallback}, incremental={is_incremental}") - - # Evaluate checks against summary - check_results = [] - for check in checks: - eval_result = eval_text(summary, check) - check_results.append({ - "check": check, - "passed": eval_result.passed, - "score": eval_result.score, - "details": eval_result.details, - }) - - total_score = sum(r["score"] for r in check_results) / max(len(check_results), 1) - passed_count = sum(1 for r in check_results if r["passed"]) - - print(f" Result: {passed_count}/{len(check_results)} checks passed, score={total_score:.2f}") - - for r in check_results: - if not r["passed"]: - desc = r["check"].get("description", "") - keywords = r["check"].get("must_contain", r["check"].get("must_contain_any", [])) - print(f" FAIL: {desc} -- missing {keywords}") - - report = { - "name": name, - "summary": summary, - "is_fallback": is_fallback, - "is_incremental": is_incremental, - "input_chars": input_chars, - "summary_chars": len(summary), - "checks": check_results, - "score": total_score, - "passed": passed_count, - "total": len(check_results), - } - - _write_result(inspection_dir, report) - return report - - -def _write_result(inspection_dir: str, report: dict): - """Write _result.json (without full summary to keep file small) and optional _summary.txt.""" - result_path = os.path.join(inspection_dir, "_result.json") - result_out = {k: v for k, v in report.items() if k != "summary"} - with open(result_path, "w", encoding="utf-8") as f: - json.dump(result_out, f, ensure_ascii=False, indent=2, default=str) - print(f" Result saved to {result_path}") - - -def main(): - parser = argparse.ArgumentParser( - description="Standalone Summary Inspector -- quick compression quality check" - ) - parser.add_argument( - "-n", "--name", - type=str, - default=None, - help="Run a specific inspection by name (directory under inspections/)", - ) - parser.add_argument( - "--config", - type=str, - default=None, - help="Path to a JSON file with ContextManagerConfig field overrides", - ) - parser.add_argument( - "--save-summary", - action="store_true", - default=False, - help="Also save the raw summary text to _summary.txt alongside the result", - ) - args = parser.parse_args() - - # Discover inspections - if args.name: - inspection_dirs = [os.path.join(INSPECTIONS_DIR, args.name)] - if not os.path.isdir(inspection_dirs[0]): - print(f"ERROR: inspection directory not found: {inspection_dirs[0]}") - sys.exit(1) - else: - inspection_dirs = sorted(glob.glob(os.path.join(INSPECTIONS_DIR, "*/history.json"))) - inspection_dirs = [os.path.dirname(p) for p in inspection_dirs] - - if not inspection_dirs: - print(f"No inspections found under {INSPECTIONS_DIR}/*/\n" - f"Create one with: mkdir -p {INSPECTIONS_DIR}/my_test\n" - f"Then add history.json and checks.json") - sys.exit(1) - - # Build config - config_overrides = {} - if args.config: - with open(args.config, "r", encoding="utf-8") as f: - config_overrides = json.load(f) - - config = build_config(config_overrides) - config.enabled = True - - # Create model - model = create_model() - - # Run inspection for each - all_results = [] - for inspection_dir in inspection_dirs: - name = os.path.basename(inspection_dir) - print(f"\n===== Inspecting: {name} =====") - - report = run_inspection(inspection_dir, model, config) - all_results.append(report) - - # Optionally save raw summary text - if args.save_summary and report.get("summary"): - summary_path = os.path.join(inspection_dir, "_summary.txt") - with open(summary_path, "w", encoding="utf-8") as f: - f.write(report["summary"]) - print(f" Summary saved to {summary_path}") - - # Print overall summary - print("\n===== Overall =====") - for r in all_results: - if r.get("skipped"): - print(f" {r['name']}: SKIPPED ({r['reason']})") - else: - print(f" {r['name']}: {r.get('passed', 0)}/{r.get('total', 0)} passed, score={r.get('score', 0):.2f}") - - active = [r for r in all_results if not r.get("skipped")] - if active: - avg_score = sum(r.get("score", 0) for r in active) / max(len(active), 1) - print(f"\n Average score: {avg_score:.2f}") - - -if __name__ == "__main__": +# -*- coding: utf-8 -*- +""" +Standalone Summary Inspector — quick evaluation of compression prompt/schema quality. + +Completely independent from test_benchmark.py and the cases/ directory. +Uses compress_history_offline to compress history and checks whether the +resulting summary retains key information. No agent runs needed — just +one LLM call per inspection + text-based checks. + +Use case: + - Iterate on summary prompt / schema in summary_config.py + - Verify that key facts survive compression without running full agent loops + - Compare different ContextManagerConfig settings side-by-side + +Directory layout (independent from cases/ and reports/): + + inspections/ + └── / + ├── history.json # [{"role": "user|assistant", "content": "..."}] + └── checks.json # [{"description": "...", "must_contain": [...]}] + +Result is written to inspections//_result.json (co-located with input). + +Usage: + python summary_inspector.py # all inspections + python summary_inspector.py -n example_infra # single inspection + python summary_inspector.py --config my_config.json # custom config overrides + python summary_inspector.py --save-summary # also save raw summary .txt +""" + +import argparse +import json +import os +import sys +import glob + +# ============ Path Setup ============ +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import paths # noqa: F401 — side-effect: adds sdk/, backend/ to sys.path + +from dotenv import load_dotenv +load_dotenv() + +from nexent.core.agents.agent_context import compress_history_offline, ContextManagerConfig +from nexent.core.agents.agent_model import ModelConfig +from nexent.core.models.openai_llm import OpenAIModel + +from eval_utils import eval_text + + +# ============ Config ============ +LLM_API_KEY = os.getenv("LLM_API_KEY") +LLM_MODEL_NAME = os.getenv("LLM_MODEL_NAME") +LLM_API_URL = os.getenv("LLM_API_URL") + +INSPECTIONS_DIR = "./inspections" + + +def create_model(temperature: float = 0.1): + """Create an LLM model for offline compression.""" + from nexent.core.utils.observer import MessageObserver + + model_config = ModelConfig( + cite_name="inspector_model", + api_key=LLM_API_KEY, + model_name=LLM_MODEL_NAME, + url=LLM_API_URL, + temperature=temperature, + ssl_verify=False, + ) + return OpenAIModel( + observer=MessageObserver(), + model_id=model_config.model_name, + api_key=model_config.api_key, + api_base=model_config.url, + temperature=model_config.temperature, + top_p=model_config.top_p, + ssl_verify=model_config.ssl_verify, + ) + + +def history_to_pairs(history: list) -> list[tuple[str, str]]: + """Convert [{role, content}] to [(user_text, assistant_text)] pairs. + + Consecutive user messages are merged; same for assistant messages, + so the output is a clean alternating sequence of pairs. + """ + pairs = [] + current_user = [] + current_assistant = [] + + for entry in history: + role = entry["role"] + content = entry["content"] + if role == "user": + if current_assistant: + pairs.append(( + "\n".join(current_user).strip(), + "\n".join(current_assistant).strip(), + )) + current_user = [] + current_assistant = [] + current_user.append(content) + elif role == "assistant": + current_assistant.append(content) + + if current_user and current_assistant: + pairs.append(( + "\n".join(current_user).strip(), + "\n".join(current_assistant).strip(), + )) + + return pairs + + +def build_config(overrides: dict = None) -> ContextManagerConfig: + """Build ContextManagerConfig with optional field overrides.""" + config = ContextManagerConfig() + if not overrides: + return config + + for key, value in overrides.items(): + if hasattr(config, key): + setattr(config, key, value) + else: + print(f"WARNING: unknown config field '{key}', ignoring") + + return config + + +def run_inspection( + inspection_dir: str, + model, + config: ContextManagerConfig, +) -> dict: + """Run summary inspection for a single inspection set. + + Reads: + - /history.json + - /checks.json + + Writes: + - /_result.json + - /_summary.txt (optional, if --save-summary) + + Returns: + dict with name, summary, checks, score, and compression metadata. + """ + name = os.path.basename(inspection_dir) + + # Load history + history_path = os.path.join(inspection_dir, "history.json") + if not os.path.exists(history_path): + print(f" SKIP: history.json not found in {inspection_dir}") + return {"name": name, "skipped": True, "reason": "no history.json"} + + with open(history_path, "r", encoding="utf-8") as f: + history = json.load(f) + + # Load checks + checks_path = os.path.join(inspection_dir, "checks.json") + if not os.path.exists(checks_path): + print(f" SKIP: checks.json not found in {inspection_dir}") + return {"name": name, "skipped": True, "reason": "no checks.json"} + + with open(checks_path, "r", encoding="utf-8") as f: + checks = json.load(f) + + if not checks: + print(f" SKIP: checks.json is empty for {name}") + return {"name": name, "skipped": True, "reason": "empty checks"} + + # Convert history to pairs + pairs = history_to_pairs(history) + print(f" History: {len(history)} messages -> {len(pairs)} pairs") + + # Compress + result = compress_history_offline(pairs=pairs, model=model, config=config) + summary = result.get("summary") or "" + is_fallback = result.get("is_fallback", False) + is_incremental = result.get("is_incremental", False) + input_chars = result.get("input_chars", 0) + + if not summary: + print(f" FAILED: compression returned no summary (fallback={is_fallback})") + report = { + "name": name, + "summary": None, + "is_fallback": is_fallback, + "input_chars": input_chars, + "checks": [], + "score": 0.0, + } + _write_result(inspection_dir, report) + return report + + print(f" Summary: {len(summary)} chars, fallback={is_fallback}, incremental={is_incremental}") + + # Evaluate checks against summary + check_results = [] + for check in checks: + eval_result = eval_text(summary, check) + check_results.append({ + "check": check, + "passed": eval_result.passed, + "score": eval_result.score, + "details": eval_result.details, + }) + + total_score = sum(r["score"] for r in check_results) / max(len(check_results), 1) + passed_count = sum(1 for r in check_results if r["passed"]) + + print(f" Result: {passed_count}/{len(check_results)} checks passed, score={total_score:.2f}") + + for r in check_results: + if not r["passed"]: + desc = r["check"].get("description", "") + keywords = r["check"].get("must_contain", r["check"].get("must_contain_any", [])) + print(f" FAIL: {desc} -- missing {keywords}") + + report = { + "name": name, + "summary": summary, + "is_fallback": is_fallback, + "is_incremental": is_incremental, + "input_chars": input_chars, + "summary_chars": len(summary), + "checks": check_results, + "score": total_score, + "passed": passed_count, + "total": len(check_results), + } + + _write_result(inspection_dir, report) + return report + + +def _write_result(inspection_dir: str, report: dict): + """Write _result.json (without full summary to keep file small) and optional _summary.txt.""" + result_path = os.path.join(inspection_dir, "_result.json") + result_out = {k: v for k, v in report.items() if k != "summary"} + with open(result_path, "w", encoding="utf-8") as f: + json.dump(result_out, f, ensure_ascii=False, indent=2, default=str) + print(f" Result saved to {result_path}") + + +def main(): + parser = argparse.ArgumentParser( + description="Standalone Summary Inspector -- quick compression quality check" + ) + parser.add_argument( + "-n", "--name", + type=str, + default=None, + help="Run a specific inspection by name (directory under inspections/)", + ) + parser.add_argument( + "--config", + type=str, + default=None, + help="Path to a JSON file with ContextManagerConfig field overrides", + ) + parser.add_argument( + "--save-summary", + action="store_true", + default=False, + help="Also save the raw summary text to _summary.txt alongside the result", + ) + args = parser.parse_args() + + # Discover inspections + if args.name: + inspection_dirs = [os.path.join(INSPECTIONS_DIR, args.name)] + if not os.path.isdir(inspection_dirs[0]): + print(f"ERROR: inspection directory not found: {inspection_dirs[0]}") + sys.exit(1) + else: + inspection_dirs = sorted(glob.glob(os.path.join(INSPECTIONS_DIR, "*/history.json"))) + inspection_dirs = [os.path.dirname(p) for p in inspection_dirs] + + if not inspection_dirs: + print(f"No inspections found under {INSPECTIONS_DIR}/*/\n" + f"Create one with: mkdir -p {INSPECTIONS_DIR}/my_test\n" + f"Then add history.json and checks.json") + sys.exit(1) + + # Build config + config_overrides = {} + if args.config: + with open(args.config, "r", encoding="utf-8") as f: + config_overrides = json.load(f) + + config = build_config(config_overrides) + config.enabled = True + + # Create model + model = create_model() + + # Run inspection for each + all_results = [] + for inspection_dir in inspection_dirs: + name = os.path.basename(inspection_dir) + print(f"\n===== Inspecting: {name} =====") + + report = run_inspection(inspection_dir, model, config) + all_results.append(report) + + # Optionally save raw summary text + if args.save_summary and report.get("summary"): + summary_path = os.path.join(inspection_dir, "_summary.txt") + with open(summary_path, "w", encoding="utf-8") as f: + f.write(report["summary"]) + print(f" Summary saved to {summary_path}") + + # Print overall summary + print("\n===== Overall =====") + for r in all_results: + if r.get("skipped"): + print(f" {r['name']}: SKIPPED ({r['reason']})") + else: + print(f" {r['name']}: {r.get('passed', 0)}/{r.get('total', 0)} passed, score={r.get('score', 0):.2f}") + + active = [r for r in all_results if not r.get("skipped")] + if active: + avg_score = sum(r.get("score", 0) for r in active) / max(len(active), 1) + print(f"\n Average score: {avg_score:.2f}") + + +if __name__ == "__main__": main() \ No newline at end of file diff --git a/sdk/benchmark/manual_cases/test_benchmark.py b/sdk/benchmark/manual_cases/test_benchmark.py index 75dd5f6fb..214e1bc4c 100644 --- a/sdk/benchmark/manual_cases/test_benchmark.py +++ b/sdk/benchmark/manual_cases/test_benchmark.py @@ -1,726 +1,726 @@ -import asyncio -import copy -import glob -import json -import os -import sys -import argparse - -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -import paths # noqa: F401 — side-effect: adds sdk/, backend/ to sys.path - -from agent_runner import ( - build_agent_run_info_with_custom_prompt, - run_agent_with_tracking, - parse_conversation_to_history, - AgentHistory, - ContextManagerConfig, -) - -from nexent.core.agents.agent_context import ContextManager -from nexent.core.utils.token_estimation import estimate_tokens_text - -from eval_utils import eval_text, average_score - -# Lean benchmark system prompt — generic, not task-specific. -# Strips the verbose platform scaffolding (File URL Guide, Reference Marks, -# safety principles, etc.) to minimize token overhead while retaining the -# core execution loop instructions the agent needs to function. -BENCHMARK_SYSTEM_PROMPT = """You are a helpful assistant. Answer the user's questions based on the conversation history and your knowledge. - -- Be precise and concise. -- When the answer depends on information from earlier conversation, refer to it accurately. -- Do not fabricate information you do not know. -- Use final_answer to submit your response. - -Now start!""" - - -# --- Custom summary schema and prompts for knowledge-discussion benchmarks --- -# These override the default 10-field Hermes schema from summary_config.py -# with a deduplicated 6-field schema (~620 word budget) that merges -# completed_work + resolved_questions into "progress" and restricts -# key_facts to values NOT already stated in progress, eliminating -# the 3-field redundancy that caused output bloat in incremental updates. -# -# KEY DESIGN PRINCIPLE for incremental compression: the output must be -# approximately the SAME size as the initial summary (~620 words). The -# incremental prompt treats old+new as a unified corpus and REWRITES the -# entire summary from scratch, rather than appending to the old one. -# This prevents output-token linear growth that would itself exceed -# token_threshold and defeat the purpose of compression. - -BENCHMARK_SUMMARY_SYSTEM_PROMPT = ( - "You are a summarization agent creating a compact working-memory checkpoint. " - "Treat the conversation turns below as source material, not as a transcript to preserve. " - "Your job is to produce a fixed-size JSON summary that preserves only the information " - "needed to continue the conversation correctly later.\n\n" - - "Output rules:\n" - "1. Produce only strict JSON. Do not add greeting, preamble, markdown, or explanation.\n" - "2. Write in the same language as the user's most recent message. Do not translate unless needed.\n" - "3. Never include API keys, tokens, passwords, secrets, credentials, or connection strings. " - "Replace any such values with [REDACTED].\n\n" - - "Compression goal:\n" - "The summary is working memory, not a historical log. " - "Do not list every question, every answer, or every conversation turn. " - "Group information by theme and keep only facts that are likely to matter for future continuation.\n\n" - - "Field constraints:\n" - "1. 'active_task' must describe only the current unfulfilled user request; if none, write 'None'.\n" - "2. 'goal' must describe the current overall objective in <=25 words.\n" - "3. 'state' must contain at most 6 numbered items. Never create item 7 or higher. " - "Each item must be <=45 words. Merge related topics into one item. " - "Do not organize by conversation order; organize by semantic importance.\n" - "4. 'decisions' must contain at most 5 short confirmed conclusions or choices. " - "Do not repeat facts already fully stated in 'state'.\n" - "5. 'open_items' must contain only unresolved questions or pending user requests. " - "If none, write 'None'.\n" - "6. 'verbatim_facts' may contain at most 12 raw values, formulas, thresholds, exact model names, " - "or identifiers that must be copied exactly later. " - "Before output, remove any item whose exact value already appears in 'state' or 'decisions'. " - "If no extra raw facts remain, write 'None'.\n\n" - - "Information priority:\n" - "Critical current task and constraints > final conclusions > decisions > exact values needed later > " - "background context. Drop vague descriptions, repeated facts, superseded intermediate reasoning, " - "and completed Q&A that no longer affects future work.\n\n" - - "Budget:\n" - "The total output must not exceed 620 words. Prefer shorter output. " - "If the content is too large, compress in this order: " - "(1) merge related state items; " - "(2) remove completed historical details; " - "(3) keep only the most diagnostic numbers; " - "(4) move only non-duplicated raw values to 'verbatim_facts'; " - "(5) write 'None' for fields with no current utility.\n\n" - - "Return strict JSON only." -) - - -BENCHMARK_INCREMENTAL_SUMMARY_SYSTEM_PROMPT = ( - "You are a summarization agent rewriting a compact working-memory checkpoint. " - "You receive a Previous Summary and New Conversations. Produce one fresh JSON summary " - "that preserves only the information needed to continue the conversation correctly. " - "Do not preserve discussion history for its own sake. The previous summary is source material, " - "not text to copy.\n\n" - - "Hard constraints:\n" - "1. The output must be no longer than the previous summary and must not exceed 620 words.\n" - "2. The 'state' field must contain at most 6 numbered items. Never create item 7 or higher.\n" - "3. When new information is added, older lower-utility information MUST be merged, generalized, or deleted.\n" - "4. Do not append to the previous summary. Rewrite by theme, not by conversation order.\n" - "5. Completed Q&A should become conclusions, not separate historical entries.\n" - "6. Preserve exact numbers only when they are needed for future correctness. If multiple numbers support the same conclusion, keep only the most diagnostic ones.\n" - "7. 'verbatim_facts' may contain at most 12 raw values/formulas/names. Remove any item already present in 'state' or 'decisions'. If none remain, write 'None'.\n" - "8. Update active_task, state, and open_items to reflect the current state.\n" - "9. Write in the same language as the user's most recent message.\n" - "10. Never include API keys, tokens, passwords, credentials, or connection strings; replace them with [REDACTED].\n\n" - - "Output strict JSON only. No markdown." -) - -BENCHMARK_SUMMARY_SCHEMA = { - "active_task": ( - "用户当前尚未完成的最新请求；如果没有，写 'None'。" - "必须是当前任务，不是历史任务。<=25 words" - ), - - "goal": ( - "对话的总体目标或当前工作方向。" - "只保留后续继续对话所需的目标。<=25 words" - ), - - "state": ( - "当前压缩后的工作记忆，不是历史日志。" - "最多 6 条编号条目；每条 <=45 words。" - "按主题合并信息，不按对话顺序罗列。" - "包括已经确定的结论、关键设计、关键结果和必要上下文。" - ), - - "decisions": ( - "已经确认、后续可能需要引用的结论或选择。" - "最多 5 条；每条 <=25 words。" - "不得重复 state 中已经完整表达的信息。" - ), - - "open_items": ( - "尚未解决的问题、待办事项或用户明确要求继续处理的内容。" - "如果没有，写 'None'。<=30 words" - ), - - "verbatim_facts": ( - "必须逐字保留的数字、公式、模型名、阈值或专有名词。" - "最多 12 项，用分号分隔。" - "不得包含已经出现在 state 或 decisions 中的事实。" - "如果没有额外需要保留的事实，写 'None'。" - ), -} -def history_to_text(history: list[AgentHistory]) -> str: - return "\n".join([f"{h.role}: {h.content}" for h in history]) - - -async def run_multi_turn_for_benchmark( - queries: list[str], - base_history: list[AgentHistory], - cm_config: ContextManagerConfig, - max_steps: int = 20, - system_prompt: str = BENCHMARK_SYSTEM_PROMPT, -): - conversation_history = list(base_history) - results = [] - - shared_cm = None - if cm_config and cm_config.enabled: - shared_cm = ContextManager(config=cm_config, max_steps=max_steps) - - initial_tokens = estimate_tokens_text(history_to_text(conversation_history)) - - # Track per-step actual input tokens for accurate token reduction - step_input_tokens = [] - - for query in queries: - agent_run_info = build_agent_run_info_with_custom_prompt( - query, - system_prompt, - conversation_history, - max_steps=max_steps, - context_manager_config=cm_config, - ) - - if shared_cm is not None: - agent_run_info.context_manager = shared_cm - - result = await run_agent_with_tracking(agent_run_info, debug=False) - results.append(result) - - # Collect actual input token count from the last step metrics - if shared_cm is not None: - tc = shared_cm.get_token_counts() - step_input_tokens.append(tc) - - conversation_history.append(AgentHistory(role="user", content=query)) - conversation_history.append( - AgentHistory(role="assistant", content=result.final_answer) - ) - - final_tokens = estimate_tokens_text(history_to_text(conversation_history)) - - cm_stats = None - cm_token_counts = None - cm_summary = None - if shared_cm is not None: - cm_stats = shared_cm.get_all_compression_stats() - cm_token_counts = shared_cm.get_token_counts() - cm_summary = shared_cm.export_summary() - - return { - "results": results, - "conversation_history": conversation_history, - "shared_cm": shared_cm, - "initial_tokens": initial_tokens, - "final_tokens": final_tokens, - "cm_stats": cm_stats, - "cm_token_counts": cm_token_counts, - "cm_summary": cm_summary, - "step_input_tokens": step_input_tokens, - } - - -def build_precompressed_history( - frozen_history: list[AgentHistory], - cm_summary: dict, -) -> list[AgentHistory]: - """Build a pre-compressed history from the compression snapshot. - - Replaces the compressed prefix pairs with a single user message containing - the summary text, then appends the retained tail pairs verbatim. This - mirrors the actual message structure produced by compress_if_needed: - - SummaryTaskStep.to_messages() → [ChatMessage(role=USER, summary)] - followed by retained tail steps → [TaskStep, ActionStep, ...] - - There is NO assistant message after the summary — the model sees the - summary as a user message, followed directly by the next retained step. - - Args: - frozen_history: The original uncompressed conversation history. - cm_summary: The export_summary() dict from the compressed run's - ContextManager, containing summary text and boundary info. - - Returns: - A new AgentHistory list that mirrors the compressed context structure. - """ - boundary = cm_summary.get("compression_boundary", {}) - compressed_pairs = boundary.get("previous_compressed_pairs", 0) - - # Each pair = 2 AgentHistory entries (user + assistant) - compressed_entries = compressed_pairs * 2 - - summary_text = cm_summary.get("previous_summary") or "" - - # If no compression happened, return original history unchanged - if not summary_text or compressed_entries == 0: - return list(frozen_history) - - # Build pre-compressed history: - # 1. Summary as a single USER message (matching SummaryTaskStep.to_messages) - # No paired assistant message — the model sees summary then next retained step - precompressed = [ - AgentHistory( - role="user", - content=f"Summary of earlier steps in this task:\n{summary_text}", - ), - ] - - # 2. Retained tail pairs (everything after the compressed prefix) - if compressed_entries < len(frozen_history): - precompressed.extend(frozen_history[compressed_entries:]) - - return precompressed - - -async def run_probe_questions( - probes: list[dict], - precompressed_history: list[AgentHistory], - max_steps: int = 20, - system_prompt: str = BENCHMARK_SYSTEM_PROMPT, -): - """Run probe questions against a pre-compressed history snapshot. - - Each probe runs independently with compression DISABLED, because the - history has already been pre-compressed (compressed prefix replaced with - summary text, retained tail kept verbatim). This avoids redundant LLM - compression calls — the compression was done once in the compressed run, - and all probes reuse that result. - - Per CLAUDE.md rules: - - Each probe uses a deep-copied frozen snapshot - - Probes see compressed context (summary + retained tail) - - No compression triggered during probe phase - - Probes are fully independent, no shared state - """ - probe_results = [] - no_compression_config = ContextManagerConfig(enabled=False, token_threshold=10**9) - - for probe in probes: - question = probe["question"] - - # Each probe gets its own deep copy — fully independent - probe_history = copy.deepcopy(precompressed_history) - - agent_run_info = build_agent_run_info_with_custom_prompt( - question, - system_prompt, - probe_history, - max_steps=max_steps, - context_manager_config=no_compression_config, - ) - - result = await run_agent_with_tracking(agent_run_info, debug=False) - eval_result = eval_text(result.final_answer, probe) - - probe_results.append( - { - "question": question, - "answer": result.final_answer, - "passed": eval_result.passed, - "score": eval_result.score, - "details": eval_result.details, - } - ) - - return probe_results - - -async def run_baseline_probes( - probes: list[dict], - frozen_history: list[AgentHistory], - max_steps: int = 20, - system_prompt: str = BENCHMARK_SYSTEM_PROMPT, -): - """Run probe questions against full uncompressed history (baseline). - - This measures the ceiling: what can the agent answer when it sees - the complete history. probe_retention = compressed_score / baseline_score. - """ - probe_results = [] - baseline_config = ContextManagerConfig(enabled=False, token_threshold=10**9) - - for probe in probes: - question = probe["question"] - probe_history = copy.deepcopy(frozen_history) - - agent_run_info = build_agent_run_info_with_custom_prompt( - question, - system_prompt, - probe_history, - max_steps=max_steps, - context_manager_config=baseline_config, - ) - - result = await run_agent_with_tracking(agent_run_info, debug=False) - eval_result = eval_text(result.final_answer, probe) - - probe_results.append( - { - "question": question, - "answer": result.final_answer, - "passed": eval_result.passed, - "score": eval_result.score, - "details": eval_result.details, - } - ) - - return probe_results - - -def eval_summary_inspection(summary: dict, checks: list[dict]) -> list[dict]: - """Static Compression Inspection — check if the compressed summary - retains key information (user preferences, file names, plans, tool results). - - Uses dedicated summary_checks when available, NOT probe must_contain - (which has different semantics — probe keywords are for agent answers, - summary keywords are for what the compressor chose to preserve). - """ - results = [] - - prev_summary = summary.get("previous_summary") or "" - curr_summary = summary.get("current_summary") or "" - combined = prev_summary + "\n" + curr_summary - - for check in checks: - eval_result = eval_text(combined, check) - results.append( - { - "check": check, - "passed": eval_result.passed, - "score": eval_result.score, - "details": eval_result.details, - } - ) - - return results - - -def eval_task_outputs(case: dict, run_outputs: list): - eval_results = [] - - for check in case.get("task_checks", []): - turn_idx = check["turn"] - 1 - if turn_idx >= len(run_outputs): - continue - - answer = run_outputs[turn_idx].final_answer - r = eval_text(answer, check) - - eval_results.append( - { - "turn": check["turn"], - "answer": answer, - "passed": r.passed, - "score": r.score, - "details": r.details, - } - ) - - return eval_results - - -def _resolve_compressed_config(case: dict, use_default_prompts: bool = False) -> ContextManagerConfig: - """Build compressed config from case definition, with sensible defaults. - - By default uses the benchmark-optimized custom summary schema and prompts. - Set use_default_prompts=True to fall back to the original ContextManager defaults. - """ - case_cfg = case.get("compressed_config", {}) - kwargs = dict( - enabled=True, - token_threshold=case_cfg.get("token_threshold", 3600), - keep_recent_pairs=case_cfg.get("keep_recent_pairs", 1), - keep_recent_steps=case_cfg.get("keep_recent_steps", 4), - max_observation_length=case_cfg.get("max_observation_length", 20000), - ) - if not use_default_prompts: - kwargs.update( - summary_json_schema=BENCHMARK_SUMMARY_SCHEMA, - summary_system_prompt=BENCHMARK_SUMMARY_SYSTEM_PROMPT, - incremental_summary_system_prompt=BENCHMARK_INCREMENTAL_SUMMARY_SYSTEM_PROMPT, - ) - return ContextManagerConfig(**kwargs) - - -async def run_one_case(case_dir: str, use_default_prompts: bool = False): - """Load and run a single benchmark case from its directory. - - Each case directory contains: - - case.json: queries, probes, summary_checks, task_checks, compressed_config - - history.json: conversation history - - Args: - case_dir: Absolute or relative path to the case directory. - - Returns: - Report dict for this case. - """ - case_path = os.path.join(case_dir, "case.json") - with open(case_path, "r", encoding="utf-8") as f: - case = json.load(f) - - # Resolve history_file relative to the case directory; - # defaults to "history.json" in the same directory if not specified. - history_relpath = case.get("history_file", "history.json") - history_abspath = os.path.join(case_dir, history_relpath) - - base_history = parse_conversation_to_history(history_abspath) - - baseline_config = ContextManagerConfig( - enabled=False, - token_threshold=10**9, - keep_recent_pairs=1, - ) - - # P5: Allow per-case config override - compressed_config = _resolve_compressed_config(case, use_default_prompts=use_default_prompts) - - print(f"\n===== CASE: {case['id']} =====") - - baseline = await run_multi_turn_for_benchmark( - queries=case["queries"], - base_history=base_history, - cm_config=baseline_config, - ) - - compressed = await run_multi_turn_for_benchmark( - queries=case["queries"], - base_history=base_history, - cm_config=compressed_config, - ) - - baseline_task_eval = eval_task_outputs(case, baseline["results"]) - compressed_task_eval = eval_task_outputs(case, compressed["results"]) - # P1: Baseline probe — agent sees full uncompressed history - # Same frozen_history, but with compression disabled, so the agent sees - # the complete unmodified context. This establishes the ceiling for - # probe_retention = compressed_probe_score / baseline_probe_score. - baseline_probe_eval = await run_baseline_probes( - probes=case["probes"], - frozen_history=compressed["conversation_history"], - max_steps=20, - ) - - # P0: Compressed probe — agent sees pre-compressed context - # Build the pre-compressed history ONCE using the summary from the - # compressed run's ContextManager, then run each probe independently - # against it with compression disabled. This avoids redundant LLM calls - # (compression was already done in the compressed multi-turn run). - precompressed_history = build_precompressed_history( - frozen_history=compressed["conversation_history"], - cm_summary=compressed["cm_summary"] or {}, - ) - compressed_probe_eval = await run_probe_questions( - probes=case["probes"], - precompressed_history=precompressed_history, - ) - - # P3: Summary inspection uses dedicated summary_checks, not probe must_contain - summary_inspection = [] - if compressed.get("cm_summary"): - summary_checks = case.get("summary_checks", []) - if summary_checks: - summary_inspection = eval_summary_inspection( - compressed["cm_summary"], summary_checks - ) - - baseline_task_score = sum(x["score"] for x in baseline_task_eval) / max( - len(baseline_task_eval), 1 - ) - - compressed_task_score = sum(x["score"] for x in compressed_task_eval) / max( - len(compressed_task_eval), 1 - ) - - baseline_probe_score = sum(x["score"] for x in baseline_probe_eval) / max( - len(baseline_probe_eval), 1 - ) - - compressed_probe_score = sum(x["score"] for x in compressed_probe_eval) / max( - len(compressed_probe_eval), 1 - ) - - summary_score = ( - sum(x["score"] for x in summary_inspection) / max(len(summary_inspection), 1) - if summary_inspection - else None - ) - - task_success_retention = ( - compressed_task_score / baseline_task_score - if baseline_task_score > 0 - else 0.0 - ) - - probe_retention = ( - compressed_probe_score / baseline_probe_score - if baseline_probe_score > 0 - else 0.0 - ) - - # P2: Token reduction from actual input token counts - # Use the last step's token counts (final compressed vs uncompressed state) - token_reduction = 0.0 - if compressed.get("step_input_tokens") and compressed["step_input_tokens"]: - last_tc = compressed["step_input_tokens"][-1] - if last_tc and last_tc.get("last_uncompressed") is not None: - unc = last_tc["last_uncompressed"] or 1 - comp = last_tc["last_compressed"] or 0 - if unc > 0: - token_reduction = 1 - comp / unc - # Fallback to text-based estimation - if token_reduction == 0.0: - token_reduction = 1 - ( - compressed["final_tokens"] / max(baseline["final_tokens"], 1) - ) - baseline_failed = baseline_task_score == 0 - - # Compute real main-LLM input token totals - baseline_real_input = sum(r.total_input_tokens for r in baseline["results"]) - compressed_real_input = sum(r.total_input_tokens for r in compressed["results"]) - - # Compression cost: tokens spent on compression LLM calls - compression_cost = 0 - if compressed.get("cm_stats"): - compression_cost = ( - compressed["cm_stats"].get("total_input_tokens", 0) - + compressed["cm_stats"].get("total_output_tokens", 0) - ) - - # Net token reduction = gross savings - compression cost - gross_input_savings = baseline_real_input - compressed_real_input - net_input_savings = gross_input_savings - compression_cost - net_token_reduction = ( - net_input_savings / max(baseline_real_input, 1) - if baseline_real_input > 0 - else 0.0 - ) - - report = { - "case_id": case["id"], - "baseline_failed": baseline_failed, - "baseline": { - "task_score": baseline_task_score, - "probe_score": baseline_probe_score, - "final_tokens": baseline["final_tokens"], - "real_input_tokens": baseline_real_input, - }, - "compressed": { - "task_score": compressed_task_score, - "probe_score": compressed_probe_score, - "final_tokens": compressed["final_tokens"], - "cm_stats": compressed["cm_stats"], - "cm_token_counts": compressed["cm_token_counts"], - "cm_summary": compressed["cm_summary"], - "real_input_tokens": compressed_real_input, - }, - "metrics": { - "task_success_retention": task_success_retention, - "probe_retention": probe_retention, - "token_reduction": token_reduction, - "net_token_reduction": net_token_reduction, - "compression_cost_tokens": compression_cost, - "summary_score": summary_score, - }, - "task_eval": compressed_task_eval, - "probe_eval": { - "baseline": baseline_probe_eval, - "compressed": compressed_probe_eval, - }, - "summary_inspection": summary_inspection, - } - - print(json.dumps(report, ensure_ascii=False, indent=2, default=str)) - return report - - -async def main(case_names: list[str] = None, use_default_prompts: bool = False): - # Discover cases: use specified names if provided, otherwise find all cases under ./cases/*/case.json - if case_names: - case_dirs = [os.path.join("./cases", name) for name in case_names] - else: - case_dirs = sorted(glob.glob("./cases/*/case.json")) - case_dirs = [os.path.dirname(p) for p in case_dirs] - - if not case_dirs: - print("No benchmark cases found under ./cases/*/case.json") - return - - print(f"Found {len(case_dirs)} case(s): {[os.path.basename(d) for d in case_dirs]}") - - # Output directory for reports - os.makedirs("./reports", exist_ok=True) - - reports = [] - for case_dir in case_dirs: - report = await run_one_case(case_dir, use_default_prompts=use_default_prompts) - reports.append(report) - - # Write per-case report - case_id = report["case_id"] - per_case_path = os.path.join("./reports", f"{case_id}.json") - with open(per_case_path, "w", encoding="utf-8") as f: - json.dump(report, f, ensure_ascii=False, indent=2, default=str) - print(f" Report saved to {per_case_path}") - - # Exclude cases where baseline itself failed - valid_reports = [r for r in reports if not r.get("baseline_failed")] - excluded_ids = [r["case_id"] for r in reports if r.get("baseline_failed")] - if excluded_ids: - print(f"\n Excluded from average (baseline failed): {excluded_ids}") - # Write summary across all cases - summary = { - "total_cases": len(reports), - "excluded_cases": len(reports) - len(valid_reports), - "metrics": { - "avg_task_success_retention": sum( - r["metrics"]["task_success_retention"] for r in valid_reports - ) / max(len(valid_reports), 1), - "avg_probe_retention": sum( - r["metrics"]["probe_retention"] for r in valid_reports - ) / max(len(valid_reports), 1), - "avg_token_reduction": sum( - r["metrics"]["token_reduction"] for r in valid_reports - ) / max(len(valid_reports), 1), - "avg_net_token_reduction": sum( - r["metrics"]["net_token_reduction"] for r in valid_reports - ) / max(len(valid_reports), 1), - "avg_compression_cost_tokens": sum( - r["metrics"]["compression_cost_tokens"] for r in valid_reports - ) / max(len(valid_reports), 1), - "per_case": { - r["case_id"]: r["metrics"] for r in reports - }, - }, - } - summary_path = "./reports/summary.json" - with open(summary_path, "w", encoding="utf-8") as f: - json.dump(summary, f, ensure_ascii=False, indent=2, default=str) - - print(f"\nBenchmark finished. Summary saved to {summary_path}") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Run Agent Context Compression Benchmark") - parser.add_argument( - "--cases",nargs="+",default=None, - help="Specific case names to run (e.g. --cases example_infra algotithm_data)." - "if omitted, run all cases under .cases/." - ) - parser.add_argument( - "--default-summary", action="store_true", default=False, - help="Use the original ContextManager summary defaults instead of the benchmark-optimized " - "custom schema (leaner 7-field, 800-word cap, merge-condense incremental updates)." - ) - args = parser.parse_args() +import asyncio +import copy +import glob +import json +import os +import sys +import argparse + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import paths # noqa: F401 — side-effect: adds sdk/, backend/ to sys.path + +from agent_runner import ( + build_agent_run_info_with_custom_prompt, + run_agent_with_tracking, + parse_conversation_to_history, + AgentHistory, + ContextManagerConfig, +) + +from nexent.core.agents.agent_context import ContextManager +from nexent.core.utils.token_estimation import estimate_tokens_text + +from eval_utils import eval_text, average_score + +# Lean benchmark system prompt — generic, not task-specific. +# Strips the verbose platform scaffolding (File URL Guide, Reference Marks, +# safety principles, etc.) to minimize token overhead while retaining the +# core execution loop instructions the agent needs to function. +BENCHMARK_SYSTEM_PROMPT = """You are a helpful assistant. Answer the user's questions based on the conversation history and your knowledge. + +- Be precise and concise. +- When the answer depends on information from earlier conversation, refer to it accurately. +- Do not fabricate information you do not know. +- Use final_answer to submit your response. + +Now start!""" + + +# --- Custom summary schema and prompts for knowledge-discussion benchmarks --- +# These override the default 10-field Hermes schema from summary_config.py +# with a deduplicated 6-field schema (~620 word budget) that merges +# completed_work + resolved_questions into "progress" and restricts +# key_facts to values NOT already stated in progress, eliminating +# the 3-field redundancy that caused output bloat in incremental updates. +# +# KEY DESIGN PRINCIPLE for incremental compression: the output must be +# approximately the SAME size as the initial summary (~620 words). The +# incremental prompt treats old+new as a unified corpus and REWRITES the +# entire summary from scratch, rather than appending to the old one. +# This prevents output-token linear growth that would itself exceed +# token_threshold and defeat the purpose of compression. + +BENCHMARK_SUMMARY_SYSTEM_PROMPT = ( + "You are a summarization agent creating a compact working-memory checkpoint. " + "Treat the conversation turns below as source material, not as a transcript to preserve. " + "Your job is to produce a fixed-size JSON summary that preserves only the information " + "needed to continue the conversation correctly later.\n\n" + + "Output rules:\n" + "1. Produce only strict JSON. Do not add greeting, preamble, markdown, or explanation.\n" + "2. Write in the same language as the user's most recent message. Do not translate unless needed.\n" + "3. Never include API keys, tokens, passwords, secrets, credentials, or connection strings. " + "Replace any such values with [REDACTED].\n\n" + + "Compression goal:\n" + "The summary is working memory, not a historical log. " + "Do not list every question, every answer, or every conversation turn. " + "Group information by theme and keep only facts that are likely to matter for future continuation.\n\n" + + "Field constraints:\n" + "1. 'active_task' must describe only the current unfulfilled user request; if none, write 'None'.\n" + "2. 'goal' must describe the current overall objective in <=25 words.\n" + "3. 'state' must contain at most 6 numbered items. Never create item 7 or higher. " + "Each item must be <=45 words. Merge related topics into one item. " + "Do not organize by conversation order; organize by semantic importance.\n" + "4. 'decisions' must contain at most 5 short confirmed conclusions or choices. " + "Do not repeat facts already fully stated in 'state'.\n" + "5. 'open_items' must contain only unresolved questions or pending user requests. " + "If none, write 'None'.\n" + "6. 'verbatim_facts' may contain at most 12 raw values, formulas, thresholds, exact model names, " + "or identifiers that must be copied exactly later. " + "Before output, remove any item whose exact value already appears in 'state' or 'decisions'. " + "If no extra raw facts remain, write 'None'.\n\n" + + "Information priority:\n" + "Critical current task and constraints > final conclusions > decisions > exact values needed later > " + "background context. Drop vague descriptions, repeated facts, superseded intermediate reasoning, " + "and completed Q&A that no longer affects future work.\n\n" + + "Budget:\n" + "The total output must not exceed 620 words. Prefer shorter output. " + "If the content is too large, compress in this order: " + "(1) merge related state items; " + "(2) remove completed historical details; " + "(3) keep only the most diagnostic numbers; " + "(4) move only non-duplicated raw values to 'verbatim_facts'; " + "(5) write 'None' for fields with no current utility.\n\n" + + "Return strict JSON only." +) + + +BENCHMARK_INCREMENTAL_SUMMARY_SYSTEM_PROMPT = ( + "You are a summarization agent rewriting a compact working-memory checkpoint. " + "You receive a Previous Summary and New Conversations. Produce one fresh JSON summary " + "that preserves only the information needed to continue the conversation correctly. " + "Do not preserve discussion history for its own sake. The previous summary is source material, " + "not text to copy.\n\n" + + "Hard constraints:\n" + "1. The output must be no longer than the previous summary and must not exceed 620 words.\n" + "2. The 'state' field must contain at most 6 numbered items. Never create item 7 or higher.\n" + "3. When new information is added, older lower-utility information MUST be merged, generalized, or deleted.\n" + "4. Do not append to the previous summary. Rewrite by theme, not by conversation order.\n" + "5. Completed Q&A should become conclusions, not separate historical entries.\n" + "6. Preserve exact numbers only when they are needed for future correctness. If multiple numbers support the same conclusion, keep only the most diagnostic ones.\n" + "7. 'verbatim_facts' may contain at most 12 raw values/formulas/names. Remove any item already present in 'state' or 'decisions'. If none remain, write 'None'.\n" + "8. Update active_task, state, and open_items to reflect the current state.\n" + "9. Write in the same language as the user's most recent message.\n" + "10. Never include API keys, tokens, passwords, credentials, or connection strings; replace them with [REDACTED].\n\n" + + "Output strict JSON only. No markdown." +) + +BENCHMARK_SUMMARY_SCHEMA = { + "active_task": ( + "用户当前尚未完成的最新请求；如果没有，写 'None'。" + "必须是当前任务，不是历史任务。<=25 words" + ), + + "goal": ( + "对话的总体目标或当前工作方向。" + "只保留后续继续对话所需的目标。<=25 words" + ), + + "state": ( + "当前压缩后的工作记忆，不是历史日志。" + "最多 6 条编号条目；每条 <=45 words。" + "按主题合并信息，不按对话顺序罗列。" + "包括已经确定的结论、关键设计、关键结果和必要上下文。" + ), + + "decisions": ( + "已经确认、后续可能需要引用的结论或选择。" + "最多 5 条；每条 <=25 words。" + "不得重复 state 中已经完整表达的信息。" + ), + + "open_items": ( + "尚未解决的问题、待办事项或用户明确要求继续处理的内容。" + "如果没有，写 'None'。<=30 words" + ), + + "verbatim_facts": ( + "必须逐字保留的数字、公式、模型名、阈值或专有名词。" + "最多 12 项，用分号分隔。" + "不得包含已经出现在 state 或 decisions 中的事实。" + "如果没有额外需要保留的事实，写 'None'。" + ), +} +def history_to_text(history: list[AgentHistory]) -> str: + return "\n".join([f"{h.role}: {h.content}" for h in history]) + + +async def run_multi_turn_for_benchmark( + queries: list[str], + base_history: list[AgentHistory], + cm_config: ContextManagerConfig, + max_steps: int = 20, + system_prompt: str = BENCHMARK_SYSTEM_PROMPT, +): + conversation_history = list(base_history) + results = [] + + shared_cm = None + if cm_config and cm_config.enabled: + shared_cm = ContextManager(config=cm_config, max_steps=max_steps) + + initial_tokens = estimate_tokens_text(history_to_text(conversation_history)) + + # Track per-step actual input tokens for accurate token reduction + step_input_tokens = [] + + for query in queries: + agent_run_info = build_agent_run_info_with_custom_prompt( + query, + system_prompt, + conversation_history, + max_steps=max_steps, + context_manager_config=cm_config, + ) + + if shared_cm is not None: + agent_run_info.context_manager = shared_cm + + result = await run_agent_with_tracking(agent_run_info, debug=False) + results.append(result) + + # Collect actual input token count from the last step metrics + if shared_cm is not None: + tc = shared_cm.get_token_counts() + step_input_tokens.append(tc) + + conversation_history.append(AgentHistory(role="user", content=query)) + conversation_history.append( + AgentHistory(role="assistant", content=result.final_answer) + ) + + final_tokens = estimate_tokens_text(history_to_text(conversation_history)) + + cm_stats = None + cm_token_counts = None + cm_summary = None + if shared_cm is not None: + cm_stats = shared_cm.get_all_compression_stats() + cm_token_counts = shared_cm.get_token_counts() + cm_summary = shared_cm.export_summary() + + return { + "results": results, + "conversation_history": conversation_history, + "shared_cm": shared_cm, + "initial_tokens": initial_tokens, + "final_tokens": final_tokens, + "cm_stats": cm_stats, + "cm_token_counts": cm_token_counts, + "cm_summary": cm_summary, + "step_input_tokens": step_input_tokens, + } + + +def build_precompressed_history( + frozen_history: list[AgentHistory], + cm_summary: dict, +) -> list[AgentHistory]: + """Build a pre-compressed history from the compression snapshot. + + Replaces the compressed prefix pairs with a single user message containing + the summary text, then appends the retained tail pairs verbatim. This + mirrors the actual message structure produced by compress_if_needed: + + SummaryTaskStep.to_messages() → [ChatMessage(role=USER, summary)] + followed by retained tail steps → [TaskStep, ActionStep, ...] + + There is NO assistant message after the summary — the model sees the + summary as a user message, followed directly by the next retained step. + + Args: + frozen_history: The original uncompressed conversation history. + cm_summary: The export_summary() dict from the compressed run's + ContextManager, containing summary text and boundary info. + + Returns: + A new AgentHistory list that mirrors the compressed context structure. + """ + boundary = cm_summary.get("compression_boundary", {}) + compressed_pairs = boundary.get("previous_compressed_pairs", 0) + + # Each pair = 2 AgentHistory entries (user + assistant) + compressed_entries = compressed_pairs * 2 + + summary_text = cm_summary.get("previous_summary") or "" + + # If no compression happened, return original history unchanged + if not summary_text or compressed_entries == 0: + return list(frozen_history) + + # Build pre-compressed history: + # 1. Summary as a single USER message (matching SummaryTaskStep.to_messages) + # No paired assistant message — the model sees summary then next retained step + precompressed = [ + AgentHistory( + role="user", + content=f"Summary of earlier steps in this task:\n{summary_text}", + ), + ] + + # 2. Retained tail pairs (everything after the compressed prefix) + if compressed_entries < len(frozen_history): + precompressed.extend(frozen_history[compressed_entries:]) + + return precompressed + + +async def run_probe_questions( + probes: list[dict], + precompressed_history: list[AgentHistory], + max_steps: int = 20, + system_prompt: str = BENCHMARK_SYSTEM_PROMPT, +): + """Run probe questions against a pre-compressed history snapshot. + + Each probe runs independently with compression DISABLED, because the + history has already been pre-compressed (compressed prefix replaced with + summary text, retained tail kept verbatim). This avoids redundant LLM + compression calls — the compression was done once in the compressed run, + and all probes reuse that result. + + Per CLAUDE.md rules: + - Each probe uses a deep-copied frozen snapshot + - Probes see compressed context (summary + retained tail) + - No compression triggered during probe phase + - Probes are fully independent, no shared state + """ + probe_results = [] + no_compression_config = ContextManagerConfig(enabled=False, token_threshold=10**9) + + for probe in probes: + question = probe["question"] + + # Each probe gets its own deep copy — fully independent + probe_history = copy.deepcopy(precompressed_history) + + agent_run_info = build_agent_run_info_with_custom_prompt( + question, + system_prompt, + probe_history, + max_steps=max_steps, + context_manager_config=no_compression_config, + ) + + result = await run_agent_with_tracking(agent_run_info, debug=False) + eval_result = eval_text(result.final_answer, probe) + + probe_results.append( + { + "question": question, + "answer": result.final_answer, + "passed": eval_result.passed, + "score": eval_result.score, + "details": eval_result.details, + } + ) + + return probe_results + + +async def run_baseline_probes( + probes: list[dict], + frozen_history: list[AgentHistory], + max_steps: int = 20, + system_prompt: str = BENCHMARK_SYSTEM_PROMPT, +): + """Run probe questions against full uncompressed history (baseline). + + This measures the ceiling: what can the agent answer when it sees + the complete history. probe_retention = compressed_score / baseline_score. + """ + probe_results = [] + baseline_config = ContextManagerConfig(enabled=False, token_threshold=10**9) + + for probe in probes: + question = probe["question"] + probe_history = copy.deepcopy(frozen_history) + + agent_run_info = build_agent_run_info_with_custom_prompt( + question, + system_prompt, + probe_history, + max_steps=max_steps, + context_manager_config=baseline_config, + ) + + result = await run_agent_with_tracking(agent_run_info, debug=False) + eval_result = eval_text(result.final_answer, probe) + + probe_results.append( + { + "question": question, + "answer": result.final_answer, + "passed": eval_result.passed, + "score": eval_result.score, + "details": eval_result.details, + } + ) + + return probe_results + + +def eval_summary_inspection(summary: dict, checks: list[dict]) -> list[dict]: + """Static Compression Inspection — check if the compressed summary + retains key information (user preferences, file names, plans, tool results). + + Uses dedicated summary_checks when available, NOT probe must_contain + (which has different semantics — probe keywords are for agent answers, + summary keywords are for what the compressor chose to preserve). + """ + results = [] + + prev_summary = summary.get("previous_summary") or "" + curr_summary = summary.get("current_summary") or "" + combined = prev_summary + "\n" + curr_summary + + for check in checks: + eval_result = eval_text(combined, check) + results.append( + { + "check": check, + "passed": eval_result.passed, + "score": eval_result.score, + "details": eval_result.details, + } + ) + + return results + + +def eval_task_outputs(case: dict, run_outputs: list): + eval_results = [] + + for check in case.get("task_checks", []): + turn_idx = check["turn"] - 1 + if turn_idx >= len(run_outputs): + continue + + answer = run_outputs[turn_idx].final_answer + r = eval_text(answer, check) + + eval_results.append( + { + "turn": check["turn"], + "answer": answer, + "passed": r.passed, + "score": r.score, + "details": r.details, + } + ) + + return eval_results + + +def _resolve_compressed_config(case: dict, use_default_prompts: bool = False) -> ContextManagerConfig: + """Build compressed config from case definition, with sensible defaults. + + By default uses the benchmark-optimized custom summary schema and prompts. + Set use_default_prompts=True to fall back to the original ContextManager defaults. + """ + case_cfg = case.get("compressed_config", {}) + kwargs = dict( + enabled=True, + token_threshold=case_cfg.get("token_threshold", 3600), + keep_recent_pairs=case_cfg.get("keep_recent_pairs", 1), + keep_recent_steps=case_cfg.get("keep_recent_steps", 4), + max_observation_length=case_cfg.get("max_observation_length", 20000), + ) + if not use_default_prompts: + kwargs.update( + summary_json_schema=BENCHMARK_SUMMARY_SCHEMA, + summary_system_prompt=BENCHMARK_SUMMARY_SYSTEM_PROMPT, + incremental_summary_system_prompt=BENCHMARK_INCREMENTAL_SUMMARY_SYSTEM_PROMPT, + ) + return ContextManagerConfig(**kwargs) + + +async def run_one_case(case_dir: str, use_default_prompts: bool = False): + """Load and run a single benchmark case from its directory. + + Each case directory contains: + - case.json: queries, probes, summary_checks, task_checks, compressed_config + - history.json: conversation history + + Args: + case_dir: Absolute or relative path to the case directory. + + Returns: + Report dict for this case. + """ + case_path = os.path.join(case_dir, "case.json") + with open(case_path, "r", encoding="utf-8") as f: + case = json.load(f) + + # Resolve history_file relative to the case directory; + # defaults to "history.json" in the same directory if not specified. + history_relpath = case.get("history_file", "history.json") + history_abspath = os.path.join(case_dir, history_relpath) + + base_history = parse_conversation_to_history(history_abspath) + + baseline_config = ContextManagerConfig( + enabled=False, + token_threshold=10**9, + keep_recent_pairs=1, + ) + + # P5: Allow per-case config override + compressed_config = _resolve_compressed_config(case, use_default_prompts=use_default_prompts) + + print(f"\n===== CASE: {case['id']} =====") + + baseline = await run_multi_turn_for_benchmark( + queries=case["queries"], + base_history=base_history, + cm_config=baseline_config, + ) + + compressed = await run_multi_turn_for_benchmark( + queries=case["queries"], + base_history=base_history, + cm_config=compressed_config, + ) + + baseline_task_eval = eval_task_outputs(case, baseline["results"]) + compressed_task_eval = eval_task_outputs(case, compressed["results"]) + # P1: Baseline probe — agent sees full uncompressed history + # Same frozen_history, but with compression disabled, so the agent sees + # the complete unmodified context. This establishes the ceiling for + # probe_retention = compressed_probe_score / baseline_probe_score. + baseline_probe_eval = await run_baseline_probes( + probes=case["probes"], + frozen_history=compressed["conversation_history"], + max_steps=20, + ) + + # P0: Compressed probe — agent sees pre-compressed context + # Build the pre-compressed history ONCE using the summary from the + # compressed run's ContextManager, then run each probe independently + # against it with compression disabled. This avoids redundant LLM calls + # (compression was already done in the compressed multi-turn run). + precompressed_history = build_precompressed_history( + frozen_history=compressed["conversation_history"], + cm_summary=compressed["cm_summary"] or {}, + ) + compressed_probe_eval = await run_probe_questions( + probes=case["probes"], + precompressed_history=precompressed_history, + ) + + # P3: Summary inspection uses dedicated summary_checks, not probe must_contain + summary_inspection = [] + if compressed.get("cm_summary"): + summary_checks = case.get("summary_checks", []) + if summary_checks: + summary_inspection = eval_summary_inspection( + compressed["cm_summary"], summary_checks + ) + + baseline_task_score = sum(x["score"] for x in baseline_task_eval) / max( + len(baseline_task_eval), 1 + ) + + compressed_task_score = sum(x["score"] for x in compressed_task_eval) / max( + len(compressed_task_eval), 1 + ) + + baseline_probe_score = sum(x["score"] for x in baseline_probe_eval) / max( + len(baseline_probe_eval), 1 + ) + + compressed_probe_score = sum(x["score"] for x in compressed_probe_eval) / max( + len(compressed_probe_eval), 1 + ) + + summary_score = ( + sum(x["score"] for x in summary_inspection) / max(len(summary_inspection), 1) + if summary_inspection + else None + ) + + task_success_retention = ( + compressed_task_score / baseline_task_score + if baseline_task_score > 0 + else 0.0 + ) + + probe_retention = ( + compressed_probe_score / baseline_probe_score + if baseline_probe_score > 0 + else 0.0 + ) + + # P2: Token reduction from actual input token counts + # Use the last step's token counts (final compressed vs uncompressed state) + token_reduction = 0.0 + if compressed.get("step_input_tokens") and compressed["step_input_tokens"]: + last_tc = compressed["step_input_tokens"][-1] + if last_tc and last_tc.get("last_uncompressed") is not None: + unc = last_tc["last_uncompressed"] or 1 + comp = last_tc["last_compressed"] or 0 + if unc > 0: + token_reduction = 1 - comp / unc + # Fallback to text-based estimation + if token_reduction == 0.0: + token_reduction = 1 - ( + compressed["final_tokens"] / max(baseline["final_tokens"], 1) + ) + baseline_failed = baseline_task_score == 0 + + # Compute real main-LLM input token totals + baseline_real_input = sum(r.total_input_tokens for r in baseline["results"]) + compressed_real_input = sum(r.total_input_tokens for r in compressed["results"]) + + # Compression cost: tokens spent on compression LLM calls + compression_cost = 0 + if compressed.get("cm_stats"): + compression_cost = ( + compressed["cm_stats"].get("total_input_tokens", 0) + + compressed["cm_stats"].get("total_output_tokens", 0) + ) + + # Net token reduction = gross savings - compression cost + gross_input_savings = baseline_real_input - compressed_real_input + net_input_savings = gross_input_savings - compression_cost + net_token_reduction = ( + net_input_savings / max(baseline_real_input, 1) + if baseline_real_input > 0 + else 0.0 + ) + + report = { + "case_id": case["id"], + "baseline_failed": baseline_failed, + "baseline": { + "task_score": baseline_task_score, + "probe_score": baseline_probe_score, + "final_tokens": baseline["final_tokens"], + "real_input_tokens": baseline_real_input, + }, + "compressed": { + "task_score": compressed_task_score, + "probe_score": compressed_probe_score, + "final_tokens": compressed["final_tokens"], + "cm_stats": compressed["cm_stats"], + "cm_token_counts": compressed["cm_token_counts"], + "cm_summary": compressed["cm_summary"], + "real_input_tokens": compressed_real_input, + }, + "metrics": { + "task_success_retention": task_success_retention, + "probe_retention": probe_retention, + "token_reduction": token_reduction, + "net_token_reduction": net_token_reduction, + "compression_cost_tokens": compression_cost, + "summary_score": summary_score, + }, + "task_eval": compressed_task_eval, + "probe_eval": { + "baseline": baseline_probe_eval, + "compressed": compressed_probe_eval, + }, + "summary_inspection": summary_inspection, + } + + print(json.dumps(report, ensure_ascii=False, indent=2, default=str)) + return report + + +async def main(case_names: list[str] = None, use_default_prompts: bool = False): + # Discover cases: use specified names if provided, otherwise find all cases under ./cases/*/case.json + if case_names: + case_dirs = [os.path.join("./cases", name) for name in case_names] + else: + case_dirs = sorted(glob.glob("./cases/*/case.json")) + case_dirs = [os.path.dirname(p) for p in case_dirs] + + if not case_dirs: + print("No benchmark cases found under ./cases/*/case.json") + return + + print(f"Found {len(case_dirs)} case(s): {[os.path.basename(d) for d in case_dirs]}") + + # Output directory for reports + os.makedirs("./reports", exist_ok=True) + + reports = [] + for case_dir in case_dirs: + report = await run_one_case(case_dir, use_default_prompts=use_default_prompts) + reports.append(report) + + # Write per-case report + case_id = report["case_id"] + per_case_path = os.path.join("./reports", f"{case_id}.json") + with open(per_case_path, "w", encoding="utf-8") as f: + json.dump(report, f, ensure_ascii=False, indent=2, default=str) + print(f" Report saved to {per_case_path}") + + # Exclude cases where baseline itself failed + valid_reports = [r for r in reports if not r.get("baseline_failed")] + excluded_ids = [r["case_id"] for r in reports if r.get("baseline_failed")] + if excluded_ids: + print(f"\n Excluded from average (baseline failed): {excluded_ids}") + # Write summary across all cases + summary = { + "total_cases": len(reports), + "excluded_cases": len(reports) - len(valid_reports), + "metrics": { + "avg_task_success_retention": sum( + r["metrics"]["task_success_retention"] for r in valid_reports + ) / max(len(valid_reports), 1), + "avg_probe_retention": sum( + r["metrics"]["probe_retention"] for r in valid_reports + ) / max(len(valid_reports), 1), + "avg_token_reduction": sum( + r["metrics"]["token_reduction"] for r in valid_reports + ) / max(len(valid_reports), 1), + "avg_net_token_reduction": sum( + r["metrics"]["net_token_reduction"] for r in valid_reports + ) / max(len(valid_reports), 1), + "avg_compression_cost_tokens": sum( + r["metrics"]["compression_cost_tokens"] for r in valid_reports + ) / max(len(valid_reports), 1), + "per_case": { + r["case_id"]: r["metrics"] for r in reports + }, + }, + } + summary_path = "./reports/summary.json" + with open(summary_path, "w", encoding="utf-8") as f: + json.dump(summary, f, ensure_ascii=False, indent=2, default=str) + + print(f"\nBenchmark finished. Summary saved to {summary_path}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Run Agent Context Compression Benchmark") + parser.add_argument( + "--cases",nargs="+",default=None, + help="Specific case names to run (e.g. --cases example_infra algotithm_data)." + "if omitted, run all cases under .cases/." + ) + parser.add_argument( + "--default-summary", action="store_true", default=False, + help="Use the original ContextManager summary defaults instead of the benchmark-optimized " + "custom schema (leaner 7-field, 800-word cap, merge-condense incremental updates)." + ) + args = parser.parse_args() asyncio.run(main(case_names=args.cases, use_default_prompts=args.default_summary)) \ No newline at end of file diff --git a/sdk/benchmark/paths.py b/sdk/benchmark/paths.py index 88faf31f2..4bb67ade8 100644 --- a/sdk/benchmark/paths.py +++ b/sdk/benchmark/paths.py @@ -1,61 +1,61 @@ -# -*- coding: utf-8 -*- -"""Robust path resolution for benchmark scripts. - -Finds the project root by searching upward for a .git entry (directory -or file), then derives SDK_DIR and BACKEND_PATH from it. This makes -path setup resilient to file relocation within the project tree and to -git worktrees (which store a .git file rather than directory at root). -""" -import os -import sys - - -def _find_project_root(start: str = None) -> str: - """Walk upward from *start* until a .git entry is found. - - Accepts ``.git`` as either a directory (normal checkout) or a file - (git worktree, where ``.git`` is a pointer file to the gitdir). - """ - current = os.path.abspath(start or os.path.dirname(__file__)) - while True: - if os.path.exists(os.path.join(current, ".git")): - return current - parent = os.path.dirname(current) - if parent == current: - raise RuntimeError( - f"Could not find project root (.git) starting from {start or __file__}" - ) - current = parent - - -def setup_paths() -> dict: - """Resolve and inject project paths into sys.path. - - Returns a dict with resolved paths: - project_root, sdk_dir, backend_dir - - Adds the following to sys.path (idempotent): - - sdk_dir (for ``from nexent import ...``) - - project_root (for ``from backend.utils import ...``) - - backend_dir (for ``from utils.prompt_template_utils import ...``) - """ - project_root = _find_project_root() - sdk_dir = os.path.join(project_root, "sdk") - backend_dir = os.path.join(project_root, "backend") - - for p in (sdk_dir, project_root, backend_dir): - if p not in sys.path: - sys.path.insert(0, p) - - return { - "project_root": project_root, - "sdk_dir": sdk_dir, - "backend_dir": backend_dir, - } - - -# Convenience: resolve on import so callers can do `from paths import PROJECT_ROOT` -_resolved = setup_paths() -PROJECT_ROOT = _resolved["project_root"] -SDK_DIR = _resolved["sdk_dir"] +# -*- coding: utf-8 -*- +"""Robust path resolution for benchmark scripts. + +Finds the project root by searching upward for a .git entry (directory +or file), then derives SDK_DIR and BACKEND_PATH from it. This makes +path setup resilient to file relocation within the project tree and to +git worktrees (which store a .git file rather than directory at root). +""" +import os +import sys + + +def _find_project_root(start: str = None) -> str: + """Walk upward from *start* until a .git entry is found. + + Accepts ``.git`` as either a directory (normal checkout) or a file + (git worktree, where ``.git`` is a pointer file to the gitdir). + """ + current = os.path.abspath(start or os.path.dirname(__file__)) + while True: + if os.path.exists(os.path.join(current, ".git")): + return current + parent = os.path.dirname(current) + if parent == current: + raise RuntimeError( + f"Could not find project root (.git) starting from {start or __file__}" + ) + current = parent + + +def setup_paths() -> dict: + """Resolve and inject project paths into sys.path. + + Returns a dict with resolved paths: + project_root, sdk_dir, backend_dir + + Adds the following to sys.path (idempotent): + - sdk_dir (for ``from nexent import ...``) + - project_root (for ``from backend.utils import ...``) + - backend_dir (for ``from utils.prompt_template_utils import ...``) + """ + project_root = _find_project_root() + sdk_dir = os.path.join(project_root, "sdk") + backend_dir = os.path.join(project_root, "backend") + + for p in (sdk_dir, project_root, backend_dir): + if p not in sys.path: + sys.path.insert(0, p) + + return { + "project_root": project_root, + "sdk_dir": sdk_dir, + "backend_dir": backend_dir, + } + + +# Convenience: resolve on import so callers can do `from paths import PROJECT_ROOT` +_resolved = setup_paths() +PROJECT_ROOT = _resolved["project_root"] +SDK_DIR = _resolved["sdk_dir"] BACKEND_DIR = _resolved["backend_dir"] \ No newline at end of file diff --git a/sdk/nexent/core/agents/agent_context/manager.py b/sdk/nexent/core/agents/agent_context/manager.py index 483b3b6d3..8b6229659 100644 --- a/sdk/nexent/core/agents/agent_context/manager.py +++ b/sdk/nexent/core/agents/agent_context/manager.py @@ -27,6 +27,13 @@ msg_char_count, msg_token_count, ) +from nexent.monitor import ( + get_monitoring_manager, + OPENINFERENCE_SPAN_KIND_CHAIN, + OPENINFERENCE_SPAN_KIND_LLM, + OPENINFERENCE_INPUT_VALUE, + OPENINFERENCE_OUTPUT_VALUE, +) from .budget import ( action_content, @@ -54,6 +61,18 @@ logger = logging.getLogger("agent_context") +_handlers_registered = False + + +def _ensure_handlers_registered(): + """Lazily register all context item handlers (idempotent).""" + global _handlers_registered + if not _handlers_registered: + from ..context.handlers import register_all + register_all() + _handlers_registered = True + + class ContextManager: def __init__(self, config: Optional[ContextManagerConfig] = None, max_steps: Optional[int] = None): @@ -153,37 +172,151 @@ def compress_if_needed( current_run_start_idx, context_overhead_tokens: int = 0, ) -> List[ChatMessage]: - if not self.config.enabled: - return original_messages + monitoring_manager = get_monitoring_manager() + with monitoring_manager.trace_operation( + "context.compress_if_needed", + OPENINFERENCE_SPAN_KIND_CHAIN, + **{ + "context.context_overhead_tokens": context_overhead_tokens, + "context.current_run_start_idx": current_run_start_idx, + OPENINFERENCE_INPUT_VALUE: { + "memory_steps": len(memory.steps) if hasattr(memory, 'steps') else 0, + "original_message_count": len(original_messages), + "current_run_start_idx": current_run_start_idx, + }, + }, + ): + if not self.config.enabled: + if monitoring_manager.is_enabled: + monitoring_manager.set_openinference_output({"decision": "disabled", "message_count": len(original_messages)}) + return original_messages + + soft_input_budget_tokens = self._soft_input_budget_tokens() + hard_input_budget_tokens = self._hard_input_budget_tokens() + soft_history_budget_tokens = max(0, soft_input_budget_tokens - context_overhead_tokens) + hard_history_budget_tokens = max(0, hard_input_budget_tokens - context_overhead_tokens) + + if estimate_tokens(memory, self.config.chars_per_token) <= soft_history_budget_tokens: + self._last_uncompressed_token_count = self._msg_token_count(original_messages) + self._last_compressed_token_count = self._last_uncompressed_token_count + + if monitoring_manager.is_enabled: + monitoring_manager.add_span_event( + "context.compress.skipped", + { + "context.reason": "under_budget", + "context.estimated_tokens": estimate_tokens(memory, self.config.chars_per_token), + "context.soft_budget": soft_history_budget_tokens, + }, + ) + monitoring_manager.set_openinference_output({"decision": "skipped", "message_count": len(original_messages), "reason": "under_budget"}) - soft_input_budget_tokens = self._soft_input_budget_tokens() - hard_input_budget_tokens = self._hard_input_budget_tokens() - soft_history_budget_tokens = max(0, soft_input_budget_tokens - context_overhead_tokens) - hard_history_budget_tokens = max(0, hard_input_budget_tokens - context_overhead_tokens) + return original_messages - if estimate_tokens(memory, self.config.chars_per_token) <= soft_history_budget_tokens: - self._last_uncompressed_token_count = self._msg_token_count(original_messages) - self._last_compressed_token_count = self._last_uncompressed_token_count - return original_messages + with self._lock: + # Run detection + if (self._last_run_start_idx is not None + and current_run_start_idx != self._last_run_start_idx): + self._current_summary_cache = None + self._last_run_start_idx = current_run_start_idx + + if self._effective_tokens(memory, current_run_start_idx) <= soft_history_budget_tokens: + # Stable-phase bypass + self._step_local_log.clear() + + prev_steps = memory.steps[:current_run_start_idx] + curr_steps = memory.steps[current_run_start_idx:] + + prev_summary_step = None + prev_tail_steps = list(prev_steps) + prev_pairs = extract_pairs(prev_steps) + if prev_pairs: + is_valid, covered_idx = is_prev_cache_valid(prev_pairs, self._previous_summary_cache) + if is_valid: + prev_summary_step = SummaryTaskStep( + task=self._previous_summary_cache.summary_text + ) + uncovered = prev_pairs[covered_idx:] + prev_tail_steps = self._renderer.pairs_to_steps(uncovered) + + curr_kept_steps = list(curr_steps) + if curr_steps: + curr_task = curr_steps[0] if isinstance(curr_steps[0], TaskStep) else None + curr_action_steps = [s for s in curr_steps if isinstance(s, ActionStep)] + if curr_action_steps: + is_valid, covered_idx = is_curr_cache_valid(curr_action_steps, self._current_summary_cache) + if is_valid: + uncovered = curr_action_steps[covered_idx:] + curr_kept_steps = ( + ([curr_task] if curr_task else []) + + [SummaryTaskStep(task=self._current_summary_cache.summary_text)] + + list(uncovered) + ) + + record = CompressionCallRecord( + call_type="stable_bypass", cache_hit=True, + details={"reason": "stable_period_effective_under_threshold"}, + ) + self.compression_calls_log.append(record) + self._step_local_log.append(record) + + compressed_msgs = self._renderer.build_messages( + memory, prev_summary_step, prev_tail_steps, curr_kept_steps + ) + self._last_uncompressed_token_count = self._msg_token_count(original_messages) + self._last_compressed_token_count = self._msg_token_count(compressed_msgs) + if monitoring_manager.is_enabled: + monitoring_manager.set_openinference_output({"decision": "stable_bypass", "message_count": len(compressed_msgs), "cache_hit": True}) + return compressed_msgs - with self._lock: - # Run detection - if (self._last_run_start_idx is not None - and current_run_start_idx != self._last_run_start_idx): - self._current_summary_cache = None - self._last_run_start_idx = current_run_start_idx - - if self._effective_tokens(memory, current_run_start_idx) <= soft_history_budget_tokens: - # Stable-phase bypass self._step_local_log.clear() + self._last_uncompressed_token_count = self._msg_token_count(original_messages) prev_steps = memory.steps[:current_run_start_idx] curr_steps = memory.steps[current_run_start_idx:] - prev_summary_step = None - prev_tail_steps = list(prev_steps) + prev_tokens = self._effective_prev_tokens(prev_steps) + curr_tokens = self._effective_curr_tokens(curr_steps) + + compress_prev = prev_tokens > soft_history_budget_tokens * 0.6 + compress_curr = curr_tokens > soft_history_budget_tokens * 0.4 + + total_effective_tokens = prev_tokens + curr_tokens + context_overhead_tokens + if compress_prev or compress_curr: + logger.info( + f"Context compression triggered: total_tokens={total_effective_tokens}, " + f"soft_budget={soft_input_budget_tokens}, " + f"hard_budget={hard_input_budget_tokens}, " + f"context_overhead_tokens={context_overhead_tokens}, " + f"prev_tokens={prev_tokens} (compress={compress_prev}), " + f"curr_tokens={curr_tokens} (compress={compress_curr})" + ) + + # --------------- Previous phase --------------- + prev_summary_step: Optional[SummaryTaskStep] = None + prev_tail_steps: List[MemoryStep] = list(prev_steps) prev_pairs = extract_pairs(prev_steps) - if prev_pairs: + + if compress_prev and prev_pairs: + keep_n = min(self.config.keep_recent_pairs, len(prev_pairs)) + pairs_to_compress = prev_pairs[:-keep_n] if keep_n > 0 else prev_pairs + pairs_to_keep = prev_pairs[-keep_n:] if keep_n > 0 else [] + if pairs_to_compress: + result = self._prev_compressor.compress( + pairs_to_compress, self._previous_summary_cache, model + ) + summary_text = result.summary_text + if summary_text: + if "[CONTEXT COMPACTION" in summary_text: + prev_summary_step = SummaryTaskStep(task=summary_text, prefix="Context fallback, Truncated raw history:") + else: + prev_summary_step = SummaryTaskStep(task=summary_text) + prev_tail_steps = self._renderer.pairs_to_steps(pairs_to_keep) + if result.new_cache: + self._previous_summary_cache = result.new_cache + self.compression_calls_log.extend(result.records) + self._step_local_log.extend(result.records) + elif prev_pairs: is_valid, covered_idx = is_prev_cache_valid(prev_pairs, self._previous_summary_cache) if is_valid: prev_summary_step = SummaryTaskStep( @@ -192,11 +325,49 @@ def compress_if_needed( uncovered = prev_pairs[covered_idx:] prev_tail_steps = self._renderer.pairs_to_steps(uncovered) - curr_kept_steps = list(curr_steps) + # --------------- Current phase --------------- + curr_kept_steps: List[MemoryStep] = list(curr_steps) + if curr_steps: curr_task = curr_steps[0] if isinstance(curr_steps[0], TaskStep) else None curr_action_steps = [s for s in curr_steps if isinstance(s, ActionStep)] - if curr_action_steps: + + if compress_curr and curr_action_steps: + keep_n = min(self.config.keep_recent_steps, len(curr_action_steps)) + if keep_n > 0 and keep_n < len(curr_action_steps): + boundary = curr_action_steps[-keep_n] + prev_a = curr_action_steps[-keep_n - 1] + if (getattr(boundary, "observations", None) is not None + and getattr(prev_a, "tool_calls", None) is not None): + keep_n += 1 + + actions_to_compress = ( + curr_action_steps[:-keep_n] if keep_n > 0 else list(curr_action_steps) + ) + actions_to_keep = ( + curr_action_steps[-keep_n:] if keep_n > 0 else [] + ) + if actions_to_compress: + result = self._curr_compressor.compress( + curr_task, actions_to_compress, + self._current_summary_cache, model, + ) + curr_summary_text = result.summary_text + if curr_summary_text: + if "[CONTEXT COMPACTION" in curr_summary_text: + curr_summary_step = SummaryTaskStep(task=curr_summary_text, prefix="Truncated recent action steps:") + else: + curr_summary_step = SummaryTaskStep(task=curr_summary_text) + curr_kept_steps = ( + ([curr_task] if curr_task else []) + + [curr_summary_step] + + list(actions_to_keep) + ) + if result.new_cache: + self._current_summary_cache = result.new_cache + self.compression_calls_log.extend(result.records) + self._step_local_log.extend(result.records) + elif curr_action_steps: is_valid, covered_idx = is_curr_cache_valid(curr_action_steps, self._current_summary_cache) if is_valid: uncovered = curr_action_steps[covered_idx:] @@ -206,140 +377,25 @@ def compress_if_needed( + list(uncovered) ) - record = CompressionCallRecord( - call_type="stable_bypass", cache_hit=True, - details={"reason": "stable_period_effective_under_threshold"}, - ) - self.compression_calls_log.append(record) - self._step_local_log.append(record) - - compressed_msgs = self._renderer.build_messages( + final_messages = self._renderer.build_messages( memory, prev_summary_step, prev_tail_steps, curr_kept_steps ) - self._last_uncompressed_token_count = self._msg_token_count(original_messages) - self._last_compressed_token_count = self._msg_token_count(compressed_msgs) - return compressed_msgs - - self._step_local_log.clear() - self._last_uncompressed_token_count = self._msg_token_count(original_messages) - - prev_steps = memory.steps[:current_run_start_idx] - curr_steps = memory.steps[current_run_start_idx:] - - prev_tokens = self._effective_prev_tokens(prev_steps) - curr_tokens = self._effective_curr_tokens(curr_steps) - - compress_prev = prev_tokens > soft_history_budget_tokens * 0.6 - compress_curr = curr_tokens > soft_history_budget_tokens * 0.4 - - total_effective_tokens = prev_tokens + curr_tokens + context_overhead_tokens - if compress_prev or compress_curr: - logger.info( - f"Context compression triggered: total_tokens={total_effective_tokens}, " - f"soft_budget={soft_input_budget_tokens}, " - f"hard_budget={hard_input_budget_tokens}, " - f"context_overhead_tokens={context_overhead_tokens}, " - f"prev_tokens={prev_tokens} (compress={compress_prev}), " - f"curr_tokens={curr_tokens} (compress={compress_curr})" - ) - - # --------------- Previous phase --------------- - prev_summary_step: Optional[SummaryTaskStep] = None - prev_tail_steps: List[MemoryStep] = list(prev_steps) - prev_pairs = extract_pairs(prev_steps) - - if compress_prev and prev_pairs: - keep_n = min(self.config.keep_recent_pairs, len(prev_pairs)) - pairs_to_compress = prev_pairs[:-keep_n] if keep_n > 0 else prev_pairs - pairs_to_keep = prev_pairs[-keep_n:] if keep_n > 0 else [] - if pairs_to_compress: - result = self._prev_compressor.compress( - pairs_to_compress, self._previous_summary_cache, model - ) - summary_text = result.summary_text - if summary_text: - if "[CONTEXT COMPACTION" in summary_text: - prev_summary_step = SummaryTaskStep(task=summary_text, prefix="Context fallback, Truncated raw history:") - else: - prev_summary_step = SummaryTaskStep(task=summary_text) - prev_tail_steps = self._renderer.pairs_to_steps(pairs_to_keep) - if result.new_cache: - self._previous_summary_cache = result.new_cache - self.compression_calls_log.extend(result.records) - self._step_local_log.extend(result.records) - elif prev_pairs: - is_valid, covered_idx = is_prev_cache_valid(prev_pairs, self._previous_summary_cache) - if is_valid: - prev_summary_step = SummaryTaskStep( - task=self._previous_summary_cache.summary_text - ) - uncovered = prev_pairs[covered_idx:] - prev_tail_steps = self._renderer.pairs_to_steps(uncovered) - - # --------------- Current phase --------------- - curr_kept_steps: List[MemoryStep] = list(curr_steps) - - if curr_steps: - curr_task = curr_steps[0] if isinstance(curr_steps[0], TaskStep) else None - curr_action_steps = [s for s in curr_steps if isinstance(s, ActionStep)] - - if compress_curr and curr_action_steps: - keep_n = min(self.config.keep_recent_steps, len(curr_action_steps)) - if keep_n > 0 and keep_n < len(curr_action_steps): - boundary = curr_action_steps[-keep_n] - prev_a = curr_action_steps[-keep_n - 1] - if (getattr(boundary, "observations", None) is not None - and getattr(prev_a, "tool_calls", None) is not None): - keep_n += 1 - - actions_to_compress = ( - curr_action_steps[:-keep_n] if keep_n > 0 else list(curr_action_steps) - ) - actions_to_keep = ( - curr_action_steps[-keep_n:] if keep_n > 0 else [] + final_tokens = self._msg_token_count(final_messages) + self._last_compressed_token_count = final_tokens + if final_tokens > hard_history_budget_tokens: + logger.warning( + f"Still exceeds hard input budget after compression: {final_tokens} > {hard_history_budget_tokens}. " + f"Consider reducing keep_recent_pairs ({self.config.keep_recent_pairs}) " + f"or keep_recent_steps({self.config.keep_recent_steps})" ) - if actions_to_compress: - result = self._curr_compressor.compress( - curr_task, actions_to_compress, - self._current_summary_cache, model, - ) - curr_summary_text = result.summary_text - if curr_summary_text: - if "[CONTEXT COMPACTION" in curr_summary_text: - curr_summary_step = SummaryTaskStep(task=curr_summary_text, prefix="Truncated recent action steps:") - else: - curr_summary_step = SummaryTaskStep(task=curr_summary_text) - curr_kept_steps = ( - ([curr_task] if curr_task else []) - + [curr_summary_step] - + list(actions_to_keep) - ) - if result.new_cache: - self._current_summary_cache = result.new_cache - self.compression_calls_log.extend(result.records) - self._step_local_log.extend(result.records) - elif curr_action_steps: - is_valid, covered_idx = is_curr_cache_valid(curr_action_steps, self._current_summary_cache) - if is_valid: - uncovered = curr_action_steps[covered_idx:] - curr_kept_steps = ( - ([curr_task] if curr_task else []) - + [SummaryTaskStep(task=self._current_summary_cache.summary_text)] - + list(uncovered) - ) - - final_messages = self._renderer.build_messages( - memory, prev_summary_step, prev_tail_steps, curr_kept_steps - ) - final_tokens = self._msg_token_count(final_messages) - self._last_compressed_token_count = final_tokens - if final_tokens > hard_history_budget_tokens: - logger.warning( - f"Still exceeds hard input budget after compression: {final_tokens} > {hard_history_budget_tokens}. " - f"Consider reducing keep_recent_pairs ({self.config.keep_recent_pairs}) " - f"or keep_recent_steps({self.config.keep_recent_steps})" - ) - return final_messages + if monitoring_manager.is_enabled: + monitoring_manager.set_openinference_output({ + "decision": "compressed", + "message_count": len(final_messages), + "uncompressed_tokens": self._last_uncompressed_token_count, + "compressed_tokens": self._last_compressed_token_count, + }) + return final_messages # ============================================================ # Token Estimation @@ -446,16 +502,25 @@ def prepare_run_context( system_prompt=stable_text or fallback_system_prompt ) source_components = tuple(self._component_source(components)) + if source_components: + budget = self._calculate_component_budget() + strategy = self._get_strategy() + selected_components = tuple(strategy.select_components( + source_components, budget, self.config.component_budgets + )) + else: + selected_components = () + selected_component_types = tuple( str(getattr(component, "component_type", "unknown")) - for component in source_components + for component in selected_components ) return ManagedRunContext( component_messages=tuple(component_messages), stable_messages=tuple(stable_messages), dynamic_messages=tuple(dynamic_messages), selected_component_types=selected_component_types, - components=source_components, + selected_components=selected_components, ) def assemble_final_context( @@ -469,67 +534,140 @@ def assemble_final_context( task: Optional[str] = None, final_answer_templates: Optional[Dict[str, Any]] = None, run_context: Optional[ManagedRunContext] = None, + conversation_id: Optional[int] = None, ) -> "FinalContext": - if run_context is None: - run_context = self.prepare_run_context(memory, fallback_system_prompt="") - - tools = self._canonical_tools(tools or ()) - purpose_stable, purpose_dynamic = self._purpose_messages( - purpose=purpose, - task=task, - final_answer_templates=final_answer_templates, - ) + monitoring_manager = get_monitoring_manager() + with monitoring_manager.trace_operation( + "context.assemble_final_context", + OPENINFERENCE_SPAN_KIND_CHAIN, + **{ + "context.purpose": purpose, + "context.conversation_id": conversation_id, + "context.use_context_items": self.config.use_context_items, + "context.current_run_start_idx": current_run_start_idx, + OPENINFERENCE_INPUT_VALUE: { + "purpose": purpose, + "memory_steps": len(memory.steps) if hasattr(memory, 'steps') else 0, + "has_tools": bool(tools), + "has_run_context": run_context is not None, + }, + }, + ): + if run_context is None: + run_context = self.prepare_run_context(memory, fallback_system_prompt="") + + context_items_for_evidence: tuple = () + if self.config.use_context_items: + projected_items = self.project_context_items(run_context.selected_components) + + history_projector = self.config.history_projector + if history_projector is not None and conversation_id is not None: + try: + history_items = history_projector.project( + conversation_id=conversation_id, + purpose=purpose if purpose in ("model_context", "resume", "chat") else "model_context", + ) + projected_items.extend(history_items) + except Exception: + logger.warning("History projection failed, continuing without history items", exc_info=True) + + context_items_for_evidence = tuple(projected_items) + + seen_components = set() + item_messages = [] + for item in projected_items: + source_component = item.metadata.get("_source_component") + if source_component and id(source_component) not in seen_components: + seen_components.add(id(source_component)) + for msg in source_component.to_messages(): + item_messages.append(msg) + + for item in projected_items: + if item.metadata.get("_source_component") is None: + from ..context.item_handler_registry import ItemHandlerRegistry + handler = ItemHandlerRegistry.get(item.item_type) + for msg in handler.to_messages(item): + item_messages.append(msg) + + stable_from_items = [m for m in item_messages if message_role(m) in ("system", "developer")] + dynamic_from_items = [m for m in item_messages if message_role(m) not in ("system", "developer")] + + run_context = ManagedRunContext( + component_messages=tuple(item_messages), + stable_messages=tuple(stable_from_items), + dynamic_messages=tuple(dynamic_from_items), + selected_component_types=run_context.selected_component_types, + selected_components=run_context.selected_components, + ) - original_messages = self._messages_from_memory(memory) - stable_messages = [*run_context.stable_messages, *purpose_stable] - dynamic_messages = [*run_context.dynamic_messages, *purpose_dynamic] + tools = self._canonical_tools(tools or ()) + purpose_stable, purpose_dynamic = self._purpose_messages( + purpose=purpose, + task=task, + final_answer_templates=final_answer_templates, + ) - context_overhead_tokens = ( - self._msg_token_count(dynamic_messages) - + self._estimate_tools_tokens(tools) - + self._msg_token_count(purpose_stable) - ) - compressed_messages = self.compress_if_needed( - model, - memory, - original_messages, - current_run_start_idx, - context_overhead_tokens=context_overhead_tokens, - ) - history_messages = self._without_leading_stable_messages(compressed_messages) - messages = [ - *stable_messages, - *dynamic_messages, - *history_messages, - ] + original_messages = self._messages_from_memory(memory) + stable_messages = [*run_context.stable_messages, *purpose_stable] + dynamic_messages = [*run_context.dynamic_messages, *purpose_dynamic] - self._last_compressed_token_count = self._msg_token_count(messages) + self._estimate_tools_tokens(tools) + context_overhead_tokens = ( + self._msg_token_count(dynamic_messages) + + self._estimate_tools_tokens(tools) + + self._msg_token_count(purpose_stable) + ) + compressed_messages = self.compress_if_needed( + model, + memory, + original_messages, + current_run_start_idx, + context_overhead_tokens=context_overhead_tokens, + ) + history_messages = self._without_leading_stable_messages(compressed_messages) + messages = [ + *stable_messages, + *dynamic_messages, + *history_messages, + ] - fingerprint = self._fingerprint({"messages": stable_messages, "tools": tools}) - component_fingerprints = self._stable_component_fingerprints( - purpose_stable, - components=run_context.components, - ) - if tools: - component_fingerprints["tools"] = self._fingerprint(tools) - reasons = self._change_reasons(fingerprint, component_fingerprints) - self._previous_stable_fingerprint = fingerprint - self._previous_stable_components = component_fingerprints - - from ...context_runtime.contracts import ContextEvidence, FinalContext - - return FinalContext( - messages=messages, - tools=tools, - evidence=ContextEvidence( - selected_component_types=run_context.selected_component_types, - stable_message_count=len(stable_messages), - dynamic_message_count=len(messages) - len(stable_messages), - compression_records=tuple(self._step_local_log or ()), - stable_prefix_fingerprint=fingerprint, - prefix_change_reasons=tuple(reasons), - ), - ) + self._last_compressed_token_count = self._msg_token_count(messages) + self._estimate_tools_tokens(tools) + + fingerprint = self._fingerprint({"messages": stable_messages, "tools": tools}) + component_fingerprints = self._stable_component_fingerprints( + purpose_stable, + components=run_context.selected_components, + ) + if tools: + component_fingerprints["tools"] = self._fingerprint(tools) + reasons = self._change_reasons(fingerprint, component_fingerprints) + self._previous_stable_fingerprint = fingerprint + self._previous_stable_components = component_fingerprints + + if monitoring_manager.is_enabled: + monitoring_manager.set_openinference_output({ + "message_count": len(messages), + "stable_count": len(stable_messages), + "dynamic_count": len(dynamic_messages), + "tool_count": len(tools), + "context_items": len(context_items_for_evidence), + "total_tokens": self._last_compressed_token_count, + }) + + from ...context_runtime.contracts import ContextEvidence, FinalContext + + return FinalContext( + messages=messages, + tools=tools, + evidence=ContextEvidence( + selected_component_types=run_context.selected_component_types, + stable_message_count=len(stable_messages), + dynamic_message_count=len(messages) - len(stable_messages), + compression_records=tuple(self._step_local_log or ()), + stable_prefix_fingerprint=fingerprint, + prefix_change_reasons=tuple(reasons), + context_items=context_items_for_evidence, + ), + ) def _purpose_messages( self, @@ -758,6 +896,39 @@ def build_context_messages( return messages + def project_context_items( + self, + components: Optional[Sequence[Any]] = None, + ) -> List[Any]: + """Project registered components into fine-grained ContextItem candidates.""" + monitoring_manager = get_monitoring_manager() + _ensure_handlers_registered() + from ..context.projector import ContextProjector + + source_components = self._component_source(components) + + with monitoring_manager.trace_operation( + "context.project_items", + OPENINFERENCE_SPAN_KIND_CHAIN, + **{ + "context.component_count": len(source_components), + OPENINFERENCE_INPUT_VALUE: [ + getattr(c, "component_type", type(c).__name__) for c in source_components + ], + }, + ): + projector = ContextProjector() + items = projector.project(list(source_components)) + + if monitoring_manager.is_enabled: + monitoring_manager.set_openinference_output({ + "item_count": len(items), + "item_types": [item.item_type.value for item in items], + }) + + return items + + def build_system_prompt(self, token_budget: Optional[int] = None) -> List: return self.build_context_messages(token_budget) diff --git a/sdk/nexent/core/agents/agent_context/summary_step.py b/sdk/nexent/core/agents/agent_context/summary_step.py index 77e294631..724ae2f16 100644 --- a/sdk/nexent/core/agents/agent_context/summary_step.py +++ b/sdk/nexent/core/agents/agent_context/summary_step.py @@ -26,4 +26,4 @@ class ManagedRunContext: stable_messages: Tuple[dict, ...] = () dynamic_messages: Tuple[dict, ...] = () selected_component_types: Tuple[str, ...] = () - components: Tuple[Any, ...] = () + selected_components: Tuple[Any, ...] = () diff --git a/sdk/nexent/core/agents/agent_model.py b/sdk/nexent/core/agents/agent_model.py index 7da3048a3..ebbbf3626 100644 --- a/sdk/nexent/core/agents/agent_model.py +++ b/sdk/nexent/core/agents/agent_model.py @@ -240,6 +240,10 @@ class AgentConfig(BaseModel): description="Layered ReAct self-verification configuration", default_factory=AgentVerificationConfig, ) + conversation_id: Optional[int] = Field( + description="Conversation ID for history projection in managed context runtime", + default=None, + ) class AgentHistory(BaseModel): diff --git a/sdk/nexent/core/agents/context/__init__.py b/sdk/nexent/core/agents/context/__init__.py new file mode 100644 index 000000000..aad3f7e5e --- /dev/null +++ b/sdk/nexent/core/agents/context/__init__.py @@ -0,0 +1,13 @@ +"""Context management module for fine-grained context assembly, policy, and reduction.""" + +from .context_item import ( + AuthorityTier, + ContextItem, + ContextItemType, + RepresentationTier, +) +from .item_handler import ContextItemHandler +from .item_handler_registry import ItemHandlerRegistry +from .history_projector import HistoryProjector +from .policy_models import MemoryDecision, SelectionDecision +from .reducer_models import ReductionResult diff --git a/sdk/nexent/core/agents/context/context_item.py b/sdk/nexent/core/agents/context/context_item.py new file mode 100644 index 000000000..ba8c81e33 --- /dev/null +++ b/sdk/nexent/core/agents/context/context_item.py @@ -0,0 +1,64 @@ +"""Context item data models for fine-grained context management.""" + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Dict, List, Optional + + +class ContextItemType(str, Enum): + """Enumeration of all context item types.""" + + SYSTEM_PROMPT = "system_prompt" + TOOL = "tool" + SKILL = "skill" + MEMORY = "memory" + KNOWLEDGE_BASE = "knowledge_base" + MANAGED_AGENT = "managed_agent" + EXTERNAL_AGENT = "external_agent" + HISTORY_TURN = "history_turn" + TOOL_CALL_RESULT = "tool_call_result" + + +class RepresentationTier(str, Enum): + """Fidelity levels for context item representations.""" + + FULL = "full" + COMPRESSED = "compressed" + STRUCTURED = "structured" + POINTER = "pointer" + + +class AuthorityTier(str, Enum): + """Authority tiers for context item provenance tracking.""" + + PLATFORM = "platform" + TENANT = "tenant" + USER = "user" + TOOL_RESULT = "tool_result" + RETRIEVED_MEMORY = "retrieved_memory" + SUMMARY = "summary" + AGENT_INFERENCE = "agent_inference" + + +@dataclass +class ContextItem: + """Bounded, source-traced context candidate unit. + + Each ContextItem represents a discrete piece of context that can be + selected, scored, reduced, and tracked independently. Items carry + provenance metadata (source_refs, authority_tier) and fidelity + constraints (minimum_fidelity, current_representation) to support + policy-driven context management. + """ + + item_id: str + item_type: ContextItemType + source_refs: List[str] = field(default_factory=list) + authority_tier: AuthorityTier = AuthorityTier.AGENT_INFERENCE + minimum_fidelity: RepresentationTier = RepresentationTier.STRUCTURED + current_representation: RepresentationTier = RepresentationTier.FULL + content: Any = None + token_estimate: int = 0 + metadata: Dict[str, Any] = field(default_factory=dict) + lifecycle_status: str = "active" + recompute_cost: Optional[int] = None diff --git a/sdk/nexent/core/agents/context/handlers/__init__.py b/sdk/nexent/core/agents/context/handlers/__init__.py new file mode 100644 index 000000000..c79fd8980 --- /dev/null +++ b/sdk/nexent/core/agents/context/handlers/__init__.py @@ -0,0 +1,31 @@ +"""Context item handler implementations for all supported item types.""" + +from .system_prompt_handler import SystemPromptHandler +from .tool_handler import ToolHandler +from .skill_handler import SkillHandler +from .memory_handler import MemoryHandler +from .knowledge_base_handler import KnowledgeBaseHandler +from .managed_agent_handler import ManagedAgentHandler +from .external_agent_handler import ExternalAgentHandler +from .history_turn_handler import HistoryTurnHandler +from .tool_call_result_handler import ToolCallResultHandler + +ALL_HANDLERS = [ + SystemPromptHandler, + ToolHandler, + SkillHandler, + MemoryHandler, + KnowledgeBaseHandler, + ManagedAgentHandler, + ExternalAgentHandler, + HistoryTurnHandler, + ToolCallResultHandler, +] + + +def register_all() -> None: + """Register all built-in handlers with the ItemHandlerRegistry.""" + from ..item_handler_registry import ItemHandlerRegistry + + for handler_cls in ALL_HANDLERS: + ItemHandlerRegistry.register(handler_cls()) diff --git a/sdk/nexent/core/agents/context/handlers/external_agent_handler.py b/sdk/nexent/core/agents/context/handlers/external_agent_handler.py new file mode 100644 index 000000000..850fca351 --- /dev/null +++ b/sdk/nexent/core/agents/context/handlers/external_agent_handler.py @@ -0,0 +1,38 @@ +"""Handler for external agent context items.""" + +from typing import Any, Dict, List + +from ..context_item import ContextItem, ContextItemType, RepresentationTier +from ..item_handler import ContextItemHandler +from ..reducer_models import ReductionResult + + +class ExternalAgentHandler(ContextItemHandler): + """Scores and reduces external A2A agent definitions.""" + + def supported_types(self) -> List[ContextItemType]: + return [ContextItemType.EXTERNAL_AGENT] + + def score(self, item: ContextItem, query: str, context: Dict[str, Any]) -> float: + # TODO(W13): Implement scoring: + # score = keyword_overlap(query, description) * 0.5 + # + priority * 0.5 + return 1.0 + + def reduce( + self, item: ContextItem, target: RepresentationTier, budget: int + ) -> ReductionResult: + # TODO(W8): Implement tiered reduction: + # STRUCTURED: name + routing metadata (description + capability tags) + # POINTER: name + capability tags only + # Deterministic, no LLM call needed. + return ReductionResult( + representation=item.current_representation, + source_fingerprint="", + token_count=item.token_estimate, + generator="passthrough", + generator_version="0.1.0", + admissible=True, + loss_metadata={}, + content=item.content, + ) diff --git a/sdk/nexent/core/agents/context/handlers/history_turn_handler.py b/sdk/nexent/core/agents/context/handlers/history_turn_handler.py new file mode 100644 index 000000000..a08e5e9f0 --- /dev/null +++ b/sdk/nexent/core/agents/context/handlers/history_turn_handler.py @@ -0,0 +1,51 @@ +"""Handler for history turn context items.""" + +from typing import Any, Dict, List + +from ..context_item import ContextItem, ContextItemType, RepresentationTier +from ..item_handler import ContextItemHandler +from ..reducer_models import ReductionResult + + +class HistoryTurnHandler(ContextItemHandler): + """Scores and reduces conversation history turns.""" + + def supported_types(self) -> List[ContextItemType]: + return [ContextItemType.HISTORY_TURN] + + def score(self, item: ContextItem, query: str, context: Dict[str, Any]) -> float: + # TODO(W13): Implement weighted scoring: + # score = recency * 0.5 + # + has_pending_action * 0.3 + # + keyword_overlap * 0.2 + # Signals: item.metadata["message_id"], item.metadata["step_index"], + # item.metadata["has_pending_action"] + return 1.0 + + def reduce( + self, item: ContextItem, target: RepresentationTier, budget: int + ) -> ReductionResult: + # TODO(W8): Implement tiered reduction: + # COMPRESSED: LLM summary (reuse existing compression logic) + # STRUCTURED: user query summary + assistant conclusion (deterministic) + return ReductionResult( + representation=item.current_representation, + source_fingerprint="", + token_count=item.token_estimate, + generator="passthrough", + generator_version="0.1.0", + admissible=True, + loss_metadata={}, + content=item.content, + ) + + def to_messages(self, item: ContextItem) -> List[Dict[str, Any]]: + content = item.content or {} + messages = [] + user_query = content.get("user_query", "") + if user_query: + messages.append({"role": "user", "content": [{"type": "text", "text": user_query}]}) + assistant_response = content.get("assistant_response", "") + if assistant_response: + messages.append({"role": "assistant", "content": [{"type": "text", "text": assistant_response}]}) + return messages diff --git a/sdk/nexent/core/agents/context/handlers/knowledge_base_handler.py b/sdk/nexent/core/agents/context/handlers/knowledge_base_handler.py new file mode 100644 index 000000000..e998950bc --- /dev/null +++ b/sdk/nexent/core/agents/context/handlers/knowledge_base_handler.py @@ -0,0 +1,37 @@ +"""Handler for knowledge base context items.""" + +from typing import Any, Dict, List + +from ..context_item import ContextItem, ContextItemType, RepresentationTier +from ..item_handler import ContextItemHandler +from ..reducer_models import ReductionResult + + +class KnowledgeBaseHandler(ContextItemHandler): + """Scores and reduces knowledge base retrieval results.""" + + def supported_types(self) -> List[ContextItemType]: + return [ContextItemType.KNOWLEDGE_BASE] + + def score(self, item: ContextItem, query: str, context: Dict[str, Any]) -> float: + # TODO(W13): Implement scoring: + # score = relevance_score from KB retrieval + # Signals: item.metadata["relevance_score"] + return 1.0 + + def reduce( + self, item: ContextItem, target: RepresentationTier, budget: int + ) -> ReductionResult: + # TODO(W8): Implement tiered reduction: + # COMPRESSED: LLM summary + # STRUCTURED: KB ID + title + relevance score (deterministic) + return ReductionResult( + representation=item.current_representation, + source_fingerprint="", + token_count=item.token_estimate, + generator="passthrough", + generator_version="0.1.0", + admissible=True, + loss_metadata={}, + content=item.content, + ) diff --git a/sdk/nexent/core/agents/context/handlers/managed_agent_handler.py b/sdk/nexent/core/agents/context/handlers/managed_agent_handler.py new file mode 100644 index 000000000..6ec21b328 --- /dev/null +++ b/sdk/nexent/core/agents/context/handlers/managed_agent_handler.py @@ -0,0 +1,38 @@ +"""Handler for managed agent context items.""" + +from typing import Any, Dict, List + +from ..context_item import ContextItem, ContextItemType, RepresentationTier +from ..item_handler import ContextItemHandler +from ..reducer_models import ReductionResult + + +class ManagedAgentHandler(ContextItemHandler): + """Scores and reduces internal managed sub-agent definitions.""" + + def supported_types(self) -> List[ContextItemType]: + return [ContextItemType.MANAGED_AGENT] + + def score(self, item: ContextItem, query: str, context: Dict[str, Any]) -> float: + # TODO(W13): Implement scoring: + # score = keyword_overlap(query, description) * 0.5 + # + priority * 0.5 + return 1.0 + + def reduce( + self, item: ContextItem, target: RepresentationTier, budget: int + ) -> ReductionResult: + # TODO(W8): Implement tiered reduction: + # STRUCTURED: name + routing metadata (description + capability tags) + # POINTER: name + capability tags only + # Deterministic, no LLM call needed. + return ReductionResult( + representation=item.current_representation, + source_fingerprint="", + token_count=item.token_estimate, + generator="passthrough", + generator_version="0.1.0", + admissible=True, + loss_metadata={}, + content=item.content, + ) diff --git a/sdk/nexent/core/agents/context/handlers/memory_handler.py b/sdk/nexent/core/agents/context/handlers/memory_handler.py new file mode 100644 index 000000000..184c226ac --- /dev/null +++ b/sdk/nexent/core/agents/context/handlers/memory_handler.py @@ -0,0 +1,41 @@ +"""Handler for memory context items.""" + +from typing import Any, Dict, List + +from ..context_item import ContextItem, ContextItemType, RepresentationTier +from ..item_handler import ContextItemHandler +from ..reducer_models import ReductionResult + + +class MemoryHandler(ContextItemHandler): + """Scores and reduces long-term memory (mem0) search results.""" + + def supported_types(self) -> List[ContextItemType]: + return [ContextItemType.MEMORY] + + def score(self, item: ContextItem, query: str, context: Dict[str, Any]) -> float: + # TODO(W13): Implement weighted scoring: + # score = mem0_relevance_score * 0.5 + # + recency * 0.2 + # + authority_weight * 0.3 + # Signals: item.metadata["mem0_score"], item.metadata["created_at"], + # item.authority_tier (user_agent > agent > user > tenant) + return 1.0 + + def reduce( + self, item: ContextItem, target: RepresentationTier, budget: int + ) -> ReductionResult: + # TODO(W8): Implement tiered reduction: + # COMPRESSED: LLM summary (reuse existing compress_if_needed prompt) + # STRUCTURED: extract core facts as key-value pairs (deterministic) + # POINTER: memory level + timestamp + first 50 chars preview (deterministic) + return ReductionResult( + representation=item.current_representation, + source_fingerprint="", + token_count=item.token_estimate, + generator="passthrough", + generator_version="0.1.0", + admissible=True, + loss_metadata={}, + content=item.content, + ) diff --git a/sdk/nexent/core/agents/context/handlers/skill_handler.py b/sdk/nexent/core/agents/context/handlers/skill_handler.py new file mode 100644 index 000000000..96a464ae0 --- /dev/null +++ b/sdk/nexent/core/agents/context/handlers/skill_handler.py @@ -0,0 +1,39 @@ +"""Handler for skill context items.""" + +from typing import Any, Dict, List + +from ..context_item import ContextItem, ContextItemType, RepresentationTier +from ..item_handler import ContextItemHandler +from ..reducer_models import ReductionResult + + +class SkillHandler(ContextItemHandler): + """Scores and reduces skill summaries based on query relevance.""" + + def supported_types(self) -> List[ContextItemType]: + return [ContextItemType.SKILL] + + def score(self, item: ContextItem, query: str, context: Dict[str, Any]) -> float: + # TODO(W13): Implement scoring: + # score = keyword_overlap(query, description) * 0.6 + # + priority * 0.4 + # Skills have no usage history, rely on query matching. + return 1.0 + + def reduce( + self, item: ContextItem, target: RepresentationTier, budget: int + ) -> ReductionResult: + # TODO(W8): Implement tiered reduction: + # STRUCTURED: name + first sentence of description (truncate at 100 chars) + # POINTER: keep only name + # Deterministic, no LLM call needed. + return ReductionResult( + representation=item.current_representation, + source_fingerprint="", + token_count=item.token_estimate, + generator="passthrough", + generator_version="0.1.0", + admissible=True, + loss_metadata={}, + content=item.content, + ) diff --git a/sdk/nexent/core/agents/context/handlers/system_prompt_handler.py b/sdk/nexent/core/agents/context/handlers/system_prompt_handler.py new file mode 100644 index 000000000..c00fd14ec --- /dev/null +++ b/sdk/nexent/core/agents/context/handlers/system_prompt_handler.py @@ -0,0 +1,34 @@ +"""Handler for system prompt context items.""" + +from typing import Any, Dict, List + +from ..context_item import ContextItem, ContextItemType, RepresentationTier +from ..item_handler import ContextItemHandler +from ..reducer_models import ReductionResult + + +class SystemPromptHandler(ContextItemHandler): + """System prompts are mandatory and irreducible.""" + + def supported_types(self) -> List[ContextItemType]: + return [ContextItemType.SYSTEM_PROMPT] + + def score(self, item: ContextItem, query: str, context: Dict[str, Any]) -> float: + # TODO: Not selectable -- mandatory item, score=inf to guarantee inclusion + return float("inf") + + def reduce( + self, item: ContextItem, target: RepresentationTier, budget: int + ) -> ReductionResult: + # TODO: Not reducible -- minimum_fidelity=FULL, reject any reduction attempt + # Suggested: raise MinimumFidelityViolation if target != FULL + return ReductionResult( + representation=RepresentationTier.FULL, + source_fingerprint="", + token_count=item.token_estimate, + generator="passthrough", + generator_version="0.1.0", + admissible=True, + loss_metadata={}, + content=item.content, + ) diff --git a/sdk/nexent/core/agents/context/handlers/tool_call_result_handler.py b/sdk/nexent/core/agents/context/handlers/tool_call_result_handler.py new file mode 100644 index 000000000..19c0edbf0 --- /dev/null +++ b/sdk/nexent/core/agents/context/handlers/tool_call_result_handler.py @@ -0,0 +1,46 @@ +"""Handler for tool call result context items.""" + +from typing import Any, Dict, List + +from ..context_item import ContextItem, ContextItemType, RepresentationTier +from ..item_handler import ContextItemHandler +from ..reducer_models import ReductionResult + + +class ToolCallResultHandler(ContextItemHandler): + """Scores and reduces tool call results from the current conversation.""" + + def supported_types(self) -> List[ContextItemType]: + return [ContextItemType.TOOL_CALL_RESULT] + + def score(self, item: ContextItem, query: str, context: Dict[str, Any]) -> float: + # TODO(W13): Implement weighted scoring: + # score = recency * 0.4 + # + is_active_tool * 0.4 + # + result_relevance * 0.2 + # Signals: item.metadata["message_id"], item.metadata["tool_name"] + return 1.0 + + def reduce( + self, item: ContextItem, target: RepresentationTier, budget: int + ) -> ReductionResult: + # TODO(W8): Implement tiered reduction: + # STRUCTURED: tool name + result summary (deterministic) + # POINTER: tool name + status only (deterministic) + return ReductionResult( + representation=item.current_representation, + source_fingerprint="", + token_count=item.token_estimate, + generator="passthrough", + generator_version="0.1.0", + admissible=True, + loss_metadata={}, + content=item.content, + ) + + def to_messages(self, item: ContextItem) -> List[Dict[str, Any]]: + content = item.content or {} + tool_call = content.get("tool_call", "") + execution_result = content.get("execution_result", "") + text = f"[Tool Call]\n{tool_call}\n\n[Execution Result]\n{execution_result}" + return [{"role": "user", "content": [{"type": "text", "text": text}]}] diff --git a/sdk/nexent/core/agents/context/handlers/tool_handler.py b/sdk/nexent/core/agents/context/handlers/tool_handler.py new file mode 100644 index 000000000..cf4506bc5 --- /dev/null +++ b/sdk/nexent/core/agents/context/handlers/tool_handler.py @@ -0,0 +1,41 @@ +"""Handler for tool context items.""" + +from typing import Any, Dict, List + +from ..context_item import ContextItem, ContextItemType, RepresentationTier +from ..item_handler import ContextItemHandler +from ..reducer_models import ReductionResult + + +class ToolHandler(ContextItemHandler): + """Scores and reduces tool definitions based on query relevance and usage.""" + + def supported_types(self) -> List[ContextItemType]: + return [ContextItemType.TOOL] + + def score(self, item: ContextItem, query: str, context: Dict[str, Any]) -> float: + # TODO(W13): Implement weighted scoring: + # score = priority * 0.4 + # + keyword_overlap(query, description) * 0.3 + # + usage_frequency * 0.3 + # Boost if tool was called in current run. + # Signals: item.metadata["priority"], item.metadata["usage_count"] + return 1.0 + + def reduce( + self, item: ContextItem, target: RepresentationTier, budget: int + ) -> ReductionResult: + # TODO(W8): Implement tiered reduction: + # STRUCTURED: template trim - keep name + one-line description + param names + # POINTER: keep only name + param count metadata + # Deterministic, no LLM call needed. + return ReductionResult( + representation=item.current_representation, + source_fingerprint="", + token_count=item.token_estimate, + generator="passthrough", + generator_version="0.1.0", + admissible=True, + loss_metadata={}, + content=item.content, + ) diff --git a/sdk/nexent/core/agents/context/history_projector.py b/sdk/nexent/core/agents/context/history_projector.py new file mode 100644 index 000000000..67723f134 --- /dev/null +++ b/sdk/nexent/core/agents/context/history_projector.py @@ -0,0 +1,277 @@ +"""HistoryProjector: projects conversation history from database into ContextItem instances.""" + +import json +from typing import Any, Callable, Dict, List, Optional + +from .context_item import AuthorityTier, ContextItem, ContextItemType, RepresentationTier +from .item_handler_registry import ItemHandlerRegistry +from nexent.monitor import ( + get_monitoring_manager, + OPENINFERENCE_SPAN_KIND_CHAIN, + OPENINFERENCE_SPAN_KIND_RETRIEVER, + OPENINFERENCE_INPUT_VALUE, + OPENINFERENCE_OUTPUT_VALUE, +) + + +class HistoryProjector: + """Projects conversation history from database into ContextItem instances. + + Uses dependency injection for database queries to maintain SDK/backend separation. + Produces HISTORY_TURN and TOOL_CALL_RESULT ContextItems. + """ + + def __init__(self, query_units_fn: Callable[[int, Optional[int]], List[Dict[str, Any]]]): + """Initialize with a database query function. + + Args: + query_units_fn: Function that takes (conversation_id, message_id) and returns + a list of unit dictionaries ordered by message_id, step_index, unit_index. + This is typically get_message_units_by_message_id from conversation_db. + """ + self.query_units_fn = query_units_fn + + def project( + self, + conversation_id: int, + message_id: Optional[int] = None, + purpose: str = "model_context", + ) -> List[ContextItem]: + """Project conversation history into ContextItem instances. + + Args: + conversation_id: The conversation to project. + message_id: Optional specific message to project (None = all messages). + purpose: Projection purpose - "model_context" or "chat". + + Returns: + List of ContextItem instances. + + Raises: + ValueError: If purpose is not recognized. + """ + monitoring_manager = get_monitoring_manager() + with monitoring_manager.trace_operation( + "context.history_project", + OPENINFERENCE_SPAN_KIND_CHAIN, + **{ + "context.conversation_id": conversation_id, + "context.message_id": message_id, + "context.purpose": purpose, + OPENINFERENCE_INPUT_VALUE: { + "conversation_id": conversation_id, + "message_id": message_id, + "purpose": purpose, + }, + }, + ): + with monitoring_manager.trace_operation( + "context.history_query", + OPENINFERENCE_SPAN_KIND_RETRIEVER, + **{ + "context.conversation_id": conversation_id, + "context.message_id": message_id, + OPENINFERENCE_INPUT_VALUE: { + "conversation_id": conversation_id, + "message_id": message_id, + }, + }, + ): + units = self.query_units_fn(conversation_id, message_id) + if monitoring_manager.is_enabled: + unit_types = {} + for u in units: + ut = u.get("unit_type", "unknown") + unit_types[ut] = unit_types.get(ut, 0) + 1 + monitoring_manager.set_openinference_output({ + "unit_count": len(units), + "unit_types": unit_types, + }) + + if purpose == "model_context": + items = self._project_model_context(units) + elif purpose == "chat": + items = self._project_chat_context(units) + else: + raise ValueError(f"Unknown purpose: {purpose}") + + if monitoring_manager.is_enabled: + monitoring_manager.set_openinference_output({ + "item_count": len(items), + "item_types": [item.item_type.value for item in items], + }) + + return items + + def _project_model_context(self, units: List[Dict[str, Any]]) -> List[ContextItem]: + """Project units into HISTORY_TURN and TOOL_CALL_RESULT items for model context. + + Groups units by message_id and step_index, then: + - Pairs user messages with assistant final_answer -> HISTORY_TURN + - Extracts merged tool_call rows -> TOOL_CALL_RESULT + - Excludes model_output_thinking and model_output_deep_thinking + """ + items: List[ContextItem] = [] + + messages = self._group_by_message(units) + + for message_id, steps in messages.items(): + for step_index, step_units in steps.items(): + history_turn = self._extract_history_turn(step_units, message_id, step_index) + if history_turn: + items.append(history_turn) + + tool_results = self._extract_tool_call_results(step_units, message_id, step_index) + items.extend(tool_results) + + for item in items: + ItemHandlerRegistry.get(item.item_type) + + return items + + def _project_chat_context(self, units: List[Dict[str, Any]]) -> List[ContextItem]: + """Project units into HISTORY_TURN items for chat display. + + Similar to model_context but includes thinking units for full transparency. + """ + items: List[ContextItem] = [] + + messages = self._group_by_message(units) + + for message_id, steps in messages.items(): + for step_index, step_units in steps.items(): + chat_turn = self._extract_chat_turn(step_units, message_id, step_index) + if chat_turn: + items.append(chat_turn) + + return items + + def _group_by_message( + self, units: List[Dict[str, Any]] + ) -> Dict[int, Dict[int, List[Dict[str, Any]]]]: + """Group units by message_id and step_index. + + Returns: + Dict[message_id, Dict[step_index, List[unit]]] + """ + messages: Dict[int, Dict[int, List[Dict[str, Any]]]] = {} + for unit in units: + message_id = unit.get("message_id") or 0 + step_index = unit.get("step_index") or 0 + if message_id not in messages: + messages[message_id] = {} + if step_index not in messages[message_id]: + messages[message_id][step_index] = [] + messages[message_id][step_index].append(unit) + return messages + + def _extract_history_turn( + self, + units: List[Dict[str, Any]], + message_id: int, + step_index: int, + ) -> Optional[ContextItem]: + """Extract a HISTORY_TURN item from user input and final_answer units. + + Returns None if no user input or final_answer found. + """ + user_units = [u for u in units if u.get("unit_type") == "user_input"] + answer_units = [u for u in units if u.get("unit_type") == "final_answer"] + + if not user_units or not answer_units: + return None + + user_content = user_units[0].get("unit_content", "") + answer_content = answer_units[0].get("unit_content", "") + + content = { + "user_query": user_content, + "assistant_response": answer_content, + } + + source_refs = [ + f"unit:{user_units[0]['unit_id']}", + f"unit:{answer_units[0]['unit_id']}", + ] + + return ContextItem( + item_id=f"history_turn:{message_id}:{step_index}", + item_type=ContextItemType.HISTORY_TURN, + source_refs=source_refs, + authority_tier=AuthorityTier.AGENT_INFERENCE, + minimum_fidelity=RepresentationTier.STRUCTURED, + current_representation=RepresentationTier.FULL, + content=content, + token_estimate=(len(user_content) + len(answer_content)) // 4, + metadata={"message_id": message_id, "step_index": step_index}, + ) + + def _extract_tool_call_results( + self, + units: List[Dict[str, Any]], + message_id: int, + step_index: int, + ) -> List[ContextItem]: + """Extract TOOL_CALL_RESULT items from merged tool_call rows.""" + items: List[ContextItem] = [] + for unit in units: + if unit.get("unit_type") != "tool_call": + continue + try: + content = json.loads(unit.get("unit_content", "{}")) + except (json.JSONDecodeError, TypeError): + continue + items.append(ContextItem( + item_id=f"tool_call_result:{unit['unit_id']}", + item_type=ContextItemType.TOOL_CALL_RESULT, + source_refs=[f"unit:{unit['unit_id']}"], + authority_tier=AuthorityTier.TOOL_RESULT, + minimum_fidelity=RepresentationTier.STRUCTURED, + current_representation=RepresentationTier.FULL, + content=content, + token_estimate=len(unit.get("unit_content", "")) // 4, + metadata={"message_id": message_id, "step_index": step_index}, + )) + return items + + def _extract_chat_turn( + self, + units: List[Dict[str, Any]], + message_id: int, + step_index: int, + ) -> Optional[ContextItem]: + """Extract a chat turn including thinking for display purposes.""" + relevant_types = { + "user_input", + "model_output_thinking", + "model_output_code", + "final_answer", + } + relevant_units = [u for u in units if u.get("unit_type") in relevant_types] + + if not relevant_units: + return None + + content = { + "units": [ + { + "type": u.get("unit_type"), + "content": u.get("unit_content", ""), + } + for u in relevant_units + ] + } + + source_refs = [f"unit:{u['unit_id']}" for u in relevant_units] + + return ContextItem( + item_id=f"chat_turn:{message_id}:{step_index}", + item_type=ContextItemType.HISTORY_TURN, + source_refs=source_refs, + authority_tier=AuthorityTier.AGENT_INFERENCE, + minimum_fidelity=RepresentationTier.FULL, + current_representation=RepresentationTier.FULL, + content=content, + token_estimate=sum(len(u.get("unit_content", "")) for u in relevant_units) // 4, + metadata={"message_id": message_id, "step_index": step_index, "includes_thinking": True}, + ) diff --git a/sdk/nexent/core/agents/context/item_handler.py b/sdk/nexent/core/agents/context/item_handler.py new file mode 100644 index 000000000..148fb3f72 --- /dev/null +++ b/sdk/nexent/core/agents/context/item_handler.py @@ -0,0 +1,44 @@ +"""Abstract base class for context item handlers.""" + +from abc import ABC, abstractmethod +from typing import Any, Dict, List + +from .context_item import ContextItem, ContextItemType, RepresentationTier +from .reducer_models import ReductionResult + + +class ContextItemHandler(ABC): + """Base handler providing default passthrough implementations for scoring and reduction.""" + + @abstractmethod + def supported_types(self) -> List[ContextItemType]: + """Return the list of ContextItemType values this handler supports.""" + ... + + def score(self, item: ContextItem, query: str, context: Dict[str, Any]) -> float: + """Score an item's relevance. Default passthrough returns 1.0.""" + return 1.0 + + def reduce( + self, item: ContextItem, target: RepresentationTier, budget: int + ) -> ReductionResult: + """Reduce an item to a target representation. Default passthrough returns content unchanged.""" + return ReductionResult( + representation=item.current_representation, + source_fingerprint="", + token_count=item.token_estimate, + generator="passthrough", + generator_version="0.1.0", + admissible=True, + loss_metadata={}, + content=item.content, + ) + + def to_messages(self, item: ContextItem) -> List[Dict[str, Any]]: + """Convert a ContextItem into message dicts for LLM consumption. + + Default implementation wraps item content as a user-role text message. + Subclasses should override for type-specific formatting. + """ + text = str(item.content) if not isinstance(item.content, str) else item.content + return [{"role": "user", "content": [{"type": "text", "text": text}]}] diff --git a/sdk/nexent/core/agents/context/item_handler_registry.py b/sdk/nexent/core/agents/context/item_handler_registry.py new file mode 100644 index 000000000..b1fe832b0 --- /dev/null +++ b/sdk/nexent/core/agents/context/item_handler_registry.py @@ -0,0 +1,42 @@ +"""Registry mapping ContextItemType to ContextItemHandler instances.""" + +from typing import Dict + +from .context_item import ContextItemType +from .item_handler import ContextItemHandler + + +class ItemHandlerRegistry: + """Central registry for context item handlers. + + Ensures every ContextItemType has exactly one registered handler. + """ + + _handlers: Dict[ContextItemType, ContextItemHandler] = {} + + @classmethod + def register(cls, handler: ContextItemHandler) -> None: + """Register a handler for all its supported types.""" + for item_type in handler.supported_types(): + cls._handlers[item_type] = handler + + @classmethod + def get(cls, item_type: ContextItemType) -> ContextItemHandler: + """Retrieve the handler for a given item type. + + Raises: + KeyError: If no handler is registered for the item type. + """ + if item_type not in cls._handlers: + raise KeyError(f"No handler registered for ContextItemType.{item_type.value}") + return cls._handlers[item_type] + + @classmethod + def all_types_covered(cls) -> bool: + """Return True if every ContextItemType has a registered handler.""" + return all(item_type in cls._handlers for item_type in ContextItemType) + + @classmethod + def reset(cls) -> None: + """Clear all registered handlers. Useful for testing.""" + cls._handlers = {} diff --git a/sdk/nexent/core/agents/context/policy_models.py b/sdk/nexent/core/agents/context/policy_models.py new file mode 100644 index 000000000..6b596e35b --- /dev/null +++ b/sdk/nexent/core/agents/context/policy_models.py @@ -0,0 +1,42 @@ +"""Policy decision models for context selection and memory operations.""" + +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + +from .context_item import RepresentationTier + + +@dataclass(frozen=True) +class SelectionDecision: + """Immutable record of a context selection policy decision. + + Captures which items were selected or excluded, their representation + requirements, budget allocations, and the reason codes explaining + each decision. + """ + + selected_item_ids: List[str] + excluded_item_ids: List[str] + representation_requirements: Dict[str, RepresentationTier] + budget_allocations: Dict[str, int] + remaining_budget: int + conflicts: List[Dict[str, Any]] + reason_codes: List[str] + policy_version: str + decision_fingerprint: str + + +@dataclass(frozen=True) +class MemoryDecision: + """Immutable record of a memory operation policy decision. + + Captures the allowed operation, scopes, excluded candidates, + conflict resolutions, and any confirmation requirements. + """ + + operation: str + allowed_scopes: List[str] + excluded_candidates: List[str] + conflict_decisions: List[Dict[str, Any]] + confirmation_required: Optional[Dict[str, Any]] + reason_codes: List[str] diff --git a/sdk/nexent/core/agents/context/projector.py b/sdk/nexent/core/agents/context/projector.py new file mode 100644 index 000000000..ec8dfef5f --- /dev/null +++ b/sdk/nexent/core/agents/context/projector.py @@ -0,0 +1,236 @@ +"""ContextProjector: converts ContextComponent instances into fine-grained ContextItem candidates.""" + +import uuid +from typing import Any, Dict, List + +from ..agent_model import ( + ContextComponent, + ExternalAgentsComponent, + KnowledgeBaseComponent, + ManagedAgentsComponent, + MemoryComponent, + SkillsComponent, + SystemPromptComponent, + ToolsComponent, +) +from .context_item import AuthorityTier, ContextItem, ContextItemType, RepresentationTier +from .item_handler_registry import ItemHandlerRegistry + +_CHARS_PER_TOKEN = 1.5 + +_MINIMUM_FIDELITY_MAP: Dict[ContextItemType, RepresentationTier] = { + ContextItemType.SYSTEM_PROMPT: RepresentationTier.FULL, + ContextItemType.TOOL: RepresentationTier.STRUCTURED, + ContextItemType.SKILL: RepresentationTier.STRUCTURED, + ContextItemType.MEMORY: RepresentationTier.STRUCTURED, + ContextItemType.KNOWLEDGE_BASE: RepresentationTier.COMPRESSED, + ContextItemType.MANAGED_AGENT: RepresentationTier.STRUCTURED, + ContextItemType.EXTERNAL_AGENT: RepresentationTier.STRUCTURED, +} + + +def _estimate_item_tokens(content: Any) -> int: + return int(len(str(content)) / _CHARS_PER_TOKEN) + + +def _make_item_id(item_type: ContextItemType, suffix: str = "") -> str: + tag = item_type.value + if suffix: + return f"{tag}:{suffix}" + return f"{tag}:{uuid.uuid4().hex[:8]}" + + +class ContextProjector: + """Projects ContextComponent instances into fine-grained ContextItem candidates.""" + + def project(self, components: List[ContextComponent]) -> List[ContextItem]: + """Convert a list of ContextComponent into ContextItem list. + + Each component type maps to one or more ContextItem instances with + appropriate authority tiers and minimum fidelity constraints. + After creation, each item is validated against the ItemHandlerRegistry. + """ + items: List[ContextItem] = [] + for component in components: + items.extend(self._project_component(component)) + + for item in items: + ItemHandlerRegistry.get(item.item_type) + + return items + + def _project_component(self, component: ContextComponent) -> List[ContextItem]: + if isinstance(component, SystemPromptComponent): + return self._project_system_prompt(component) + if isinstance(component, ToolsComponent): + return self._project_tools(component) + if isinstance(component, SkillsComponent): + return self._project_skills(component) + if isinstance(component, MemoryComponent): + return self._project_memory(component) + if isinstance(component, KnowledgeBaseComponent): + return self._project_knowledge_base(component) + if isinstance(component, ManagedAgentsComponent): + return self._project_managed_agents(component) + if isinstance(component, ExternalAgentsComponent): + return self._project_external_agents(component) + return [] + + def _project_system_prompt(self, component: SystemPromptComponent) -> List[ContextItem]: + return [ + ContextItem( + item_id=_make_item_id(ContextItemType.SYSTEM_PROMPT, component.template_name or ""), + item_type=ContextItemType.SYSTEM_PROMPT, + source_refs=[component.template_name] if component.template_name else [], + authority_tier=AuthorityTier.PLATFORM, + minimum_fidelity=_MINIMUM_FIDELITY_MAP[ContextItemType.SYSTEM_PROMPT], + current_representation=RepresentationTier.FULL, + content=component.content, + token_estimate=component.token_estimate or component.estimate_tokens(), + metadata={ + "template_name": component.template_name, + "_source_component": component, + **component.metadata, + }, + ) + ] + + def _project_tools(self, component: ToolsComponent) -> List[ContextItem]: + items: List[ContextItem] = [] + for tool in component.tools: + tool_name = tool.get("name", "unknown") + items.append( + ContextItem( + item_id=_make_item_id(ContextItemType.TOOL, tool_name), + item_type=ContextItemType.TOOL, + source_refs=[tool_name], + authority_tier=AuthorityTier.PLATFORM, + minimum_fidelity=_MINIMUM_FIDELITY_MAP[ContextItemType.TOOL], + current_representation=RepresentationTier.FULL, + content=tool, + token_estimate=_estimate_item_tokens(tool), + metadata={ + "priority": component.priority, + "_source_component": component, + **component.metadata, + }, + ) + ) + return items + + def _project_skills(self, component: SkillsComponent) -> List[ContextItem]: + items: List[ContextItem] = [] + for skill in component.skills: + skill_name = skill.get("name", "unknown") + items.append( + ContextItem( + item_id=_make_item_id(ContextItemType.SKILL, skill_name), + item_type=ContextItemType.SKILL, + source_refs=[skill_name], + authority_tier=AuthorityTier.PLATFORM, + minimum_fidelity=_MINIMUM_FIDELITY_MAP[ContextItemType.SKILL], + current_representation=RepresentationTier.FULL, + content=skill, + token_estimate=_estimate_item_tokens(skill), + metadata={ + "priority": component.priority, + "_source_component": component, + **component.metadata, + }, + ) + ) + return items + + def _project_memory(self, component: MemoryComponent) -> List[ContextItem]: + items: List[ContextItem] = [] + for idx, memory in enumerate(component.memories): + memory_content = memory.get("content", "") + memory_type = memory.get("memory_type", "user") + items.append( + ContextItem( + item_id=_make_item_id(ContextItemType.MEMORY, f"{memory_type}:{idx}"), + item_type=ContextItemType.MEMORY, + source_refs=[f"mem0:{memory_type}"], + authority_tier=AuthorityTier.RETRIEVED_MEMORY, + minimum_fidelity=_MINIMUM_FIDELITY_MAP[ContextItemType.MEMORY], + current_representation=RepresentationTier.FULL, + content=memory, + token_estimate=_estimate_item_tokens(memory_content), + metadata={ + "memory_type": memory_type, + "search_query": component.search_query, + "_source_component": component, + **memory.get("metadata", {}), + **component.metadata, + }, + ) + ) + return items + + def _project_knowledge_base(self, component: KnowledgeBaseComponent) -> List[ContextItem]: + return [ + ContextItem( + item_id=_make_item_id( + ContextItemType.KNOWLEDGE_BASE, ":".join(component.kb_ids) + ), + item_type=ContextItemType.KNOWLEDGE_BASE, + source_refs=component.kb_ids, + authority_tier=AuthorityTier.RETRIEVED_MEMORY, + minimum_fidelity=_MINIMUM_FIDELITY_MAP[ContextItemType.KNOWLEDGE_BASE], + current_representation=RepresentationTier.FULL, + content=component.summary, + token_estimate=component.token_estimate or component.estimate_tokens(), + metadata={ + "kb_ids": component.kb_ids, + "_source_component": component, + **component.metadata, + }, + ) + ] + + def _project_managed_agents(self, component: ManagedAgentsComponent) -> List[ContextItem]: + items: List[ContextItem] = [] + for agent in component.agents: + agent_name = agent.get("name", "unknown") + items.append( + ContextItem( + item_id=_make_item_id(ContextItemType.MANAGED_AGENT, agent_name), + item_type=ContextItemType.MANAGED_AGENT, + source_refs=[agent_name], + authority_tier=AuthorityTier.PLATFORM, + minimum_fidelity=_MINIMUM_FIDELITY_MAP[ContextItemType.MANAGED_AGENT], + current_representation=RepresentationTier.FULL, + content=agent, + token_estimate=_estimate_item_tokens(agent), + metadata={ + "priority": component.priority, + "_source_component": component, + **component.metadata, + }, + ) + ) + return items + + def _project_external_agents(self, component: ExternalAgentsComponent) -> List[ContextItem]: + items: List[ContextItem] = [] + for agent in component.agents: + agent_name = agent.get("name", "unknown") + agent_id = agent.get("agent_id", "") + items.append( + ContextItem( + item_id=_make_item_id(ContextItemType.EXTERNAL_AGENT, agent_id or agent_name), + item_type=ContextItemType.EXTERNAL_AGENT, + source_refs=[agent_id or agent_name], + authority_tier=AuthorityTier.PLATFORM, + minimum_fidelity=_MINIMUM_FIDELITY_MAP[ContextItemType.EXTERNAL_AGENT], + current_representation=RepresentationTier.FULL, + content=agent, + token_estimate=_estimate_item_tokens(agent), + metadata={ + "priority": component.priority, + "_source_component": component, + **component.metadata, + }, + ) + ) + return items diff --git a/sdk/nexent/core/agents/context/reason_codes.py b/sdk/nexent/core/agents/context/reason_codes.py new file mode 100644 index 000000000..502c60c3c --- /dev/null +++ b/sdk/nexent/core/agents/context/reason_codes.py @@ -0,0 +1,24 @@ +"""Reason code constants for context policy decisions. + +These codes provide traceability for why items were selected, excluded, +or transformed during context management operations. +""" + +# Selection reason codes +SELECTED_MANDATORY_MINIMUM = "selected_mandatory_minimum" +SELECTED_BUDGET_UPGRADE = "selected_budget_upgrade" + +# Exclusion reason codes +EXCLUDED_BUDGET = "excluded_budget" +EXCLUDED_POLICY_DISABLED = "excluded_policy_disabled" +EXCLUDED_LOWER_AUTHORITY = "excluded_lower_authority" + +# Memory operation reason codes +MEMORY_OPERATION_ALLOWED = "memory_operation_allowed" +MEMORY_OPERATION_DENIED = "memory_operation_denied" +CONFIRMATION_REQUIRED = "confirmation_required" + +# Reduction reason codes +MINIMUM_FIDELITY_VIOLATION = "minimum_fidelity_violation" +REDUCER_FAILED = "reducer_failed" +REPRESENTATION_STALE = "representation_stale" diff --git a/sdk/nexent/core/agents/context/reducer_models.py b/sdk/nexent/core/agents/context/reducer_models.py new file mode 100644 index 000000000..f635ef38a --- /dev/null +++ b/sdk/nexent/core/agents/context/reducer_models.py @@ -0,0 +1,25 @@ +"""Reduction result models for context item representation transforms.""" + +from dataclasses import dataclass, field +from typing import Any, Dict + +from .context_item import RepresentationTier + + +@dataclass(frozen=True) +class ReductionResult: + """Immutable record of a context item reduction operation. + + Captures the output representation tier, token count, generator + metadata, admissibility flag, and any loss metadata describing + what information was discarded during reduction. + """ + + representation: RepresentationTier + source_fingerprint: str + token_count: int + generator: str + generator_version: str + admissible: bool + loss_metadata: Dict[str, Any] + content: Any diff --git a/sdk/nexent/core/agents/nexent_agent.py b/sdk/nexent/core/agents/nexent_agent.py index d65768c15..7987c5c0e 100644 --- a/sdk/nexent/core/agents/nexent_agent.py +++ b/sdk/nexent/core/agents/nexent_agent.py @@ -460,6 +460,7 @@ def create_single_agent(self, agent_config: AgentConfig): context_runtime = ManagedContextRuntime( context_manager, components=getattr(agent_config, 'context_components', None) or [], + conversation_id=getattr(agent_config, 'conversation_id', None), ) else: from ..context_runtime.legacy.runtime import LegacyContextRuntime diff --git a/sdk/nexent/core/agents/summary_config.py b/sdk/nexent/core/agents/summary_config.py index fcca60eb5..162905364 100644 --- a/sdk/nexent/core/agents/summary_config.py +++ b/sdk/nexent/core/agents/summary_config.py @@ -101,6 +101,25 @@ class ContextManagerConfig: inject_app_context: bool = True """Whether to inject APP_NAME, APP_DESCRIPTION, time, user_id.""" + # === NEW: ContextItem Projection === + use_context_items: bool = False + """Whether to project components into fine-grained ContextItems before assembly. + + When True, uses ContextProjector to convert ContextComponent instances into + ContextItem candidates, then converts them back to messages for the existing + assembly pipeline. This enables W12 history projections and W13 policy decisions + in subsequent PRs. Default False for backward compatibility. + """ + + history_projector: Any = None + """Optional HistoryProjector instance for projecting DB conversation history into ContextItems. + + When set and use_context_items=True, the projector is called during build_context_messages + to produce HISTORY_TURN and TOOL_CALL_RESULT items from persisted + conversation history. Must be injected by the backend service layer since the SDK + cannot import database modules directly. + """ + # === NEW: Per-Component Token Budgets === component_budgets: Dict[str, int] = field(default_factory=lambda: { "system_prompt": 4000, diff --git a/sdk/nexent/core/context_runtime/contracts.py b/sdk/nexent/core/context_runtime/contracts.py index 32bf44ae4..dc3cae474 100644 --- a/sdk/nexent/core/context_runtime/contracts.py +++ b/sdk/nexent/core/context_runtime/contracts.py @@ -16,6 +16,8 @@ class ContextEvidence: compression_records: tuple[Any, ...] = () stable_prefix_fingerprint: str | None = None prefix_change_reasons: tuple[str, ...] = () + context_items: tuple[Any, ...] = () + """Projected ContextItem candidates when use_context_items is enabled.""" @dataclass(frozen=True) diff --git a/sdk/nexent/core/context_runtime/managed/runtime.py b/sdk/nexent/core/context_runtime/managed/runtime.py index e66887dea..296c984c8 100644 --- a/sdk/nexent/core/context_runtime/managed/runtime.py +++ b/sdk/nexent/core/context_runtime/managed/runtime.py @@ -8,14 +8,21 @@ from typing import Any, Sequence from ..contracts import FinalContext +from nexent.monitor import ( + get_monitoring_manager, + OPENINFERENCE_SPAN_KIND_CHAIN, + OPENINFERENCE_INPUT_VALUE, + OPENINFERENCE_OUTPUT_VALUE, +) class ManagedContextRuntime: """Adapter for the ContextManager-owned managed path.""" - def __init__(self, context_manager: Any, components: Sequence[Any] | None = None): + def __init__(self, context_manager: Any, components: Sequence[Any] | None = None, conversation_id: int | None = None): self.context_manager = context_manager self.components = list(components or ()) + self.conversation_id = conversation_id self._run_context = None def replace_components(self, components: Sequence[Any] | None) -> None: @@ -47,14 +54,39 @@ def prepare_step( current_run_start_idx: int, tools: Sequence[Any] | None = None, ) -> FinalContext: - return self.context_manager.assemble_final_context( - model=model, - memory=memory, - current_run_start_idx=current_run_start_idx, - tools=tools, - purpose="step", - run_context=self._ensure_run_context(memory), - ) + monitoring_manager = get_monitoring_manager() + with monitoring_manager.trace_operation( + "context.prepare_step", + OPENINFERENCE_SPAN_KIND_CHAIN, + **{ + "context.current_run_start_idx": current_run_start_idx, + "context.conversation_id": self.conversation_id, + OPENINFERENCE_INPUT_VALUE: { + "current_run_start_idx": current_run_start_idx, + "memory_steps": len(memory.steps) if hasattr(memory, 'steps') else 0, + "tool_count": len(tools) if tools else 0, + "conversation_id": self.conversation_id, + }, + }, + ): + result = self.context_manager.assemble_final_context( + model=model, + memory=memory, + current_run_start_idx=current_run_start_idx, + tools=tools, + purpose="step", + run_context=self._ensure_run_context(memory), + conversation_id=self.conversation_id, + ) + if monitoring_manager.is_enabled: + monitoring_manager.set_openinference_output({ + "message_count": len(result.messages), + "tool_count": len(result.tools), + "stable_count": result.evidence.stable_message_count, + "dynamic_count": result.evidence.dynamic_message_count, + "context_items": len(result.evidence.context_items), + }) + return result def prepare_final_answer( self, @@ -66,16 +98,43 @@ def prepare_final_answer( final_answer_templates: dict, tools: Sequence[Any] | None = None, ) -> FinalContext: - return self.context_manager.assemble_final_context( - model=model, - memory=memory, - current_run_start_idx=current_run_start_idx, - tools=tools, - purpose="final_answer", - task=task, - final_answer_templates=final_answer_templates, - run_context=self._ensure_run_context(memory), - ) + monitoring_manager = get_monitoring_manager() + with monitoring_manager.trace_operation( + "context.prepare_final_answer", + OPENINFERENCE_SPAN_KIND_CHAIN, + **{ + "context.current_run_start_idx": current_run_start_idx, + "context.conversation_id": self.conversation_id, + "context.task": task, + OPENINFERENCE_INPUT_VALUE: { + "task": task[:200] if task else "", + "current_run_start_idx": current_run_start_idx, + "memory_steps": len(memory.steps) if hasattr(memory, 'steps') else 0, + "tool_count": len(tools) if tools else 0, + "conversation_id": self.conversation_id, + }, + }, + ): + result = self.context_manager.assemble_final_context( + model=model, + memory=memory, + current_run_start_idx=current_run_start_idx, + tools=tools, + purpose="final_answer", + task=task, + final_answer_templates=final_answer_templates, + run_context=self._ensure_run_context(memory), + conversation_id=self.conversation_id, + ) + if monitoring_manager.is_enabled: + monitoring_manager.set_openinference_output({ + "message_count": len(result.messages), + "tool_count": len(result.tools), + "stable_count": result.evidence.stable_message_count, + "dynamic_count": result.evidence.dynamic_message_count, + "context_items": len(result.evidence.context_items), + }) + return result def render_summary_messages(self, *, memory: Any) -> list[Any]: """Return display-only memory messages without compression side effects.""" diff --git a/sdk/nexent/monitor/monitoring.py b/sdk/nexent/monitor/monitoring.py index e9381665d..3caf21048 100644 --- a/sdk/nexent/monitor/monitoring.py +++ b/sdk/nexent/monitor/monitoring.py @@ -936,6 +936,21 @@ def trace_operation( span_attrs = { OPENINFERENCE_SPAN_KIND: span_kind, } + + # Process input.value and output.value through payload preview (same as trace_llm_request) + input_value = attributes.pop(OPENINFERENCE_INPUT_VALUE, None) + output_value = attributes.pop(OPENINFERENCE_OUTPUT_VALUE, None) + if input_value is not None: + input_preview = self._trace_payload_preview(input_value) + if input_preview != "": + span_attrs[OPENINFERENCE_INPUT_VALUE] = input_preview + span_attrs.update(self._trace_payload_attributes("input", input_value)) + if output_value is not None: + output_preview = self._trace_payload_preview(output_value) + if output_preview != "": + span_attrs[OPENINFERENCE_OUTPUT_VALUE] = output_preview + span_attrs.update(self._trace_payload_attributes("output", output_value)) + span_attrs.update(attributes) with self._tracer.start_as_current_span( diff --git a/test/backend/agents/test_create_agent_info.py b/test/backend/agents/test_create_agent_info.py index 940a26c56..7bcf40cd0 100644 --- a/test/backend/agents/test_create_agent_info.py +++ b/test/backend/agents/test_create_agent_info.py @@ -1907,7 +1907,8 @@ async def test_create_agent_config_basic(self): context_components=ANY, capacity_snapshot=ANY, safe_input_budget_snapshot=ANY, - verification_config=ANY + verification_config=ANY, + conversation_id=None ) @pytest.mark.asyncio @@ -1982,7 +1983,8 @@ async def test_create_agent_config_with_sub_agents(self): context_components=ANY, capacity_snapshot=ANY, safe_input_budget_snapshot=ANY, - verification_config=ANY + verification_config=ANY, + conversation_id=None ) @pytest.mark.asyncio @@ -2244,7 +2246,8 @@ async def test_create_agent_config_model_id_none(self): context_components=ANY, capacity_snapshot=None, safe_input_budget_snapshot=None, - verification_config=ANY + verification_config=ANY, + conversation_id=None ) @pytest.mark.asyncio @@ -3394,6 +3397,7 @@ async def test_create_agent_run_info_success(self): allow_memory_search=True, version_no=1, tool_params=None, + conversation_id=None, ) mock_get_mcp.assert_called_once_with(tenant_id="tenant_1", is_need_auth=True) mock_filter.assert_called_once_with("agent_config", { @@ -3931,6 +3935,7 @@ async def test_create_agent_run_info_forwards_allow_memory_false(self): allow_memory_search=False, version_no=1, tool_params=None, + conversation_id=None, ) @pytest.mark.asyncio @@ -3978,6 +3983,7 @@ async def test_create_agent_run_info_is_debug_true(self): allow_memory_search=True, version_no=0, # Debug mode uses draft version 0 tool_params=None, + conversation_id=None, ) @pytest.mark.asyncio @@ -4031,6 +4037,7 @@ async def test_create_agent_run_info_no_published_version_fallback(self): allow_memory_search=True, version_no=0, # Fallback to draft version 0 tool_params=None, + conversation_id=None, ) # Verify that get_remote_mcp_server_list was called with is_need_auth=True mock_get_mcp.assert_called_once_with(tenant_id="tenant_1", is_need_auth=True) diff --git a/test/backend/database/test_conversation_db.py b/test/backend/database/test_conversation_db.py index fc86c8b06..22ee6d56a 100644 --- a/test/backend/database/test_conversation_db.py +++ b/test/backend/database/test_conversation_db.py @@ -129,6 +129,7 @@ class ConversationMessageUnit: conversation_id = MagicMock(name="ConversationMessageUnit.conversation_id") delete_flag = MagicMock(name="ConversationMessageUnit.delete_flag") unit_status = MagicMock(name="ConversationMessageUnit.unit_status") + step_index = MagicMock(name="ConversationMessageUnit.step_index") class ConversationSourceSearch: @@ -200,6 +201,7 @@ def _add_update_tracking(data, user_id): get_message, get_message_id_by_index, get_message_units, + get_message_units_by_message, get_source_images_by_conversation, get_source_images_by_message, get_source_searches_by_conversation, @@ -2292,3 +2294,97 @@ def as_dict_side_effect(record): assert result is not None assert result['message_records'][0]['units'] == [] + + +# ============================================================================= +# Tests for get_message_units_by_message +# ============================================================================= + + +def test_get_message_units_by_message_without_message_id(monkeypatch, mock_session_ctx): + """get_message_units_by_message returns all units when message_id is None.""" + session, ctx = mock_session_ctx + mock_records = [MagicMock(), MagicMock(), MagicMock()] + session.scalars.return_value.all.return_value = mock_records + + def as_dict_side_effect(record): + return {"unit_id": id(record)} + + monkeypatch.setattr("backend.database.conversation_db.get_db_session", lambda: ctx) + monkeypatch.setattr("backend.database.conversation_db.as_dict", as_dict_side_effect) + + result = get_message_units_by_message(42) + + assert len(result) == 3 + session.scalars.assert_called_once() + + +def test_get_message_units_by_message_with_message_id(monkeypatch, mock_session_ctx): + """get_message_units_by_message filters by message_id when provided.""" + session, ctx = mock_session_ctx + mock_records = [MagicMock(), MagicMock()] + session.scalars.return_value.all.return_value = mock_records + + def as_dict_side_effect(record): + return {"unit_id": id(record), "message_id": 5} + + monkeypatch.setattr("backend.database.conversation_db.get_db_session", lambda: ctx) + monkeypatch.setattr("backend.database.conversation_db.as_dict", as_dict_side_effect) + + result = get_message_units_by_message(42, message_id=5) + + assert len(result) == 2 + session.scalars.assert_called_once() + + +def test_get_message_units_by_message_empty_result(monkeypatch, mock_session_ctx): + """get_message_units_by_message returns empty list when no units found.""" + session, ctx = mock_session_ctx + session.scalars.return_value.all.return_value = [] + + monkeypatch.setattr("backend.database.conversation_db.get_db_session", lambda: ctx) + + result = get_message_units_by_message(42, message_id=99) + + assert result == [] + + +def test_get_message_units_by_message_string_conversation_id(monkeypatch, mock_session_ctx): + """get_message_units_by_message handles conversation_id passed as string.""" + session, ctx = mock_session_ctx + mock_records = [MagicMock()] + session.scalars.return_value.all.return_value = mock_records + + def as_dict_side_effect(record): + return {"unit_id": 1} + + monkeypatch.setattr("backend.database.conversation_db.get_db_session", lambda: ctx) + monkeypatch.setattr("backend.database.conversation_db.as_dict", as_dict_side_effect) + + result = get_message_units_by_message(42) + + assert len(result) == 1 + + +def test_get_message_units_by_message_returns_dicts(monkeypatch, mock_session_ctx): + """get_message_units_by_message converts records to dicts via as_dict.""" + session, ctx = mock_session_ctx + mock_record = MagicMock(unit_id=10, unit_type="final_answer", unit_content="Hello") + session.scalars.return_value.all.return_value = [mock_record] + + def as_dict_side_effect(record): + return { + "unit_id": record.unit_id, + "unit_type": record.unit_type, + "unit_content": record.unit_content, + } + + monkeypatch.setattr("backend.database.conversation_db.get_db_session", lambda: ctx) + monkeypatch.setattr("backend.database.conversation_db.as_dict", as_dict_side_effect) + + result = get_message_units_by_message(1, message_id=2) + + assert len(result) == 1 + assert result[0]["unit_id"] == 10 + assert result[0]["unit_type"] == "final_answer" + assert result[0]["unit_content"] == "Hello" diff --git a/test/backend/services/test_agent_service.py b/test/backend/services/test_agent_service.py index 50a0ef32b..15a64be42 100644 --- a/test/backend/services/test_agent_service.py +++ b/test/backend/services/test_agent_service.py @@ -33,6 +33,8 @@ def model_dump(self, **kwargs): sys.modules['nexent.core.agents'] = MagicMock() sys.modules['nexent.core.agents.agent_model'] = nexent_agent_model_mock sys.modules['nexent.core.agents.run_agent'] = MagicMock() +sys.modules['nexent.core.agents.context'] = MagicMock() +sys.modules['nexent.core.agents.context.history_projector'] = MagicMock() # Mock other nexent submodules sys.modules['nexent.memory'] = MagicMock() @@ -209,6 +211,12 @@ class MODEL_OUTPUT_THINKING: value = "model_output_thinking" class MODEL_OUTPUT_DEEP_THINKING: value = "model_output_deep_thinking" + class STEP_COUNT: + value = "step_count" + class TOOL: + value = "tool" + class EXECUTION_LOGS: + value = "execution_logs" sys.modules['nexent.core.utils.observer'] = MagicMock() sys.modules['nexent.core.utils.observer'].ProcessType = MockProcessType @@ -3850,6 +3858,7 @@ async def test_prepare_agent_run( override_model_id=None, requested_output_tokens=4096, tool_params=None, + conversation_id=123, ) mock_agent_run_manager.register_agent_run.assert_called_once_with( 123, mock_run_info, "test_user") @@ -4448,8 +4457,8 @@ async def fake_agent_run(*_, **__): # Track save_message calls to verify streaming message creation save_message_calls = [] - def fake_save_message(req, user_id, tenant_id, status="completed"): - save_message_calls.append((req, user_id, tenant_id, status)) + def fake_save_message(req, user_id, tenant_id, status="completed", **kwargs): + save_message_calls.append((req, user_id, tenant_id, status, kwargs)) return 4242 monkeypatch.setattr( @@ -4645,9 +4654,14 @@ def capture_task(coro): ): collected.append(out) + # Give the finally block time to create and execute the background task + await asyncio.sleep(0.1) + # Ensure background task completed if task_holder["task"] is not None: await task_holder["task"] + # Give the task time to complete after awaiting + await asyncio.sleep(0.01) assert add_calls["called"] is True assert add_calls["args"]["messages"] == [ @@ -13842,3 +13856,301 @@ async def mock_title_gen(*args, **kwargs): # Should complete successfully assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_stream_agent_chunks_tool_call_merge(monkeypatch): + """TOOL + EXECUTION_LOGS chunks should be merged into a single tool_call row.""" + from backend.services import agent_service + + agent_request = MagicMock() + agent_request.agent_id = 1 + agent_request.conversation_id = 999 + agent_request.query = "test" + agent_request.history = [] + agent_request.minio_files = [] + agent_request.is_debug = False + + async def fake_agent_run(*_, **__): + yield json.dumps({"type": "step_count", "content": "**Step 1**"}) + yield json.dumps({"type": "tool", "content": "search('query')"}) + yield json.dumps({"type": "execution_logs", "content": "result: found 3 items"}) + yield json.dumps({"type": "final_answer", "content": "Done"}) + + monkeypatch.setattr( + "backend.services.agent_service.agent_run", fake_agent_run, raising=False + ) + + monkeypatch.setattr( + "backend.services.agent_service.save_message", + lambda *a, **kw: 4242, + raising=False, + ) + + saved_units = [] + + class FakeFuture: + def __init__(self, unit_id): + self._unit_id = unit_id + def result(self): + return self._unit_id + + def fake_submit(fn, *args, **kwargs): + if "unit_type" in kwargs: + unit_id = len(saved_units) + 100 + saved_units.append({ + "unit_type": kwargs.get("unit_type"), + "unit_content": kwargs.get("unit_content"), + "step_index": kwargs.get("step_index"), + }) + return FakeFuture(unit_id) + return MagicMock() + + monkeypatch.setattr( + "backend.services.agent_service.submit", fake_submit, raising=False + ) + + monkeypatch.setattr( + "backend.services.agent_service.update_unit_status", + lambda *a, **kw: None, + raising=False, + ) + monkeypatch.setattr( + "backend.services.agent_service.update_unit_content", + lambda *a, **kw: None, + raising=False, + ) + monkeypatch.setattr( + "backend.services.agent_service.update_message_content", + lambda *a, **kw: None, + raising=False, + ) + monkeypatch.setattr( + "backend.services.agent_service.update_message_status", + lambda *a, **kw: None, + raising=False, + ) + monkeypatch.setattr( + "backend.services.agent_service.agent_run_manager.unregister_agent_run", + lambda *a, **kw: None, + raising=False, + ) + monkeypatch.setattr( + "backend.services.agent_service.generate_conversation_title_service", + AsyncMock(), + raising=False, + ) + + collected = [] + async for out in agent_service._stream_agent_chunks( + agent_request, "u", "t", MagicMock(), MagicMock() + ): + collected.append(out) + + tool_call_units = [u for u in saved_units if u["unit_type"] == "tool_call"] + assert len(tool_call_units) == 1, f"Expected 1 tool_call unit, got {len(tool_call_units)}: {saved_units}" + + merged = json.loads(tool_call_units[0]["unit_content"]) + assert merged["tool_call"] == "search('query')" + assert merged["execution_result"] == "result: found 3 items" + + standalone_tools = [u for u in saved_units if u["unit_type"] == "tool"] + standalone_logs = [u for u in saved_units if u["unit_type"] == "execution_logs"] + assert len(standalone_tools) == 0 + assert len(standalone_logs) == 0 + + +@pytest.mark.asyncio +async def test_stream_agent_chunks_orphaned_tool_flush(monkeypatch): + """Orphaned TOOL chunk (no EXECUTION_LOGS) should be flushed as standalone 'tool' row at end of stream.""" + from backend.services import agent_service + + agent_request = MagicMock() + agent_request.agent_id = 1 + agent_request.conversation_id = 999 + agent_request.query = "test" + agent_request.history = [] + agent_request.minio_files = [] + agent_request.is_debug = False + + async def fake_agent_run(*_, **__): + yield json.dumps({"type": "step_count", "content": "**Step 1**"}) + yield json.dumps({"type": "tool", "content": "search('query')"}) + + monkeypatch.setattr( + "backend.services.agent_service.agent_run", fake_agent_run, raising=False + ) + + monkeypatch.setattr( + "backend.services.agent_service.save_message", + lambda *a, **kw: 4242, + raising=False, + ) + + saved_units = [] + + class FakeFuture: + def __init__(self, unit_id): + self._unit_id = unit_id + def result(self): + return self._unit_id + + def fake_submit(fn, *args, **kwargs): + if "unit_type" in kwargs: + unit_id = len(saved_units) + 100 + saved_units.append({ + "unit_type": kwargs.get("unit_type"), + "unit_content": kwargs.get("unit_content"), + }) + return FakeFuture(unit_id) + return MagicMock() + + monkeypatch.setattr( + "backend.services.agent_service.submit", fake_submit, raising=False + ) + + monkeypatch.setattr( + "backend.services.agent_service.update_unit_status", + lambda *a, **kw: None, + raising=False, + ) + monkeypatch.setattr( + "backend.services.agent_service.update_unit_content", + lambda *a, **kw: None, + raising=False, + ) + monkeypatch.setattr( + "backend.services.agent_service.update_message_content", + lambda *a, **kw: None, + raising=False, + ) + monkeypatch.setattr( + "backend.services.agent_service.update_message_status", + lambda *a, **kw: None, + raising=False, + ) + monkeypatch.setattr( + "backend.services.agent_service.agent_run_manager.unregister_agent_run", + lambda *a, **kw: None, + raising=False, + ) + monkeypatch.setattr( + "backend.services.agent_service.generate_conversation_title_service", + AsyncMock(), + raising=False, + ) + + collected = [] + async for out in agent_service._stream_agent_chunks( + agent_request, "u", "t", MagicMock(), MagicMock() + ): + collected.append(out) + + standalone_tools = [u for u in saved_units if u["unit_type"] == "tool"] + assert len(standalone_tools) == 1, f"Expected 1 standalone tool, got: {saved_units}" + assert standalone_tools[0]["unit_content"] == "search('query')" + + tool_call_units = [u for u in saved_units if u["unit_type"] == "tool_call"] + assert len(tool_call_units) == 0 + + +@pytest.mark.asyncio +async def test_stream_agent_chunks_multiple_tool_calls(monkeypatch): + """Multiple TOOL + EXECUTION_LOGS pairs should each produce one tool_call row.""" + from backend.services import agent_service + + agent_request = MagicMock() + agent_request.agent_id = 1 + agent_request.conversation_id = 999 + agent_request.query = "test" + agent_request.history = [] + agent_request.minio_files = [] + agent_request.is_debug = False + + async def fake_agent_run(*_, **__): + yield json.dumps({"type": "step_count", "content": "**Step 1**"}) + yield json.dumps({"type": "tool", "content": "search('a')"}) + yield json.dumps({"type": "execution_logs", "content": "result a"}) + yield json.dumps({"type": "tool", "content": "search('b')"}) + yield json.dumps({"type": "execution_logs", "content": "result b"}) + yield json.dumps({"type": "final_answer", "content": "Done"}) + + monkeypatch.setattr( + "backend.services.agent_service.agent_run", fake_agent_run, raising=False + ) + + monkeypatch.setattr( + "backend.services.agent_service.save_message", + lambda *a, **kw: 4242, + raising=False, + ) + + saved_units = [] + + class FakeFuture: + def __init__(self, unit_id): + self._unit_id = unit_id + def result(self): + return self._unit_id + + def fake_submit(fn, *args, **kwargs): + if "unit_type" in kwargs: + unit_id = len(saved_units) + 100 + saved_units.append({ + "unit_type": kwargs.get("unit_type"), + "unit_content": kwargs.get("unit_content"), + }) + return FakeFuture(unit_id) + return MagicMock() + + monkeypatch.setattr( + "backend.services.agent_service.submit", fake_submit, raising=False + ) + + monkeypatch.setattr( + "backend.services.agent_service.update_unit_status", + lambda *a, **kw: None, + raising=False, + ) + monkeypatch.setattr( + "backend.services.agent_service.update_unit_content", + lambda *a, **kw: None, + raising=False, + ) + monkeypatch.setattr( + "backend.services.agent_service.update_message_content", + lambda *a, **kw: None, + raising=False, + ) + monkeypatch.setattr( + "backend.services.agent_service.update_message_status", + lambda *a, **kw: None, + raising=False, + ) + monkeypatch.setattr( + "backend.services.agent_service.agent_run_manager.unregister_agent_run", + lambda *a, **kw: None, + raising=False, + ) + monkeypatch.setattr( + "backend.services.agent_service.generate_conversation_title_service", + AsyncMock(), + raising=False, + ) + + collected = [] + async for out in agent_service._stream_agent_chunks( + agent_request, "u", "t", MagicMock(), MagicMock() + ): + collected.append(out) + + tool_call_units = [u for u in saved_units if u["unit_type"] == "tool_call"] + assert len(tool_call_units) == 2, f"Expected 2 tool_call units, got {len(tool_call_units)}: {saved_units}" + + merged_a = json.loads(tool_call_units[0]["unit_content"]) + assert merged_a["tool_call"] == "search('a')" + assert merged_a["execution_result"] == "result a" + + merged_b = json.loads(tool_call_units[1]["unit_content"]) + assert merged_b["tool_call"] == "search('b')" + assert merged_b["execution_result"] == "result b" diff --git a/test/backend/services/test_conversation_management_service.py b/test/backend/services/test_conversation_management_service.py index cb3ef3f98..dcdfbfc94 100644 --- a/test/backend/services/test_conversation_management_service.py +++ b/test/backend/services/test_conversation_management_service.py @@ -324,6 +324,7 @@ def test_save_message_unit_inserts_single_row(self, mock_create_message_unit): unit_content="print('hi')", user_id=self.user_id, unit_status="streaming", + step_index=None, ) @patch('backend.services.conversation_management_service.create_source_image') diff --git a/test/pytest.ini b/test/pytest.ini index 21e178bdd..587fb42d2 100644 --- a/test/pytest.ini +++ b/test/pytest.ini @@ -8,3 +8,7 @@ asyncio_default_fixture_loop_scope = function filterwarnings = # Disable all warnings ignore + +# Custom markers +markers = + local_only: marks tests that require external resources (API keys, network) and should only run locally (deselect with '-m "not local_only"') diff --git a/test/sdk/core/agents/context/test_context_item.py b/test/sdk/core/agents/context/test_context_item.py new file mode 100644 index 000000000..cfd0b3e78 --- /dev/null +++ b/test/sdk/core/agents/context/test_context_item.py @@ -0,0 +1,127 @@ +"""Tests for context item data models and enumerations.""" + +from dataclasses import fields + +from nexent.core.agents.context.context_item import ( + AuthorityTier, + ContextItem, + ContextItemType, + RepresentationTier, +) + + +class TestContextItemTypeEnum: + """Tests for ContextItemType enumeration.""" + + def test_context_item_type_enum_has_all_values(self): + expected = { + "SYSTEM_PROMPT", + "TOOL", + "SKILL", + "MEMORY", + "KNOWLEDGE_BASE", + "MANAGED_AGENT", + "EXTERNAL_AGENT", + "HISTORY_TURN", + "TOOL_CALL_RESULT", + } + actual = {member.name for member in ContextItemType} + assert actual == expected + + def test_context_item_type_is_string_enum(self): + for member in ContextItemType: + assert isinstance(member.value, str) + + +class TestRepresentationTierEnum: + """Tests for RepresentationTier enumeration.""" + + def test_representation_tier_enum_has_all_values(self): + expected = {"FULL", "COMPRESSED", "STRUCTURED", "POINTER"} + actual = {member.name for member in RepresentationTier} + assert actual == expected + + +class TestAuthorityTierEnum: + """Tests for AuthorityTier enumeration.""" + + def test_authority_tier_enum_has_all_values(self): + expected = { + "PLATFORM", + "TENANT", + "USER", + "TOOL_RESULT", + "RETRIEVED_MEMORY", + "SUMMARY", + "AGENT_INFERENCE", + } + actual = {member.name for member in AuthorityTier} + assert actual == expected + + +class TestContextItem: + """Tests for ContextItem dataclass.""" + + def test_context_item_default_values(self): + item = ContextItem(item_id="test-1", item_type=ContextItemType.TOOL) + + assert item.item_id == "test-1" + assert item.item_type == ContextItemType.TOOL + assert item.source_refs == [] + assert item.authority_tier == AuthorityTier.AGENT_INFERENCE + assert item.minimum_fidelity == RepresentationTier.STRUCTURED + assert item.current_representation == RepresentationTier.FULL + assert item.content is None + assert item.token_estimate == 0 + assert item.metadata == {} + assert item.lifecycle_status == "active" + assert item.recompute_cost is None + + def test_context_item_custom_values(self): + item = ContextItem( + item_id="custom-1", + item_type=ContextItemType.MEMORY, + source_refs=["ref-1", "ref-2"], + authority_tier=AuthorityTier.USER, + minimum_fidelity=RepresentationTier.FULL, + current_representation=RepresentationTier.COMPRESSED, + content="some memory content", + token_estimate=150, + metadata={"key": "value"}, + lifecycle_status="archived", + recompute_cost=50, + ) + + assert item.item_id == "custom-1" + assert item.item_type == ContextItemType.MEMORY + assert item.source_refs == ["ref-1", "ref-2"] + assert item.authority_tier == AuthorityTier.USER + assert item.minimum_fidelity == RepresentationTier.FULL + assert item.current_representation == RepresentationTier.COMPRESSED + assert item.content == "some memory content" + assert item.token_estimate == 150 + assert item.metadata == {"key": "value"} + assert item.lifecycle_status == "archived" + assert item.recompute_cost == 50 + + def test_context_item_serialization(self): + """Verify dataclass fields round-trip correctly.""" + original = ContextItem( + item_id="serial-1", + item_type=ContextItemType.SKILL, + source_refs=["src-a"], + authority_tier=AuthorityTier.PLATFORM, + minimum_fidelity=RepresentationTier.POINTER, + current_representation=RepresentationTier.STRUCTURED, + content={"skill": "data"}, + token_estimate=200, + metadata={"version": 2}, + lifecycle_status="active", + recompute_cost=10, + ) + + field_values = {f.name: getattr(original, f.name) for f in fields(original)} + reconstructed = ContextItem(**field_values) + + for f in fields(original): + assert getattr(original, f.name) == getattr(reconstructed, f.name) diff --git a/test/sdk/core/agents/context/test_handlers.py b/test/sdk/core/agents/context/test_handlers.py new file mode 100644 index 000000000..78cbae34f --- /dev/null +++ b/test/sdk/core/agents/context/test_handlers.py @@ -0,0 +1,141 @@ +"""Tests for all context item handler implementations.""" + +import math + +import pytest + +from nexent.core.agents.context.context_item import ( + AuthorityTier, + ContextItem, + ContextItemType, + RepresentationTier, +) +from nexent.core.agents.context.reducer_models import ReductionResult + +from nexent.core.agents.context.handlers import ( + ExternalAgentHandler, + HistoryTurnHandler, + KnowledgeBaseHandler, + ManagedAgentHandler, + MemoryHandler, + SkillHandler, + SystemPromptHandler, + ToolCallResultHandler, + ToolHandler, +) + + +ALL_HANDLERS = [ + SystemPromptHandler, + ToolHandler, + SkillHandler, + MemoryHandler, + KnowledgeBaseHandler, + ManagedAgentHandler, + ExternalAgentHandler, + HistoryTurnHandler, + ToolCallResultHandler, +] + +MANDATORY_HANDLERS = [SystemPromptHandler] + +NON_MANDATORY_HANDLERS = [ + ToolHandler, + SkillHandler, + MemoryHandler, + KnowledgeBaseHandler, + ManagedAgentHandler, + ExternalAgentHandler, + HistoryTurnHandler, + ToolCallResultHandler, +] + + +def _make_item(handler_cls): + """Create a ContextItem matching the handler's first supported type.""" + handler = handler_cls() + item_type = handler.supported_types()[0] + return ContextItem( + item_id=f"test-{item_type.value}", + item_type=item_type, + content=f"content for {item_type.value}", + token_estimate=100, + ) + + +class TestHandlerSupportedTypes: + """Tests for handler supported_types() method.""" + + @pytest.mark.parametrize("handler_cls", ALL_HANDLERS, ids=lambda h: h.__name__) + def test_handler_supported_types(self, handler_cls): + handler = handler_cls() + types = handler.supported_types() + assert len(types) > 0 + for t in types: + assert isinstance(t, ContextItemType) + + +class TestHandlerScore: + """Tests for handler score() method.""" + + @pytest.mark.parametrize("handler_cls", ALL_HANDLERS, ids=lambda h: h.__name__) + def test_handler_score_returns_float(self, handler_cls): + handler = handler_cls() + item = _make_item(handler_cls) + result = handler.score(item, "test query", {}) + assert isinstance(result, float) + + @pytest.mark.parametrize("handler_cls", MANDATORY_HANDLERS, ids=lambda h: h.__name__) + def test_mandatory_handlers_return_inf_score(self, handler_cls): + handler = handler_cls() + item = _make_item(handler_cls) + result = handler.score(item, "test query", {}) + assert math.isinf(result) + + @pytest.mark.parametrize("handler_cls", NON_MANDATORY_HANDLERS, ids=lambda h: h.__name__) + def test_non_mandatory_handlers_return_1_0_score(self, handler_cls): + handler = handler_cls() + item = _make_item(handler_cls) + result = handler.score(item, "test query", {}) + assert result == pytest.approx(1.0, rel=1e-9) + + +class TestHandlerReduce: + """Tests for handler reduce() method.""" + + @pytest.mark.parametrize("handler_cls", ALL_HANDLERS, ids=lambda h: h.__name__) + def test_handler_reduce_returns_reduction_result(self, handler_cls): + handler = handler_cls() + item = _make_item(handler_cls) + result = handler.reduce(item, RepresentationTier.FULL, 1000) + assert isinstance(result, ReductionResult) + + @pytest.mark.parametrize("handler_cls", ALL_HANDLERS, ids=lambda h: h.__name__) + def test_handler_reduce_passthrough_preserves_content(self, handler_cls): + handler = handler_cls() + item = _make_item(handler_cls) + result = handler.reduce(item, RepresentationTier.FULL, 1000) + assert result.content == item.content + assert result.admissible is True + + +class TestHandlerCoverage: + """Tests for handler type coverage and disjointness.""" + + def test_all_handler_supported_types_are_disjoint(self): + seen_types = set() + for handler_cls in ALL_HANDLERS: + handler = handler_cls() + for t in handler.supported_types(): + assert t not in seen_types, ( + f"{handler_cls.__name__} duplicates type {t.name}" + ) + seen_types.add(t) + + def test_all_context_item_types_covered_by_handlers(self): + covered = set() + for handler_cls in ALL_HANDLERS: + handler = handler_cls() + covered.update(handler.supported_types()) + + assert covered == set(ContextItemType) diff --git a/test/sdk/core/agents/context/test_item_handler_registry.py b/test/sdk/core/agents/context/test_item_handler_registry.py new file mode 100644 index 000000000..888968f3e --- /dev/null +++ b/test/sdk/core/agents/context/test_item_handler_registry.py @@ -0,0 +1,109 @@ +"""Tests for ItemHandlerRegistry.""" + +import pytest + +from nexent.core.agents.context.context_item import ( + AuthorityTier, + ContextItem, + ContextItemType, + RepresentationTier, +) +from nexent.core.agents.context.item_handler import ContextItemHandler +from nexent.core.agents.context.item_handler_registry import ItemHandlerRegistry +from nexent.core.agents.context.reducer_models import ReductionResult + +from nexent.core.agents.context.handlers import ( + ExternalAgentHandler, + HistoryTurnHandler, + KnowledgeBaseHandler, + ManagedAgentHandler, + MemoryHandler, + SkillHandler, + SystemPromptHandler, + ToolCallResultHandler, + ToolHandler, +) + + +@pytest.fixture(autouse=True) +def reset_registry(): + """Reset the shared class-level registry before each test.""" + ItemHandlerRegistry.reset() + yield + ItemHandlerRegistry.reset() + + +class _StubHandler(ContextItemHandler): + """Minimal handler for testing registry operations.""" + + def __init__(self, types): + self._types = types + + def supported_types(self): + return self._types + + +class TestItemHandlerRegistry: + """Tests for ItemHandlerRegistry class methods.""" + + def test_register_and_get_handler(self): + handler = _StubHandler([ContextItemType.TOOL]) + ItemHandlerRegistry.register(handler) + + retrieved = ItemHandlerRegistry.get(ContextItemType.TOOL) + assert retrieved is handler + + def test_get_unregistered_type_raises_key_error(self): + with pytest.raises(KeyError, match="No handler registered"): + ItemHandlerRegistry.get(ContextItemType.SKILL) + + def test_all_types_covered_returns_false_when_empty(self): + assert ItemHandlerRegistry.all_types_covered() is False + + def test_all_types_covered_returns_true_when_complete(self): + all_handler_classes = [ + SystemPromptHandler, + ToolHandler, + SkillHandler, + MemoryHandler, + KnowledgeBaseHandler, + ManagedAgentHandler, + ExternalAgentHandler, + HistoryTurnHandler, + ToolCallResultHandler, + ] + for handler_cls in all_handler_classes: + ItemHandlerRegistry.register(handler_cls()) + + assert ItemHandlerRegistry.all_types_covered() is True + + def test_reset_clears_handlers(self): + handler = _StubHandler([ContextItemType.TOOL]) + ItemHandlerRegistry.register(handler) + assert ItemHandlerRegistry.get(ContextItemType.TOOL) is handler + + ItemHandlerRegistry.reset() + + with pytest.raises(KeyError): + ItemHandlerRegistry.get(ContextItemType.TOOL) + + def test_register_all_handlers_covers_all_types(self): + """Register all built-in handlers and verify complete coverage.""" + all_handler_classes = [ + SystemPromptHandler, + ToolHandler, + SkillHandler, + MemoryHandler, + KnowledgeBaseHandler, + ManagedAgentHandler, + ExternalAgentHandler, + HistoryTurnHandler, + ToolCallResultHandler, + ] + for handler_cls in all_handler_classes: + ItemHandlerRegistry.register(handler_cls()) + + assert ItemHandlerRegistry.all_types_covered() is True + + assert isinstance(ItemHandlerRegistry.get(ContextItemType.SYSTEM_PROMPT), SystemPromptHandler) + assert isinstance(ItemHandlerRegistry.get(ContextItemType.TOOL), ToolHandler) diff --git a/test/sdk/core/agents/context/test_policy_models.py b/test/sdk/core/agents/context/test_policy_models.py new file mode 100644 index 000000000..ce8296901 --- /dev/null +++ b/test/sdk/core/agents/context/test_policy_models.py @@ -0,0 +1,110 @@ +"""Tests for policy decision models.""" + +import pytest + +from nexent.core.agents.context.context_item import ( + AuthorityTier, + ContextItem, + ContextItemType, + RepresentationTier, +) +from nexent.core.agents.context.policy_models import MemoryDecision, SelectionDecision + + +class TestSelectionDecision: + """Tests for SelectionDecision frozen dataclass.""" + + def test_selection_decision_is_frozen(self): + decision = SelectionDecision( + selected_item_ids=["item-1"], + excluded_item_ids=["item-2"], + representation_requirements={"item-1": RepresentationTier.FULL}, + budget_allocations={"item-1": 100}, + remaining_budget=500, + conflicts=[], + reason_codes=["selected_mandatory_minimum"], + policy_version="1.0", + decision_fingerprint="abc123", + ) + + with pytest.raises(AttributeError): + decision.selected_item_ids = ["other"] + + def test_selection_decision_creation(self): + decision = SelectionDecision( + selected_item_ids=["item-1", "item-2"], + excluded_item_ids=["item-3"], + representation_requirements={ + "item-1": RepresentationTier.FULL, + "item-2": RepresentationTier.COMPRESSED, + }, + budget_allocations={"item-1": 200, "item-2": 100}, + remaining_budget=300, + conflicts=[{"type": "budget", "items": ["item-1", "item-3"]}], + reason_codes=["selected_mandatory_minimum", "selected_budget_upgrade"], + policy_version="2.0", + decision_fingerprint="def456", + ) + + assert decision.selected_item_ids == ["item-1", "item-2"] + assert decision.excluded_item_ids == ["item-3"] + assert decision.remaining_budget == 300 + assert decision.policy_version == "2.0" + assert decision.decision_fingerprint == "def456" + assert len(decision.conflicts) == 1 + assert len(decision.reason_codes) == 2 + + def test_selection_decision_with_empty_lists(self): + decision = SelectionDecision( + selected_item_ids=[], + excluded_item_ids=[], + representation_requirements={}, + budget_allocations={}, + remaining_budget=1000, + conflicts=[], + reason_codes=[], + policy_version="1.0", + decision_fingerprint="empty", + ) + + assert decision.selected_item_ids == [] + assert decision.excluded_item_ids == [] + assert decision.representation_requirements == {} + assert decision.budget_allocations == {} + assert decision.conflicts == [] + assert decision.reason_codes == [] + + +class TestMemoryDecision: + """Tests for MemoryDecision frozen dataclass.""" + + def test_memory_decision_is_frozen(self): + decision = MemoryDecision( + operation="read", + allowed_scopes=["user"], + excluded_candidates=[], + conflict_decisions=[], + confirmation_required=None, + reason_codes=["memory_operation_allowed"], + ) + + with pytest.raises(AttributeError): + decision.operation = "write" + + def test_memory_decision_creation(self): + decision = MemoryDecision( + operation="write", + allowed_scopes=["user", "agent"], + excluded_candidates=["candidate-1"], + conflict_decisions=[{"conflict": "scope_overlap", "resolution": "merge"}], + confirmation_required={"reason": "sensitive_data", "scope": "tenant"}, + reason_codes=["memory_operation_allowed", "confirmation_required"], + ) + + assert decision.operation == "write" + assert decision.allowed_scopes == ["user", "agent"] + assert decision.excluded_candidates == ["candidate-1"] + assert len(decision.conflict_decisions) == 1 + assert decision.confirmation_required is not None + assert decision.confirmation_required["reason"] == "sensitive_data" + assert len(decision.reason_codes) == 2 diff --git a/test/sdk/core/agents/context/test_reason_codes.py b/test/sdk/core/agents/context/test_reason_codes.py new file mode 100644 index 000000000..8e2e794df --- /dev/null +++ b/test/sdk/core/agents/context/test_reason_codes.py @@ -0,0 +1,40 @@ +"""Tests for reason code constants.""" + +from nexent.core.agents.context import reason_codes + + +EXPECTED_CODE_NAMES = [ + "SELECTED_MANDATORY_MINIMUM", + "SELECTED_BUDGET_UPGRADE", + "EXCLUDED_BUDGET", + "EXCLUDED_POLICY_DISABLED", + "EXCLUDED_LOWER_AUTHORITY", + "MEMORY_OPERATION_ALLOWED", + "MEMORY_OPERATION_DENIED", + "CONFIRMATION_REQUIRED", + "MINIMUM_FIDELITY_VIOLATION", + "REDUCER_FAILED", + "REPRESENTATION_STALE", +] + + +class TestReasonCodes: + """Tests for reason code string constants.""" + + def test_expected_reason_codes_exist(self): + for name in EXPECTED_CODE_NAMES: + assert hasattr(reason_codes, name), f"Missing reason code: {name}" + + def test_all_reason_codes_are_strings(self): + for name in EXPECTED_CODE_NAMES: + value = getattr(reason_codes, name) + assert isinstance(value, str), f"{name} is not a string" + + def test_all_reason_codes_are_non_empty(self): + for name in EXPECTED_CODE_NAMES: + value = getattr(reason_codes, name) + assert len(value) > 0, f"{name} is an empty string" + + def test_all_reason_codes_are_unique(self): + values = [getattr(reason_codes, name) for name in EXPECTED_CODE_NAMES] + assert len(values) == len(set(values)), "Duplicate reason code values found" diff --git a/test/sdk/core/agents/context/test_reducer_models.py b/test/sdk/core/agents/context/test_reducer_models.py new file mode 100644 index 000000000..67fd8fd29 --- /dev/null +++ b/test/sdk/core/agents/context/test_reducer_models.py @@ -0,0 +1,67 @@ +"""Tests for reducer result models.""" + +import pytest + +from nexent.core.agents.context.context_item import ( + AuthorityTier, + ContextItem, + ContextItemType, + RepresentationTier, +) +from nexent.core.agents.context.reducer_models import ReductionResult + + +class TestReductionResult: + """Tests for ReductionResult frozen dataclass.""" + + def test_reduction_result_is_frozen(self): + result = ReductionResult( + representation=RepresentationTier.FULL, + source_fingerprint="fp-1", + token_count=100, + generator="passthrough", + generator_version="0.1.0", + admissible=True, + loss_metadata={}, + content="test content", + ) + + with pytest.raises(AttributeError): + result.content = "modified" + + def test_reduction_result_creation(self): + result = ReductionResult( + representation=RepresentationTier.COMPRESSED, + source_fingerprint="fp-abc", + token_count=50, + generator="llm-summarizer", + generator_version="1.2.0", + admissible=True, + loss_metadata={"dropped_sections": 3, "compression_ratio": 0.6}, + content="summarized content", + ) + + assert result.representation == RepresentationTier.COMPRESSED + assert result.source_fingerprint == "fp-abc" + assert result.token_count == 50 + assert result.generator == "llm-summarizer" + assert result.generator_version == "1.2.0" + assert result.admissible is True + assert result.loss_metadata == {"dropped_sections": 3, "compression_ratio": 0.6} + assert result.content == "summarized content" + + def test_reduction_result_with_none_content(self): + result = ReductionResult( + representation=RepresentationTier.POINTER, + source_fingerprint="fp-none", + token_count=0, + generator="passthrough", + generator_version="0.1.0", + admissible=False, + loss_metadata={"reason": "empty_input"}, + content=None, + ) + + assert result.content is None + assert result.admissible is False + assert result.token_count == 0 diff --git a/test/sdk/core/agents/test_agent_context/loader.py b/test/sdk/core/agents/test_agent_context/loader.py index b6d4944f8..a45d1b073 100644 --- a/test/sdk/core/agents/test_agent_context/loader.py +++ b/test/sdk/core/agents/test_agent_context/loader.py @@ -285,6 +285,41 @@ class _SystemPromptComponent(_ContextComponent): if token_est_key not in sys.modules: sys.modules[token_est_key] = _build_token_estimation_stub() + # Stub for nexent.monitor (used by manager.py for OTel instrumentation) + if "nexent.monitor" not in sys.modules: + _monitor_stub = ModuleType("nexent.monitor") + + class _MockMonitoringManager: + def is_enabled(self): + return False + + def trace_operation(self, *args, **kwargs): + from contextlib import contextmanager + @contextmanager + def _noop(): + yield None + return _noop() + + def trace_llm_request(self, *args, **kwargs): + from contextlib import contextmanager + @contextmanager + def _noop(): + yield None + return _noop() + + def set_openinference_output(self, *args, **kwargs): + pass + + def add_span_event(self, *args, **kwargs): + pass + + _monitor_stub.get_monitoring_manager = lambda: _MockMonitoringManager() + _monitor_stub.OPENINFERENCE_SPAN_KIND_CHAIN = "CHAIN" + _monitor_stub.OPENINFERENCE_SPAN_KIND_LLM = "LLM" + _monitor_stub.OPENINFERENCE_INPUT_VALUE = "input.value" + _monitor_stub.OPENINFERENCE_OUTPUT_VALUE = "output.value" + sys.modules["nexent.monitor"] = _monitor_stub + _register_stub_packages() diff --git a/test/sdk/core/agents/test_context_item_pipeline.py b/test/sdk/core/agents/test_context_item_pipeline.py new file mode 100644 index 000000000..4e88bc958 --- /dev/null +++ b/test/sdk/core/agents/test_context_item_pipeline.py @@ -0,0 +1,555 @@ +""" +Comprehensive verification test for context module features. + +Tests: +- All 9 handler to_messages() implementations +- ContextProjector with all 7 component types +- Semantic equivalence between use_context_items=True/False +- HistoryProjector with both projection purposes +- Integration: Full agent run with all components + history projector +""" +import os +import sys +from pathlib import Path +from typing import List, Dict, Any, Optional +from threading import Event + +_project_root = Path(__file__).parent.parent.parent.parent.parent +sys.path.insert(0, str(_project_root / "sdk")) +sys.path.insert(0, str(_project_root / "backend")) + +from utils.monitoring import monitoring_manager + +from nexent.core.agents.context.handlers import register_all +from nexent.core.agents.context.item_handler_registry import ItemHandlerRegistry +from nexent.core.agents.context.context_item import ( + ContextItem, + ContextItemType, + AuthorityTier, + RepresentationTier, +) +from nexent.core.agents.context.projector import ContextProjector +from nexent.core.agents.context.history_projector import HistoryProjector +from nexent.core.agents.agent_model import ( + SystemPromptComponent, + ToolsComponent, + SkillsComponent, + MemoryComponent, + KnowledgeBaseComponent, + ManagedAgentsComponent, + ExternalAgentsComponent, +) +from nexent.core.agents.summary_config import ContextManagerConfig +from nexent.core.agents.agent_context import ContextManager + + +def test_pr0_handlers(): + """Handler: Test all 9 handler to_messages() implementations.""" + print("\n" + "=" * 70) + print("Handler: Testing all 9 handler to_messages() implementations") + print("=" * 70) + + register_all() + + test_cases = [ + # (item_type, content, expected_message_count, expected_roles) + ( + "SystemPromptHandler", + ContextItemType.SYSTEM_PROMPT, + "You are a helpful assistant", + 1, + ["user"], + ), + ( + "ToolHandler", + ContextItemType.TOOL, + {"name": "search", "description": "Search the web"}, + 1, + ["user"], + ), + ( + "SkillHandler", + ContextItemType.SKILL, + {"name": "coding", "description": "Write code"}, + 1, + ["user"], + ), + ( + "MemoryHandler", + ContextItemType.MEMORY, + {"content": "User prefers Python", "memory_type": "user"}, + 1, + ["user"], + ), + ( + "KnowledgeBaseHandler", + ContextItemType.KNOWLEDGE_BASE, + "Retrieved knowledge about AI", + 1, + ["user"], + ), + ( + "ManagedAgentHandler", + ContextItemType.MANAGED_AGENT, + {"name": "researcher", "description": "Research agent"}, + 1, + ["user"], + ), + ( + "ExternalAgentHandler", + ContextItemType.EXTERNAL_AGENT, + {"name": "external_api", "description": "External API agent"}, + 1, + ["user"], + ), + ( + "HistoryTurnHandler (both)", + ContextItemType.HISTORY_TURN, + {"user_query": "What is AI?", "assistant_response": "AI is artificial intelligence"}, + 2, + ["user", "assistant"], + ), + ( + "HistoryTurnHandler (user only)", + ContextItemType.HISTORY_TURN, + {"user_query": "What is AI?"}, + 1, + ["user"], + ), + ( + "HistoryTurnHandler (empty)", + ContextItemType.HISTORY_TURN, + {}, + 0, + [], + ), + ( + "ToolCallResultHandler", + ContextItemType.TOOL_CALL_RESULT, + {"tool_call": "search('AI')", "execution_result": "Found 10 results"}, + 1, + ["user"], + ), + ] + + passed = 0 + failed = 0 + + for name, item_type, content, expected_count, expected_roles in test_cases: + item = ContextItem( + item_id=f"test:{name}", + item_type=item_type, + content=content, + ) + + handler = ItemHandlerRegistry.get(item_type) + messages = handler.to_messages(item) + + if len(messages) != expected_count: + print(f" ❌ {name}: expected {expected_count} messages, got {len(messages)}") + failed += 1 + continue + + actual_roles = [msg["role"] for msg in messages] + if actual_roles != expected_roles: + print(f" ❌ {name}: expected roles {expected_roles}, got {actual_roles}") + failed += 1 + continue + + print(f" ✅ {name}: {len(messages)} message(s), roles={actual_roles}") + passed += 1 + + print(f"\nHandlers: {passed} passed, {failed} failed") + return failed == 0 + + +def test_pr0_projector(): + """Projector: Test ContextProjector with all 7 component types.""" + print("\n" + "=" * 70) + print("Projector: Testing ContextProjector with all 7 component types") + print("=" * 70) + + register_all() + projector = ContextProjector() + + components = [ + SystemPromptComponent(content="You are helpful"), + ToolsComponent( + tools=[ + {"name": "search", "description": "Search"}, + {"name": "calculate", "description": "Calculate"}, + ], + formatted_description="Available tools: search, calculate", + ), + SkillsComponent( + skills=[{"name": "coding", "description": "Write code"}], + formatted_description="Skills: coding", + ), + MemoryComponent( + memories=[ + {"content": "User likes Python", "memory_type": "user"}, + {"content": "User is a developer", "memory_type": "user"}, + ], + formatted_content="User preferences: Python, developer", + ), + KnowledgeBaseComponent( + summary="AI knowledge base summary", + kb_ids=["kb_1", "kb_2"], + ), + ManagedAgentsComponent( + agents=[{"name": "researcher", "description": "Research agent"}], + formatted_description="Managed agents: researcher", + ), + ExternalAgentsComponent( + agents=[{"name": "external_api", "description": "External API"}], + formatted_description="External agents: external_api", + ), + ] + + items = projector.project(components) + + expected_types = { + ContextItemType.SYSTEM_PROMPT: 1, + ContextItemType.TOOL: 2, + ContextItemType.SKILL: 1, + ContextItemType.MEMORY: 2, + ContextItemType.KNOWLEDGE_BASE: 1, + ContextItemType.MANAGED_AGENT: 1, + ContextItemType.EXTERNAL_AGENT: 1, + } + + actual_counts = {} + for item in items: + actual_counts[item.item_type] = actual_counts.get(item.item_type, 0) + 1 + + print(f" Total items projected: {len(items)}") + + passed = 0 + failed = 0 + + for item_type, expected_count in expected_types.items(): + actual_count = actual_counts.get(item_type, 0) + if actual_count == expected_count: + print(f" ✅ {item_type.value}: {actual_count} item(s)") + passed += 1 + else: + print(f" ❌ {item_type.value}: expected {expected_count}, got {actual_count}") + failed += 1 + + has_source_ref = all("_source_component" in item.metadata for item in items) + if has_source_ref: + print(f" ✅ All items have _source_component back-reference") + passed += 1 + else: + print(f" ❌ Some items missing _source_component back-reference") + failed += 1 + + print(f"\nProjector: {passed} passed, {failed} failed") + return failed == 0 + + +def test_pr1_semantic_equivalence(): + """Equivalence: Test semantic equivalence between use_context_items=True/False.""" + print("\n" + "=" * 70) + print("Equivalence: Testing semantic equivalence (use_context_items=True vs False)") + print("=" * 70) + + register_all() + + from smolagents.memory import AgentMemory + + components = [ + SystemPromptComponent(content="You are a helpful assistant"), + ToolsComponent( + tools=[{"name": "search", "description": "Search"}], + formatted_description="Tools: search", + ), + MemoryComponent( + memories=[{"content": "User fact", "memory_type": "user"}], + formatted_content="Memory: User fact", + ), + ] + + config_false = ContextManagerConfig( + enabled=True, + use_context_items=False, + token_threshold=100000, + strategy="full", + ) + + config_true = ContextManagerConfig( + enabled=True, + use_context_items=True, + token_threshold=100000, + strategy="full", + ) + + cm_false = ContextManager(config=config_false) + cm_true = ContextManager(config=config_true) + + from smolagents.memory import SystemPromptStep + memory_false = AgentMemory(system_prompt=SystemPromptStep(system_prompt="")) + memory_true = AgentMemory(system_prompt=SystemPromptStep(system_prompt="")) + + run_ctx_false = cm_false.prepare_run_context( + memory=memory_false, + fallback_system_prompt="", + components=components, + ) + + run_ctx_true = cm_true.prepare_run_context( + memory=memory_true, + fallback_system_prompt="", + components=components, + ) + + msgs_false = list(run_ctx_false.component_messages) + msgs_true = list(run_ctx_true.component_messages) + + print(f" OLD path (use_context_items=False): {len(msgs_false)} messages") + print(f" NEW path (use_context_items=True): {len(msgs_true)} messages") + + if len(msgs_false) != len(msgs_true): + print(f" ❌ Message count mismatch") + return False + + for i, (msg_f, msg_t) in enumerate(zip(msgs_false, msgs_true)): + if msg_f["role"] != msg_t["role"]: + print(f" ❌ Message {i} role mismatch: {msg_f['role']} vs {msg_t['role']}") + return False + + content_f = msg_f["content"][0]["text"] if msg_f["content"] else "" + content_t = msg_t["content"][0]["text"] if msg_t["content"] else "" + + if content_f != content_t: + print(f" ❌ Message {i} content mismatch") + print(f" OLD: {content_f[:100]}...") + print(f" NEW: {content_t[:100]}...") + return False + + print(f" ✅ All {len(msgs_false)} messages are semantically equivalent") + print(f"\nSemantic Equivalence: PASSED") + return True + + +def test_pr2_history_projector(): + """HistoryProjector: Test HistoryProjector with both projection purposes.""" + print("\n" + "=" * 70) + print("HistoryProjector: Testing HistoryProjector with both projection purposes") + print("=" * 70) + + register_all() + + def mock_query_units(conversation_id: int, message_id: Optional[int] = None) -> List[Dict[str, Any]]: + units = [ + {"unit_id": 1, "unit_type": "user_input", "unit_content": "What is AI?", "message_id": 1, "step_index": 1}, + {"unit_id": 2, "unit_type": "final_answer", "unit_content": "AI is artificial intelligence", "message_id": 1, "step_index": 1}, + {"unit_id": 3, "unit_type": "tool_call", "unit_content": '{"tool_call": "search(\'machine learning\')", "execution_result": "Found 5 results"}', "message_id": 1, "step_index": 2}, + {"unit_id": 4, "unit_type": "model_output_thinking", "unit_content": "Let me think about this...", "message_id": 1, "step_index": 2}, + {"unit_id": 5, "unit_type": "user_input", "unit_content": "Tell me more", "message_id": 2, "step_index": 1}, + ] + if message_id is not None: + return [u for u in units if u.get("message_id") == message_id] + return units + + projector = HistoryProjector(query_units_fn=mock_query_units) + + passed = 0 + failed = 0 + + print("\n Testing purpose='model_context':") + items_mc = projector.project(conversation_id=123, purpose="model_context") + + mc_types = {} + for item in items_mc: + mc_types[item.item_type] = mc_types.get(item.item_type, 0) + 1 + + if ContextItemType.HISTORY_TURN in mc_types: + print(f" ✅ HISTORY_TURN: {mc_types[ContextItemType.HISTORY_TURN]} item(s)") + passed += 1 + else: + print(f" ❌ Missing HISTORY_TURN items") + failed += 1 + + if ContextItemType.TOOL_CALL_RESULT in mc_types: + print(f" ✅ TOOL_CALL_RESULT: {mc_types[ContextItemType.TOOL_CALL_RESULT]} item(s)") + passed += 1 + else: + print(f" ❌ Missing TOOL_CALL_RESULT items") + failed += 1 + + print("\n Testing purpose='chat':") + items_chat = projector.project(conversation_id=123, purpose="chat") + + chat_types = {} + for item in items_chat: + chat_types[item.item_type] = chat_types.get(item.item_type, 0) + 1 + + if ContextItemType.HISTORY_TURN in chat_types: + print(f" ✅ HISTORY_TURN: {chat_types[ContextItemType.HISTORY_TURN]} item(s)") + passed += 1 + else: + print(f" ❌ Missing HISTORY_TURN items") + failed += 1 + + has_thinking = any( + "thinking" in str(item.content).lower() + for item in items_chat + if item.item_type == ContextItemType.HISTORY_TURN + ) + if has_thinking: + print(f" ✅ Chat includes thinking content") + passed += 1 + else: + print(f" ⚠️ Chat may not include thinking (check implementation)") + + print(f"\nHistoryProjector: {passed} passed, {failed} failed") + return failed == 0 + + +def test_integration_full_agent(): + """Integration: Full agent run with all components + history projector.""" + print("\n" + "=" * 70) + print("Integration: Full agent run with all components + history projector") + print("=" * 70) + + try: + import asyncio + from dotenv import load_dotenv + load_dotenv(override=True) + + from nexent.core.utils.observer import MessageObserver + from nexent.core.agents.agent_model import ModelConfig, AgentConfig, AgentRunInfo + from nexent.core.agents.run_agent import agent_run + from nexent.monitor import agent_monitoring_context, AgentRunMetadata + + def mock_query_units(conversation_id: int, message_id: Optional[int] = None) -> List[Dict[str, Any]]: + return [ + {"unit_id": 1, "unit_type": "user_input", "unit_content": "Previous question", "message_id": 1, "step_index": 1}, + {"unit_id": 2, "unit_type": "final_answer", "unit_content": "Previous answer", "message_id": 1, "step_index": 1}, + ] + + history_projector = HistoryProjector(query_units_fn=mock_query_units) + + cm_config = ContextManagerConfig( + enabled=True, + use_context_items=True, + token_threshold=100000, + strategy="full", + history_projector=history_projector, + ) + + components = [ + SystemPromptComponent(content="You are a helpful assistant. Answer concisely."), + ToolsComponent( + tools=[{"name": "search", "description": "Search the web"}], + formatted_description="Available tools: search", + ), + MemoryComponent( + memories=[{"content": "User prefers English", "memory_type": "user"}], + formatted_content="User preferences: English", + ), + ] + + model_cfg = ModelConfig( + cite_name="main_model", + api_key=os.environ.get("DEFAULT_MODEL_API_KEY", ""), + model_name=os.environ.get("DEFAULT_MODEL", "qwen3.6-plus"), + url=os.environ.get("DEFAULT_MODEL_ENDPOINT", ""), + temperature=0.1, + extra_body={"chat_template_kwargs": {"enable_thinking": False}}, + ) + + agent_cfg = AgentConfig( + name="integration_test_agent", + description="Full integration test", + model_name="main_model", + tools=[], + max_steps=3, + context_manager_config=cm_config, + context_components=components, + conversation_id=12345, + ) + + observer = MessageObserver(lang="en") + stop_event = Event() + + run_info = AgentRunInfo( + query="What is 2+2?", + model_config_list=[model_cfg], + observer=observer, + agent_config=agent_cfg, + stop_event=stop_event, + ) + + metadata = AgentRunMetadata( + agent_name="integration_test_agent", + query="What is 2+2?", + tenant_id="test", + user_id="integration_test", + conversation_id=12345, + model_name=model_cfg.model_name, + ) + + async def run(): + final_answer = "" + with agent_monitoring_context(metadata): + async for message in agent_run(run_info): + msg = message if isinstance(message, dict) else eval(message) + if msg.get("type") == "final_answer": + final_answer = msg.get("content", "") + return final_answer + + answer = asyncio.run(run()) + + if answer: + print(f" ✅ Agent completed with answer: {answer[:100]}") + print(f"\nIntegration Test: PASSED") + return True + else: + print(f" ❌ Agent did not produce final answer") + print(f"\nIntegration Test: FAILED") + return False + + except Exception as e: + print(f" ❌ Integration test failed with error: {e}") + import traceback + traceback.print_exc() + print(f"\nIntegration Test: FAILED") + return False + + +def main(): + print("=" * 70) + print("Comprehensive Context Module Verification") + print("=" * 70) + + results = {} + + results["Handlers"] = test_pr0_handlers() + results["Projector"] = test_pr0_projector() + results["Equivalence"] = test_pr1_semantic_equivalence() + results["HistoryProjector"] = test_pr2_history_projector() + results["Integration"] = test_integration_full_agent() + + print("\n" + "=" * 70) + print("SUMMARY") + print("=" * 70) + + for test_name, passed in results.items(): + status = "✅ PASSED" if passed else "❌ FAILED" + print(f" {status}: {test_name}") + + all_passed = all(results.values()) + print("\n" + "=" * 70) + if all_passed: + print("ALL TESTS PASSED ✅") + else: + print("SOME TESTS FAILED ❌") + print("=" * 70) + + return 0 if all_passed else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/test/sdk/core/agents/test_context_item_projection.py b/test/sdk/core/agents/test_context_item_projection.py new file mode 100644 index 000000000..68b6d3234 --- /dev/null +++ b/test/sdk/core/agents/test_context_item_projection.py @@ -0,0 +1,476 @@ +"""Tests for ContextItem projection integration in ContextManager. + +Verifies that when use_context_items=True, ContextManager correctly: +1. Projects ContextComponent instances into ContextItem candidates +2. Converts ContextItems back to messages for the assembly pipeline +3. Includes projected items in ContextEvidence +4. Maintains backward compatibility when use_context_items=False +""" + +import pytest +from unittest.mock import MagicMock + +from nexent.core.agents.agent_context import ContextManager +from nexent.core.agents.agent_model import ( + ExternalAgentsComponent, + KnowledgeBaseComponent, + ManagedAgentsComponent, + MemoryComponent, + SkillsComponent, + SystemPromptComponent, + ToolsComponent, +) +from nexent.core.agents.summary_config import ContextManagerConfig +from nexent.core.agents.context.context_item import AuthorityTier, ContextItemType +from nexent.core.agents.context.handlers import register_all + + +@pytest.fixture(autouse=True) +def ensure_handlers_registered(): + """Ensure all handlers are registered before each test.""" + register_all() + + +class TestContextItemProjectionDisabled: + """Verify backward compatibility when use_context_items=False (default).""" + + def test_default_config_has_use_context_items_false(self): + """ContextManagerConfig defaults to use_context_items=False.""" + config = ContextManagerConfig(enabled=True, token_threshold=10000) + assert config.use_context_items is False + + def test_assemble_final_context_without_projection(self): + """When use_context_items=False, context_items in evidence is empty.""" + config = ContextManagerConfig( + enabled=True, + token_threshold=10000, + use_context_items=False, + ) + manager = ContextManager(config=config) + + manager.register_component(SystemPromptComponent(content="You are helpful")) + manager.register_component( + ToolsComponent( + tools=[{"name": "search", "description": "Search the web"}], + formatted_description="Available tools: search", + ) + ) + + memory = MagicMock() + memory.system_prompt = None + memory.steps = [] + + final = manager.assemble_final_context( + model=None, + memory=memory, + current_run_start_idx=0, + tools=[], + ) + + assert final.evidence.context_items == () + assert len(final.messages) > 0 + + +class TestContextItemProjectionEnabled: + """Verify ContextItem projection when use_context_items=True.""" + + def test_config_accepts_use_context_items_true(self): + """ContextManagerConfig can be initialized with use_context_items=True.""" + config = ContextManagerConfig( + enabled=True, + token_threshold=10000, + use_context_items=True, + ) + assert config.use_context_items is True + + def test_project_context_items_returns_list(self): + """project_context_items() returns a list of ContextItem objects.""" + config = ContextManagerConfig(enabled=True, token_threshold=10000) + manager = ContextManager(config=config) + + manager.register_component(SystemPromptComponent(content="System prompt")) + manager.register_component( + ToolsComponent( + tools=[ + {"name": "tool1", "description": "First tool"}, + {"name": "tool2", "description": "Second tool"}, + ], + formatted_description="Tools: tool1, tool2", + ) + ) + + items = manager.project_context_items() + + assert isinstance(items, list) + assert len(items) == 3 + assert items[0].item_type == ContextItemType.SYSTEM_PROMPT + assert items[1].item_type == ContextItemType.TOOL + assert items[2].item_type == ContextItemType.TOOL + + def test_assemble_final_context_with_projection_includes_context_items(self): + """When use_context_items=True, evidence.context_items contains projected items.""" + config = ContextManagerConfig( + enabled=True, + token_threshold=10000, + use_context_items=True, + ) + manager = ContextManager(config=config) + + manager.register_component(SystemPromptComponent(content="Be helpful")) + manager.register_component( + MemoryComponent( + memories=[{"content": "User prefers Python", "memory_type": "user"}], + formatted_content="User preferences: Python", + ) + ) + + memory = MagicMock() + memory.system_prompt = None + memory.steps = [] + + final = manager.assemble_final_context( + model=None, + memory=memory, + current_run_start_idx=0, + tools=[], + ) + + assert len(final.evidence.context_items) > 0 + item_types = {item.item_type for item in final.evidence.context_items} + assert ContextItemType.SYSTEM_PROMPT in item_types + assert ContextItemType.MEMORY in item_types + + def test_projection_preserves_message_roles(self): + """Projected items are converted to messages with correct roles.""" + config = ContextManagerConfig( + enabled=True, + token_threshold=10000, + use_context_items=True, + ) + manager = ContextManager(config=config) + + manager.register_component(SystemPromptComponent(content="System instruction")) + manager.register_component( + MemoryComponent( + memories=[{"content": "Retrieved fact", "memory_type": "user"}], + formatted_content="Memory: Retrieved fact", + ) + ) + manager.register_component( + KnowledgeBaseComponent( + summary="KB summary", + kb_ids=["kb1"], + ) + ) + + memory = MagicMock() + memory.system_prompt = None + memory.steps = [] + + final = manager.assemble_final_context( + model=None, + memory=memory, + current_run_start_idx=0, + tools=[], + ) + + roles = [msg["role"] for msg in final.messages] + + assert "system" in roles + assert "user" in roles + + def test_projection_uses_formatted_description(self): + config = ContextManagerConfig( + enabled=True, + token_threshold=10000, + use_context_items=True, + ) + manager = ContextManager(config=config) + + tool_dict = {"name": "calculator", "description": "Performs math"} + manager.register_component( + ToolsComponent( + tools=[tool_dict], + formatted_description="Calculator tool", + ) + ) + + memory = MagicMock() + memory.system_prompt = None + memory.steps = [] + + final = manager.assemble_final_context( + model=None, + memory=memory, + current_run_start_idx=0, + tools=[], + ) + + assert len(final.messages) > 0 + tool_msg = None + for msg in final.messages: + if msg["role"] == "system": + content = msg["content"] + if isinstance(content, list) and len(content) > 0: + text = content[0].get("text", "") + if "Calculator tool" in text: + tool_msg = msg + break + + assert tool_msg is not None + assert "Calculator tool" in tool_msg["content"][0]["text"] + + def test_projection_maintains_stable_dynamic_partition(self): + """Projected items are correctly partitioned into stable/dynamic.""" + config = ContextManagerConfig( + enabled=True, + token_threshold=10000, + use_context_items=True, + ) + manager = ContextManager(config=config) + + manager.register_component(SystemPromptComponent(content="Stable system prompt")) + manager.register_component( + MemoryComponent( + memories=[{"content": "Dynamic memory", "memory_type": "user"}], + formatted_content="Dynamic memory content", + ) + ) + + memory = MagicMock() + memory.system_prompt = None + memory.steps = [] + + final = manager.assemble_final_context( + model=None, + memory=memory, + current_run_start_idx=0, + tools=[], + ) + + assert final.evidence.stable_message_count > 0 + assert final.evidence.dynamic_message_count > 0 + + def test_empty_components_returns_empty_projection(self): + """When no components are registered, projection returns empty list.""" + config = ContextManagerConfig(enabled=True, token_threshold=10000) + manager = ContextManager(config=config) + + items = manager.project_context_items() + assert items == [] + + +class TestContextItemProjectionEdgeCases: + """Edge cases and error handling for ContextItem projection.""" + + def test_projection_with_none_content_items(self): + """Items with None content are skipped during message conversion.""" + config = ContextManagerConfig( + enabled=True, + token_threshold=10000, + use_context_items=True, + ) + manager = ContextManager(config=config) + + manager.register_component(SystemPromptComponent(content="")) + + memory = MagicMock() + memory.system_prompt = None + memory.steps = [] + + final = manager.assemble_final_context( + model=None, + memory=memory, + current_run_start_idx=0, + tools=[], + ) + + assert final is not None + assert final.evidence.context_items is not None + + def test_projection_with_multiple_tools(self): + """Multiple tools are projected as separate ContextItems.""" + config = ContextManagerConfig( + enabled=True, + token_threshold=10000, + use_context_items=True, + ) + manager = ContextManager(config=config) + + tools = [ + {"name": "tool1", "description": "First"}, + {"name": "tool2", "description": "Second"}, + {"name": "tool3", "description": "Third"}, + ] + manager.register_component( + ToolsComponent(tools=tools, formatted_description="Three tools") + ) + + items = manager.project_context_items() + + tool_items = [item for item in items if item.item_type == ContextItemType.TOOL] + assert len(tool_items) == 3 + assert tool_items[0].item_id != tool_items[1].item_id + assert tool_items[1].item_id != tool_items[2].item_id + + def test_projection_preserves_authority_tiers(self): + """Different component types have correct authority tiers.""" + config = ContextManagerConfig(enabled=True, token_threshold=10000) + manager = ContextManager(config=config) + + manager.register_component(SystemPromptComponent(content="System")) + manager.register_component( + MemoryComponent( + memories=[{"content": "Memory", "memory_type": "user"}], + formatted_content="Memory", + ) + ) + manager.register_component( + KnowledgeBaseComponent(summary="KB", kb_ids=["kb1"]) + ) + + items = manager.project_context_items() + + system_item = next( + item for item in items if item.item_type == ContextItemType.SYSTEM_PROMPT + ) + assert system_item.authority_tier == AuthorityTier.PLATFORM + + memory_item = next( + item for item in items if item.item_type == ContextItemType.MEMORY + ) + assert memory_item.authority_tier == AuthorityTier.RETRIEVED_MEMORY + + kb_item = next( + item for item in items if item.item_type == ContextItemType.KNOWLEDGE_BASE + ) + assert kb_item.authority_tier == AuthorityTier.RETRIEVED_MEMORY + + def test_projection_handles_skills_component(self): + config = ContextManagerConfig( + enabled=True, + token_threshold=10000, + use_context_items=True, + ) + manager = ContextManager(config=config) + + manager.register_component( + SkillsComponent( + skills=[{"name": "skill1", "description": "First skill"}], + formatted_description="Available skills: skill1", + ) + ) + + items = manager.project_context_items() + assert len(items) == 1 + assert items[0].item_type == ContextItemType.SKILL + assert items[0].authority_tier == AuthorityTier.PLATFORM + + memory = MagicMock() + memory.system_prompt = None + memory.steps = [] + + final = manager.assemble_final_context( + model=None, + memory=memory, + current_run_start_idx=0, + tools=[], + ) + + skill_msg = None + for msg in final.messages: + if msg["role"] == "system": + content = msg["content"] + if isinstance(content, list) and len(content) > 0: + text = content[0].get("text", "") + if "Available skills" in text: + skill_msg = msg + break + + assert skill_msg is not None + + def test_projection_handles_managed_agents_component(self): + config = ContextManagerConfig( + enabled=True, + token_threshold=10000, + use_context_items=True, + ) + manager = ContextManager(config=config) + + manager.register_component( + ManagedAgentsComponent( + agents=[{"name": "agent1", "description": "First agent"}], + formatted_description="Managed agents: agent1", + ) + ) + + items = manager.project_context_items() + assert len(items) == 1 + assert items[0].item_type == ContextItemType.MANAGED_AGENT + assert items[0].authority_tier == AuthorityTier.PLATFORM + + memory = MagicMock() + memory.system_prompt = None + memory.steps = [] + + final = manager.assemble_final_context( + model=None, + memory=memory, + current_run_start_idx=0, + tools=[], + ) + + agent_msg = None + for msg in final.messages: + if msg["role"] == "system": + content = msg["content"] + if isinstance(content, list) and len(content) > 0: + text = content[0].get("text", "") + if "Managed agents" in text: + agent_msg = msg + break + + assert agent_msg is not None + + def test_projection_handles_external_agents_component(self): + config = ContextManagerConfig( + enabled=True, + token_threshold=10000, + use_context_items=True, + ) + manager = ContextManager(config=config) + + manager.register_component( + ExternalAgentsComponent( + agents=[{"agent_id": "ext1", "name": "external1", "description": "External agent"}], + formatted_description="External agents: external1", + ) + ) + + items = manager.project_context_items() + assert len(items) == 1 + assert items[0].item_type == ContextItemType.EXTERNAL_AGENT + assert items[0].authority_tier == AuthorityTier.PLATFORM + + memory = MagicMock() + memory.system_prompt = None + memory.steps = [] + + final = manager.assemble_final_context( + model=None, + memory=memory, + current_run_start_idx=0, + tools=[], + ) + + ext_msg = None + for msg in final.messages: + if msg["role"] == "system": + content = msg["content"] + if isinstance(content, list) and len(content) > 0: + text = content[0].get("text", "") + if "External agents" in text: + ext_msg = msg + break + + assert ext_msg is not None diff --git a/test/sdk/core/agents/test_context_item_runtime_integration.py b/test/sdk/core/agents/test_context_item_runtime_integration.py new file mode 100644 index 000000000..97430e2a6 --- /dev/null +++ b/test/sdk/core/agents/test_context_item_runtime_integration.py @@ -0,0 +1,143 @@ +"""Integration test for ManagedContextRuntime with use_context_items=True.""" + +import pytest +from unittest.mock import MagicMock + +from nexent.core.agents.agent_context import ContextManager +from nexent.core.agents.agent_model import ( + MemoryComponent, + SystemPromptComponent, + ToolsComponent, +) +from nexent.core.agents.summary_config import ContextManagerConfig +from nexent.core.agents.context.handlers import register_all +from nexent.core.context_runtime.managed.runtime import ManagedContextRuntime + + +@pytest.fixture(autouse=True) +def ensure_handlers_registered(): + register_all() + + +class TestManagedContextRuntimeWithContextItems: + + def test_prepare_step_with_context_items_enabled(self): + config = ContextManagerConfig( + enabled=True, + token_threshold=10000, + use_context_items=True, + ) + manager = ContextManager(config=config) + + system_prompt = SystemPromptComponent(content="You are helpful") + tools = ToolsComponent( + tools=[{"name": "search", "description": "Search the web"}], + formatted_description="Available tools: search", + ) + memory_comp = MemoryComponent( + memories=[{"content": "User prefers Python", "memory_type": "user"}], + formatted_content="User preferences: Python", + ) + + manager.register_component(system_prompt) + manager.register_component(tools) + manager.register_component(memory_comp) + + runtime = ManagedContextRuntime(manager, components=[system_prompt, tools, memory_comp]) + + memory = MagicMock() + memory.system_prompt = None + memory.steps = [] + + runtime.prepare_run(memory=memory, fallback_system_prompt="Fallback") + + final = runtime.prepare_step( + model=None, + memory=memory, + current_run_start_idx=0, + tools=[], + ) + + assert final is not None + assert len(final.messages) > 0 + assert final.evidence.context_items is not None + assert len(final.evidence.context_items) > 0 + + roles = [msg["role"] for msg in final.messages] + assert "system" in roles + assert "user" in roles + + def test_prepare_final_answer_with_context_items_enabled(self): + config = ContextManagerConfig( + enabled=True, + token_threshold=10000, + use_context_items=True, + ) + manager = ContextManager(config=config) + + system_prompt = SystemPromptComponent(content="You are helpful") + manager.register_component(system_prompt) + + runtime = ManagedContextRuntime(manager, components=[system_prompt]) + + memory = MagicMock() + memory.system_prompt = None + memory.steps = [] + + runtime.prepare_run(memory=memory, fallback_system_prompt="Fallback") + + final_answer_templates = { + "final_answer": { + "pre_messages": "Generate final answer", + "post_messages": "Task: {{ task }}", + } + } + + final = runtime.prepare_final_answer( + model=None, + memory=memory, + current_run_start_idx=0, + task="Answer the question", + final_answer_templates=final_answer_templates, + tools=[], + ) + + assert final is not None + assert len(final.messages) > 0 + assert final.evidence.context_items is not None + + def test_backward_compatibility_with_context_items_disabled(self): + config = ContextManagerConfig( + enabled=True, + token_threshold=10000, + use_context_items=False, + ) + manager = ContextManager(config=config) + + system_prompt = SystemPromptComponent(content="You are helpful") + memory_comp = MemoryComponent( + memories=[{"content": "User prefers Python", "memory_type": "user"}], + formatted_content="User preferences: Python", + ) + + manager.register_component(system_prompt) + manager.register_component(memory_comp) + + runtime = ManagedContextRuntime(manager, components=[system_prompt, memory_comp]) + + memory = MagicMock() + memory.system_prompt = None + memory.steps = [] + + runtime.prepare_run(memory=memory, fallback_system_prompt="Fallback") + + final = runtime.prepare_step( + model=None, + memory=memory, + current_run_start_idx=0, + tools=[], + ) + + assert final is not None + assert len(final.messages) > 0 + assert final.evidence.context_items == () diff --git a/test/sdk/core/agents/test_context_strategy_selection.py b/test/sdk/core/agents/test_context_strategy_selection.py new file mode 100644 index 000000000..39dcb8967 --- /dev/null +++ b/test/sdk/core/agents/test_context_strategy_selection.py @@ -0,0 +1,535 @@ +""" +Strategy filtering fix verification tests. + +Verifies that use_context_items=True respects strategy-based component selection +instead of bypassing it. The fix adds selected_components to ManagedRunContext +and ensures assemble_final_context() uses strategy-filtered components. + +Tests: + 1. Budget-pressure equivalence between use_context_items=True/False + 2. Strategy drop propagation when use_context_items=True + 3. selected_component_types accuracy after strategy filtering + 4. Fingerprint stability across use_context_items toggle + 5. History items unaffected by strategy filtering +""" +import sys +from pathlib import Path + +_project_root = Path(__file__).parent.parent.parent.parent.parent +sys.path.insert(0, str(_project_root / "sdk")) +sys.path.insert(0, str(_project_root / "backend")) + +from typing import Any, Dict, List, Optional +from unittest.mock import MagicMock + +from smolagents.memory import AgentMemory, SystemPromptStep + +from nexent.core.agents.context.handlers import register_all +from nexent.core.agents.summary_config import ContextManagerConfig +from nexent.core.agents.agent_context import ContextManager, ManagedRunContext +from nexent.core.agents.agent_model import ( + SystemPromptComponent, + ToolsComponent, + MemoryComponent, + KnowledgeBaseComponent, +) +from nexent.core.agents.context.history_projector import HistoryProjector +from nexent.core.agents.context.context_item import ContextItemType + +register_all() + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def create_mock_model() -> MagicMock: + """Create a mock model that satisfies compress_if_needed interface.""" + model = MagicMock() + model.model_id = "test-model" + return model + + +def create_components_with_sizes( + system_prompt_chars: int = 150, + tools_chars: int = 300, + memory_chars: int = 200, + kb_chars: int = 400, +) -> List[Any]: + """Create four components with controllable text sizes. + + Token estimates use chars_per_token=1.5 (default), so: + 150 chars ~ 100 tokens, 300 chars ~ 200 tokens, etc. + """ + return [ + SystemPromptComponent( + content="S" * system_prompt_chars, + priority=100, + ), + ToolsComponent( + tools=[{"name": "search", "description": "Search the web"}], + formatted_description="T" * tools_chars, + priority=80, + ), + MemoryComponent( + memories=[{"content": "User fact", "memory_type": "user"}], + formatted_content="M" * memory_chars, + priority=50, + ), + KnowledgeBaseComponent( + summary="K" * kb_chars, + kb_ids=["kb_1"], + priority=30, + ), + ] + + +def build_memory() -> AgentMemory: + """Build a minimal AgentMemory for prepare_run_context.""" + return AgentMemory(system_prompt=SystemPromptStep(system_prompt="")) + + +# --------------------------------------------------------------------------- +# Test 1: Budget-pressure equivalence +# --------------------------------------------------------------------------- + +def test_budget_pressure_equivalence() -> bool: + """Verify use_context_items=True and False produce same component set under budget pressure.""" + print("\n" + "=" * 70) + print("Test 1: Budget-pressure equivalence") + print("=" * 70) + + components = create_components_with_sizes( + system_prompt_chars=150, # ~100 tokens + tools_chars=300, # ~200 tokens + memory_chars=200, # ~133 tokens + kb_chars=400, # ~267 tokens + ) + # Total ~700 tokens. Set budget so only ~3 components fit. + # component_budgets total (excluding conversation_history) controls selection. + config_false = ContextManagerConfig( + enabled=True, + use_context_items=False, + token_threshold=500, + strategy="token_budget", + component_budgets={ + "system_prompt": 200, + "tools": 300, + "memory": 200, + "knowledge_base": 200, + "conversation_history": 0, + }, + ) + config_true = ContextManagerConfig( + enabled=True, + use_context_items=True, + token_threshold=500, + strategy="token_budget", + component_budgets={ + "system_prompt": 200, + "tools": 300, + "memory": 200, + "knowledge_base": 200, + "conversation_history": 0, + }, + ) + + cm_false = ContextManager(config=config_false) + cm_true = ContextManager(config=config_true) + + memory_false = build_memory() + memory_true = build_memory() + + run_ctx_false = cm_false.prepare_run_context( + memory=memory_false, + fallback_system_prompt="", + components=components, + ) + run_ctx_true = cm_true.prepare_run_context( + memory=memory_true, + fallback_system_prompt="", + components=components, + ) + + types_false = set(run_ctx_false.selected_component_types) + types_true = set(run_ctx_true.selected_component_types) + + print(f" use_context_items=False selected: {sorted(types_false)}") + print(f" use_context_items=True selected: {sorted(types_true)}") + + if types_false == types_true: + print(" \u2705 PASSED: Both paths selected same component types") + return True + else: + print(f" \u274c FAILED: Mismatch - False={sorted(types_false)}, True={sorted(types_true)}") + return False + + +# --------------------------------------------------------------------------- +# Test 2: Strategy drop propagation +# --------------------------------------------------------------------------- + +def test_strategy_drop_propagation() -> bool: + """Verify low-priority components are dropped when use_context_items=True.""" + print("\n" + "=" * 70) + print("Test 2: Strategy drop propagation") + print("=" * 70) + + high_priority = SystemPromptComponent( + content="H" * 150, + priority=100, + ) + medium_priority = ToolsComponent( + tools=[{"name": "calc", "description": "Calculate"}], + formatted_description="M" * 200, + priority=50, + ) + low_priority = KnowledgeBaseComponent( + summary="L" * 300, + kb_ids=["kb_low"], + priority=10, + ) + + components = [high_priority, medium_priority, low_priority] + + # _calculate_component_budget sums component_budgets (excluding conversation_history). + # high=100 tokens, medium=133 tokens, low=200 tokens. + # Total budget = 120+150+50 = 320. After high+medium (233), low (200) won't fit (433>320). + # Per-type: knowledge_base=50 also blocks low (200>50). + config = ContextManagerConfig( + enabled=True, + use_context_items=True, + token_threshold=300, + strategy="token_budget", + component_budgets={ + "system_prompt": 120, + "tools": 150, + "knowledge_base": 50, + "conversation_history": 0, + }, + ) + + cm = ContextManager(config=config) + memory = build_memory() + + run_ctx = cm.prepare_run_context( + memory=memory, + fallback_system_prompt="", + components=components, + ) + + selected_types = set(run_ctx.selected_component_types) + print(f" Selected types: {sorted(selected_types)}") + + # Verify low-priority KB was dropped + has_high = "system_prompt" in selected_types + has_medium = "tools" in selected_types + has_low = "knowledge_base" in selected_types + + if has_high and has_medium and not has_low: + print(" \u2705 PASSED: Low-priority component correctly dropped") + return True + else: + print(f" \u274c FAILED: high={has_high}, medium={has_medium}, low={has_low}") + return False + + +# --------------------------------------------------------------------------- +# Test 3: selected_component_types accuracy +# --------------------------------------------------------------------------- + +def test_selected_component_types_accuracy() -> bool: + """Verify selected_component_types reflects strategy-filtered components, not all.""" + print("\n" + "=" * 70) + print("Test 3: selected_component_types accuracy") + print("=" * 70) + + # 4 components, budget drops 1 + components = create_components_with_sizes( + system_prompt_chars=120, # ~80 tokens, priority=100 + tools_chars=180, # ~120 tokens, priority=80 + memory_chars=150, # ~100 tokens, priority=50 + kb_chars=300, # ~200 tokens, priority=30 + ) + + # Total budget = 350 tokens. High+Medium+Memory = 300 tokens. KB (200) won't fit. + config = ContextManagerConfig( + enabled=True, + use_context_items=True, + token_threshold=350, + strategy="token_budget", + component_budgets={ + "system_prompt": 150, + "tools": 200, + "memory": 150, + "knowledge_base": 150, + "conversation_history": 0, + }, + ) + + cm = ContextManager(config=config) + memory = build_memory() + + run_ctx = cm.prepare_run_context( + memory=memory, + fallback_system_prompt="", + components=components, + ) + + selected_types = list(run_ctx.selected_component_types) + all_types = [c.component_type for c in components] + + print(f" All component types: {all_types}") + print(f" Selected component types: {selected_types}") + print(f" Total components: {len(all_types)}, Selected: {len(selected_types)}") + + # Expect exactly 3 selected (KB dropped) + expected_count = 3 + if len(selected_types) == expected_count and "knowledge_base" not in selected_types: + print(f" \u2705 PASSED: {expected_count} types selected, KB correctly excluded") + return True + else: + print(f" \u274c FAILED: Expected {expected_count} types without KB, got {len(selected_types)}: {selected_types}") + return False + + +# --------------------------------------------------------------------------- +# Test 4: Fingerprint stability +# --------------------------------------------------------------------------- + +def test_fingerprint_stability() -> bool: + """Verify fingerprint is identical when toggling use_context_items with same components.""" + print("\n" + "=" * 70) + print("Test 4: Fingerprint stability") + print("=" * 70) + + components = create_components_with_sizes( + system_prompt_chars=150, + tools_chars=200, + memory_chars=100, + kb_chars=150, + ) + + model = create_mock_model() + + # Run with use_context_items=False + config_false = ContextManagerConfig( + enabled=True, + use_context_items=False, + token_threshold=10000, + strategy="full", + ) + cm_false = ContextManager(config=config_false) + memory_false = build_memory() + + run_ctx_false = cm_false.prepare_run_context( + memory=memory_false, + fallback_system_prompt="", + components=components, + ) + + final_ctx_false = cm_false.assemble_final_context( + model=model, + memory=memory_false, + current_run_start_idx=0, + run_context=run_ctx_false, + ) + + fp_false = final_ctx_false.evidence.stable_prefix_fingerprint + + # Run with use_context_items=True (same components, same strategy) + config_true = ContextManagerConfig( + enabled=True, + use_context_items=True, + token_threshold=10000, + strategy="full", + ) + cm_true = ContextManager(config=config_true) + memory_true = build_memory() + + run_ctx_true = cm_true.prepare_run_context( + memory=memory_true, + fallback_system_prompt="", + components=components, + ) + + final_ctx_true = cm_true.assemble_final_context( + model=model, + memory=memory_true, + current_run_start_idx=0, + run_context=run_ctx_true, + ) + + fp_true = final_ctx_true.evidence.stable_prefix_fingerprint + + print(f" Fingerprint (use_context_items=False): {fp_false[:16]}...") + print(f" Fingerprint (use_context_items=True): {fp_true[:16]}...") + + if fp_false == fp_true: + print(" \u2705 PASSED: Fingerprints are identical") + return True + else: + print(" \u274c FAILED: Fingerprints differ") + return False + + +# --------------------------------------------------------------------------- +# Test 5: History items unaffected +# --------------------------------------------------------------------------- + +def test_history_items_unaffected() -> bool: + """Verify history-projected items still appear regardless of strategy filtering.""" + print("\n" + "=" * 70) + print("Test 5: History items unaffected") + print("=" * 70) + + def mock_query_units(conversation_id: int, message_id: Optional[int] = None) -> List[Dict[str, Any]]: + return [ + { + "unit_id": 1, + "unit_type": "user_input", + "unit_content": "What is AI?", + "message_id": 1, + "step_index": 1, + }, + { + "unit_id": 2, + "unit_type": "final_answer", + "unit_content": "AI is artificial intelligence", + "message_id": 1, + "step_index": 1, + }, + { + "unit_id": 3, + "unit_type": "tool_call", + "unit_content": '{"tool_call": "search(\'machine learning\')", "execution_result": "Found 5 results about ML"}', + "message_id": 1, + "step_index": 2, + }, + ] + + history_projector = HistoryProjector(query_units_fn=mock_query_units) + + # Tight budget that drops KB component + components = create_components_with_sizes( + system_prompt_chars=150, # ~100 tokens + tools_chars=200, # ~133 tokens + memory_chars=100, # ~67 tokens + kb_chars=400, # ~267 tokens - will be dropped + ) + + config = ContextManagerConfig( + enabled=True, + use_context_items=True, + token_threshold=400, + strategy="token_budget", + component_budgets={ + "system_prompt": 200, + "tools": 200, + "memory": 150, + "knowledge_base": 150, + "conversation_history": 0, + }, + history_projector=history_projector, + ) + + cm = ContextManager(config=config) + memory = build_memory() + model = create_mock_model() + + run_ctx = cm.prepare_run_context( + memory=memory, + fallback_system_prompt="", + components=components, + ) + + # Verify KB was dropped by strategy + selected_types = set(run_ctx.selected_component_types) + print(f" Strategy-selected types: {sorted(selected_types)}") + kb_dropped = "knowledge_base" not in selected_types + + # Now assemble final context with history projector + final_ctx = cm.assemble_final_context( + model=model, + memory=memory, + current_run_start_idx=0, + run_context=run_ctx, + conversation_id=12345, + ) + + context_items = final_ctx.evidence.context_items + print(f" Total context items in evidence: {len(context_items)}") + + # Check for history items (HISTORY_TURN and TOOL_CALL_RESULT) + history_turn_items = [ + item for item in context_items + if item.item_type == ContextItemType.HISTORY_TURN + ] + tool_result_items = [ + item for item in context_items + if item.item_type == ContextItemType.TOOL_CALL_RESULT + ] + + print(f" HISTORY_TURN items: {len(history_turn_items)}") + print(f" TOOL_CALL_RESULT items: {len(tool_result_items)}") + + has_history = len(history_turn_items) > 0 + has_tool_results = len(tool_result_items) > 0 + + if kb_dropped and has_history and has_tool_results: + print(" \u2705 PASSED: KB dropped by strategy, but history items still projected") + return True + else: + reasons = [] + if not kb_dropped: + reasons.append("KB was NOT dropped (unexpected)") + if not has_history: + reasons.append("No HISTORY_TURN items found") + if not has_tool_results: + reasons.append("No TOOL_CALL_RESULT items found") + print(f" \u274c FAILED: {'; '.join(reasons)}") + return False + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main() -> int: + print("=" * 70) + print("Strategy Filtering Fix Verification Tests") + print("=" * 70) + + results = { + "Budget-pressure equivalence": test_budget_pressure_equivalence(), + "Strategy drop propagation": test_strategy_drop_propagation(), + "selected_component_types accuracy": test_selected_component_types_accuracy(), + "Fingerprint stability": test_fingerprint_stability(), + "History items unaffected": test_history_items_unaffected(), + } + + print("\n" + "=" * 70) + print("SUMMARY") + print("=" * 70) + + passed_count = 0 + total_count = len(results) + + for test_name, passed in results.items(): + status = "\u2705 PASSED" if passed else "\u274c FAILED" + print(f" {status}: {test_name}") + if passed: + passed_count += 1 + + print("\n" + "=" * 70) + if passed_count == total_count: + print(f"ALL TESTS PASSED \u2705 ({passed_count}/{total_count})") + else: + print(f"SOME TESTS FAILED \u274c ({passed_count}/{total_count} passed)") + print("=" * 70) + + return 0 if passed_count == total_count else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/test/sdk/core/agents/test_history_projector.py b/test/sdk/core/agents/test_history_projector.py new file mode 100644 index 000000000..9147c82b1 --- /dev/null +++ b/test/sdk/core/agents/test_history_projector.py @@ -0,0 +1,616 @@ +"""Unit tests for HistoryProjector. + +Verifies that HistoryProjector correctly projects conversation history +units into ContextItem instances for model_context and chat purposes. +""" + +import json +import pytest +from unittest.mock import MagicMock + +from nexent.core.agents.agent_context import ContextManager +from nexent.core.agents.agent_model import ( + MemoryComponent, + SystemPromptComponent, +) +from nexent.core.agents.context.history_projector import HistoryProjector +from nexent.core.agents.context.context_item import ( + AuthorityTier, + ContextItem, + ContextItemType, + RepresentationTier, +) +from nexent.core.agents.context.handlers import register_all +from nexent.core.agents.summary_config import ContextManagerConfig + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def make_unit( + unit_id=1, + unit_type="user_input", + unit_content="hello", + message_id=1, + step_index=1, +): + """Build a minimal unit dict matching the database row shape.""" + return { + "unit_id": unit_id, + "unit_type": unit_type, + "unit_content": unit_content, + "message_id": message_id, + "step_index": step_index, + } + + +def make_tool_call_unit( + unit_id=1, + tool_call="web_search('AI')", + execution_result="result: found AI", + message_id=1, + step_index=1, +): + """Build a merged tool_call unit with JSON unit_content.""" + return { + "unit_id": unit_id, + "unit_type": "tool_call", + "unit_content": json.dumps({ + "tool_call": tool_call, + "execution_result": execution_result, + }), + "message_id": message_id, + "step_index": step_index, + } + + +def make_query_fn(units): + """Return a query_units_fn closure over a fixed list of units.""" + def query_fn(conversation_id, message_id=None): + if message_id is not None: + return [u for u in units if u.get("message_id") == message_id] + return units + return query_fn + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture(autouse=True) +def ensure_handlers(): + """Register all item handlers so ItemHandlerRegistry.get() succeeds.""" + register_all() + + +# =================================================================== +# 1. Basic Projection Tests +# =================================================================== + +class TestBasicProjection: + """Tests for core model_context projection behavior.""" + + def test_project_model_context_empty_units(self): + """Empty units list returns empty items.""" + projector = HistoryProjector(make_query_fn([])) + items = projector.project(conversation_id=1, purpose="model_context") + assert items == [] + + def test_project_model_context_single_turn(self): + """Single user_input + final_answer produces one HISTORY_TURN.""" + units = [ + make_unit(unit_id=1, unit_type="user_input", unit_content="What is AI?", message_id=1, step_index=1), + make_unit(unit_id=2, unit_type="final_answer", unit_content="AI is artificial intelligence.", message_id=1, step_index=1), + ] + projector = HistoryProjector(make_query_fn(units)) + items = projector.project(conversation_id=1, purpose="model_context") + + assert len(items) == 1 + assert items[0].item_type == ContextItemType.HISTORY_TURN + assert items[0].content["user_query"] == "What is AI?" + assert items[0].content["assistant_response"] == "AI is artificial intelligence." + + def test_project_model_context_multiple_turns(self): + """Multiple messages/steps produce multiple HISTORY_TURNs.""" + units = [ + make_unit(unit_id=1, unit_type="user_input", unit_content="Q1", message_id=1, step_index=1), + make_unit(unit_id=2, unit_type="final_answer", unit_content="A1", message_id=1, step_index=1), + make_unit(unit_id=3, unit_type="user_input", unit_content="Q2", message_id=1, step_index=2), + make_unit(unit_id=4, unit_type="final_answer", unit_content="A2", message_id=1, step_index=2), + make_unit(unit_id=5, unit_type="user_input", unit_content="Q3", message_id=2, step_index=1), + make_unit(unit_id=6, unit_type="final_answer", unit_content="A3", message_id=2, step_index=1), + ] + projector = HistoryProjector(make_query_fn(units)) + items = projector.project(conversation_id=1, purpose="model_context") + + history_turns = [i for i in items if i.item_type == ContextItemType.HISTORY_TURN] + assert len(history_turns) == 3 + + def test_project_model_context_excludes_thinking(self): + """model_output_thinking and model_output_deep_thinking are NOT included in HISTORY_TURN.""" + units = [ + make_unit(unit_id=1, unit_type="user_input", unit_content="Think hard", message_id=1, step_index=1), + make_unit(unit_id=2, unit_type="model_output_thinking", unit_content="internal thought", message_id=1, step_index=1), + make_unit(unit_id=3, unit_type="model_output_deep_thinking", unit_content="deep thought", message_id=1, step_index=1), + make_unit(unit_id=4, unit_type="final_answer", unit_content="The answer.", message_id=1, step_index=1), + ] + projector = HistoryProjector(make_query_fn(units)) + items = projector.project(conversation_id=1, purpose="model_context") + + assert len(items) == 1 + turn = items[0] + assert turn.item_type == ContextItemType.HISTORY_TURN + assert "internal thought" not in str(turn.content) + assert "deep thought" not in str(turn.content) + assert turn.content["user_query"] == "Think hard" + assert turn.content["assistant_response"] == "The answer." + + def test_project_invalid_purpose(self): + """Unknown purpose raises ValueError.""" + projector = HistoryProjector(make_query_fn([])) + with pytest.raises(ValueError, match="Unknown purpose"): + projector.project(conversation_id=1, purpose="invalid_purpose") + + +# =================================================================== +# 2. Tool Call Result Tests +# =================================================================== + +class TestToolCallResult: + """Tests for merged tool_call rows into TOOL_CALL_RESULT items.""" + + def test_tool_call_result_single(self): + """Single merged tool_call row produces one TOOL_CALL_RESULT.""" + units = [ + make_unit(unit_id=1, unit_type="user_input", unit_content="search", message_id=1, step_index=1), + make_unit(unit_id=2, unit_type="final_answer", unit_content="done", message_id=1, step_index=1), + make_tool_call_unit(unit_id=3, tool_call="web_search('AI')", execution_result="result: found AI", message_id=1, step_index=1), + ] + projector = HistoryProjector(make_query_fn(units)) + items = projector.project(conversation_id=1, purpose="model_context") + + tool_results = [i for i in items if i.item_type == ContextItemType.TOOL_CALL_RESULT] + assert len(tool_results) == 1 + assert tool_results[0].content["tool_call"] == "web_search('AI')" + assert tool_results[0].content["execution_result"] == "result: found AI" + + def test_tool_call_result_multiple(self): + """Multiple merged tool_call rows produce multiple TOOL_CALL_RESULTs.""" + units = [ + make_unit(unit_id=1, unit_type="user_input", unit_content="multi", message_id=1, step_index=1), + make_unit(unit_id=2, unit_type="final_answer", unit_content="done", message_id=1, step_index=1), + make_tool_call_unit(unit_id=3, tool_call="tool_a()", execution_result="result_a", message_id=1, step_index=1), + make_tool_call_unit(unit_id=4, tool_call="tool_b()", execution_result="result_b", message_id=1, step_index=1), + ] + projector = HistoryProjector(make_query_fn(units)) + items = projector.project(conversation_id=1, purpose="model_context") + + tool_results = [i for i in items if i.item_type == ContextItemType.TOOL_CALL_RESULT] + assert len(tool_results) == 2 + item_ids = {i.item_id for i in tool_results} + assert item_ids == {"tool_call_result:3", "tool_call_result:4"} + + def test_tool_call_result_invalid_json_skipped(self): + """tool_call row with invalid JSON is skipped.""" + units = [ + make_unit(unit_id=1, unit_type="user_input", unit_content="q", message_id=1, step_index=1), + make_unit(unit_id=2, unit_type="final_answer", unit_content="a", message_id=1, step_index=1), + { + "unit_id": 3, + "unit_type": "tool_call", + "unit_content": "not valid json", + "message_id": 1, + "step_index": 1, + }, + ] + projector = HistoryProjector(make_query_fn(units)) + items = projector.project(conversation_id=1, purpose="model_context") + + tool_results = [i for i in items if i.item_type == ContextItemType.TOOL_CALL_RESULT] + assert len(tool_results) == 0 + + def test_tool_call_result_no_tool_call_units(self): + """No tool_call units means no TOOL_CALL_RESULT items.""" + units = [ + make_unit(unit_id=1, unit_type="user_input", unit_content="q", message_id=1, step_index=1), + make_unit(unit_id=2, unit_type="final_answer", unit_content="a", message_id=1, step_index=1), + ] + projector = HistoryProjector(make_query_fn(units)) + items = projector.project(conversation_id=1, purpose="model_context") + + tool_results = [i for i in items if i.item_type == ContextItemType.TOOL_CALL_RESULT] + assert len(tool_results) == 0 + + +# =================================================================== +# 3. Chat Context Tests +# =================================================================== + +class TestChatContext: + """Tests for chat purpose which includes thinking units.""" + + def test_chat_context_includes_thinking(self): + """Chat context includes model_output_thinking units.""" + units = [ + make_unit(unit_id=1, unit_type="user_input", unit_content="Explain", message_id=1, step_index=1), + make_unit(unit_id=2, unit_type="model_output_thinking", unit_content="Let me think...", message_id=1, step_index=1), + make_unit(unit_id=3, unit_type="final_answer", unit_content="The answer is X.", message_id=1, step_index=1), + ] + projector = HistoryProjector(make_query_fn(units)) + items = projector.project(conversation_id=1, purpose="chat") + + assert len(items) == 1 + turn = items[0] + assert turn.item_type == ContextItemType.HISTORY_TURN + unit_types = [u["type"] for u in turn.content["units"]] + assert "user_input" in unit_types + assert "model_output_thinking" in unit_types + assert "final_answer" in unit_types + assert turn.metadata.get("includes_thinking") is True + + def test_chat_context_empty_units(self): + """Empty units returns empty items.""" + projector = HistoryProjector(make_query_fn([])) + items = projector.project(conversation_id=1, purpose="chat") + assert items == [] + + +# =================================================================== +# 4. Grouping Tests +# =================================================================== + +class TestGrouping: + """Tests for _group_by_message handling of None/falsy IDs.""" + + def test_group_by_message_handles_none_message_id(self): + """Units with None message_id are grouped under message_id=0.""" + units = [ + make_unit(unit_id=1, unit_type="user_input", unit_content="q", message_id=None, step_index=1), + make_unit(unit_id=2, unit_type="final_answer", unit_content="a", message_id=None, step_index=1), + ] + projector = HistoryProjector(make_query_fn(units)) + items = projector.project(conversation_id=1, purpose="model_context") + + assert len(items) == 1 + assert items[0].item_id == "history_turn:0:1" + assert items[0].metadata["message_id"] == 0 + + def test_group_by_message_handles_none_step_index(self): + """Units with None step_index are grouped under step_index=0.""" + units = [ + make_unit(unit_id=1, unit_type="user_input", unit_content="q", message_id=1, step_index=None), + make_unit(unit_id=2, unit_type="final_answer", unit_content="a", message_id=1, step_index=None), + ] + projector = HistoryProjector(make_query_fn(units)) + items = projector.project(conversation_id=1, purpose="model_context") + + assert len(items) == 1 + assert items[0].item_id == "history_turn:1:0" + assert items[0].metadata["step_index"] == 0 + + +# =================================================================== +# 5. Item Validation Tests +# =================================================================== + +class TestItemValidation: + """Verify field-level correctness of produced ContextItems.""" + + def test_history_turn_item_fields(self): + """Verify item_id format, item_type, source_refs, authority_tier, content structure.""" + units = [ + make_unit(unit_id=10, unit_type="user_input", unit_content="query text", message_id=3, step_index=2), + make_unit(unit_id=20, unit_type="final_answer", unit_content="answer text", message_id=3, step_index=2), + ] + projector = HistoryProjector(make_query_fn(units)) + items = projector.project(conversation_id=1, purpose="model_context") + + assert len(items) == 1 + item = items[0] + + assert item.item_id == "history_turn:3:2" + assert item.item_type == ContextItemType.HISTORY_TURN + assert item.source_refs == ["unit:10", "unit:20"] + assert item.authority_tier == AuthorityTier.AGENT_INFERENCE + assert item.minimum_fidelity == RepresentationTier.STRUCTURED + assert item.current_representation == RepresentationTier.FULL + assert item.content == { + "user_query": "query text", + "assistant_response": "answer text", + } + assert item.token_estimate == (len("query text") + len("answer text")) // 4 + assert item.metadata == {"message_id": 3, "step_index": 2} + + def test_tool_call_result_item_fields(self): + """Verify item_id format, item_type, source_refs, authority_tier, content structure.""" + units = [ + make_unit(unit_id=1, unit_type="user_input", unit_content="q", message_id=1, step_index=1), + make_unit(unit_id=2, unit_type="final_answer", unit_content="a", message_id=1, step_index=1), + make_tool_call_unit(unit_id=5, tool_call="search('x')", execution_result="found: x", message_id=1, step_index=1), + ] + projector = HistoryProjector(make_query_fn(units)) + items = projector.project(conversation_id=1, purpose="model_context") + + tool_items = [i for i in items if i.item_type == ContextItemType.TOOL_CALL_RESULT] + assert len(tool_items) == 1 + item = tool_items[0] + + assert item.item_id == "tool_call_result:5" + assert item.item_type == ContextItemType.TOOL_CALL_RESULT + assert item.source_refs == ["unit:5"] + assert item.authority_tier == AuthorityTier.TOOL_RESULT + assert item.minimum_fidelity == RepresentationTier.STRUCTURED + assert item.current_representation == RepresentationTier.FULL + assert item.content == { + "tool_call": "search('x')", + "execution_result": "found: x", + } + assert item.metadata == { + "message_id": 1, + "step_index": 1, + } + + +# =================================================================== +# 6. Chat Projection Completeness Tests +# =================================================================== + +class TestChatProjectionCompleteness: + """Verify chat projection produces sufficient content for UI reconstruction.""" + + def test_chat_projection_covers_all_displayable_unit_types(self): + """Chat projection includes user_input, model_output_thinking, + model_output_code, and final_answer — the correct display set. + + model_output_deep_thinking is intentionally excluded from the + current implementation's relevant_types. + """ + units = [ + make_unit(unit_id=1, unit_type="user_input", unit_content="Explain recursion", message_id=1, step_index=1), + make_unit(unit_id=2, unit_type="model_output_thinking", unit_content="Let me reason...", message_id=1, step_index=1), + make_unit(unit_id=3, unit_type="model_output_code", unit_content="def recurse(): ...", message_id=1, step_index=1), + make_unit(unit_id=4, unit_type="model_output_deep_thinking", unit_content="deep analysis", message_id=1, step_index=1), + make_unit(unit_id=5, unit_type="final_answer", unit_content="Recursion is...", message_id=1, step_index=1), + ] + projector = HistoryProjector(make_query_fn(units)) + items = projector.project(conversation_id=1, purpose="chat") + + assert len(items) == 1 + turn = items[0] + unit_types = [u["type"] for u in turn.content["units"]] + + assert "user_input" in unit_types + assert "model_output_thinking" in unit_types + assert "model_output_code" in unit_types + assert "final_answer" in unit_types + assert "model_output_deep_thinking" not in unit_types + + def test_chat_projection_preserves_unit_ordering(self): + """Units within a chat turn maintain their original ordering by position.""" + units = [ + make_unit(unit_id=1, unit_type="user_input", unit_content="Q", message_id=1, step_index=1), + make_unit(unit_id=2, unit_type="model_output_thinking", unit_content="think first", message_id=1, step_index=1), + make_unit(unit_id=3, unit_type="model_output_code", unit_content="code second", message_id=1, step_index=1), + make_unit(unit_id=4, unit_type="final_answer", unit_content="answer third", message_id=1, step_index=1), + ] + projector = HistoryProjector(make_query_fn(units)) + items = projector.project(conversation_id=1, purpose="chat") + + assert len(items) == 1 + content_units = items[0].content["units"] + types_in_order = [u["type"] for u in content_units] + + assert types_in_order == [ + "user_input", + "model_output_thinking", + "model_output_code", + "final_answer", + ] + + def test_chat_projection_includes_metadata_for_reconstruction(self): + """Metadata contains message_id, step_index, and includes_thinking flag + needed by any adapter converting to frontend format.""" + units = [ + make_unit(unit_id=1, unit_type="user_input", unit_content="Q", message_id=7, step_index=3), + make_unit(unit_id=2, unit_type="final_answer", unit_content="A", message_id=7, step_index=3), + ] + projector = HistoryProjector(make_query_fn(units)) + items = projector.project(conversation_id=1, purpose="chat") + + assert len(items) == 1 + meta = items[0].metadata + assert meta["message_id"] == 7 + assert meta["step_index"] == 3 + assert meta["includes_thinking"] is True + + def test_chat_projection_source_refs_cover_all_included_units(self): + """source_refs list contains references to ALL units included in the content.""" + units = [ + make_unit(unit_id=10, unit_type="user_input", unit_content="Q", message_id=1, step_index=1), + make_unit(unit_id=20, unit_type="model_output_thinking", unit_content="think", message_id=1, step_index=1), + make_unit(unit_id=30, unit_type="model_output_code", unit_content="code", message_id=1, step_index=1), + make_unit(unit_id=40, unit_type="final_answer", unit_content="A", message_id=1, step_index=1), + ] + projector = HistoryProjector(make_query_fn(units)) + items = projector.project(conversation_id=1, purpose="chat") + + assert len(items) == 1 + turn = items[0] + + expected_refs = {"unit:10", "unit:20", "unit:30", "unit:40"} + assert set(turn.source_refs) == expected_refs + assert len(turn.source_refs) == len(turn.content["units"]) + + +# =================================================================== +# 7. End-to-End Integration Tests +# =================================================================== + +class MockHistoryProjector: + """Mock projector for integration tests — no real DB needed.""" + + def __init__(self, items=None, should_fail=False): + self._items = items or [] + self._should_fail = should_fail + + def project(self, conversation_id, message_id=None, purpose="model_context"): + if self._should_fail: + raise RuntimeError("Simulated projection failure") + return self._items + + +class TestEndToEndIntegration: + """Full flow: ContextManager + HistoryProjector -> assemble_final_context -> FinalContext.""" + + def test_assemble_final_context_with_history_projector(self): + """FinalContext.evidence.context_items contains both component-projected + AND history-projected items when conversation_id is provided.""" + history_items = [ + ContextItem( + item_id="chat_turn:1:1", + item_type=ContextItemType.HISTORY_TURN, + source_refs=["unit:1", "unit:2"], + authority_tier=AuthorityTier.AGENT_INFERENCE, + minimum_fidelity=RepresentationTier.STRUCTURED, + current_representation=RepresentationTier.FULL, + content={"user_query": "Hello", "assistant_response": "Hi there"}, + token_estimate=4, + metadata={"message_id": 1, "step_index": 1}, + ), + ] + mock_projector = MockHistoryProjector(items=history_items) + + config = ContextManagerConfig( + enabled=True, + token_threshold=10000, + use_context_items=True, + history_projector=mock_projector, + ) + manager = ContextManager(config=config) + + manager.register_component(SystemPromptComponent(content="You are helpful")) + manager.register_component( + MemoryComponent( + memories=[{"content": "User prefers Python", "memory_type": "user"}], + formatted_content="User preferences: Python", + ) + ) + + memory = MagicMock() + memory.system_prompt = None + memory.steps = [] + + final = manager.assemble_final_context( + model=None, + memory=memory, + current_run_start_idx=0, + tools=[], + conversation_id=123, + ) + + assert final is not None + assert len(final.evidence.context_items) > 0 + + component_types = {item.item_type for item in final.evidence.context_items} + assert ContextItemType.SYSTEM_PROMPT in component_types + + history_turn_items = [ + item for item in final.evidence.context_items + if item.item_id == "chat_turn:1:1" + ] + assert len(history_turn_items) == 1 + assert history_turn_items[0].content["user_query"] == "Hello" + + def test_assemble_final_context_without_conversation_id_skips_history(self): + """Without conversation_id, only component-projected items appear.""" + history_items = [ + ContextItem( + item_id="chat_turn:1:1", + item_type=ContextItemType.HISTORY_TURN, + source_refs=["unit:1"], + authority_tier=AuthorityTier.AGENT_INFERENCE, + minimum_fidelity=RepresentationTier.STRUCTURED, + current_representation=RepresentationTier.FULL, + content={"user_query": "Should not appear"}, + token_estimate=3, + metadata={"message_id": 1, "step_index": 1}, + ), + ] + mock_projector = MockHistoryProjector(items=history_items) + + config = ContextManagerConfig( + enabled=True, + token_threshold=10000, + use_context_items=True, + history_projector=mock_projector, + ) + manager = ContextManager(config=config) + + manager.register_component(SystemPromptComponent(content="You are helpful")) + + memory = MagicMock() + memory.system_prompt = None + memory.steps = [] + + final = manager.assemble_final_context( + model=None, + memory=memory, + current_run_start_idx=0, + tools=[], + ) + + assert final is not None + assert len(final.evidence.context_items) > 0 + + history_turn_items = [ + item for item in final.evidence.context_items + if item.item_id == "chat_turn:1:1" + ] + assert len(history_turn_items) == 0 + + component_types = {item.item_type for item in final.evidence.context_items} + assert ContextItemType.SYSTEM_PROMPT in component_types + + def test_assemble_final_context_history_projector_failure_graceful(self): + """When history projector raises, assemble_final_context continues + with component items only — no crash.""" + mock_projector = MockHistoryProjector(should_fail=True) + + config = ContextManagerConfig( + enabled=True, + token_threshold=10000, + use_context_items=True, + history_projector=mock_projector, + ) + manager = ContextManager(config=config) + + manager.register_component(SystemPromptComponent(content="You are helpful")) + manager.register_component( + MemoryComponent( + memories=[{"content": "Important fact", "memory_type": "user"}], + formatted_content="Memory: Important fact", + ) + ) + + memory = MagicMock() + memory.system_prompt = None + memory.steps = [] + + final = manager.assemble_final_context( + model=None, + memory=memory, + current_run_start_idx=0, + tools=[], + conversation_id=123, + ) + + assert final is not None + assert len(final.messages) > 0 + assert len(final.evidence.context_items) > 0 + component_types = {item.item_type for item in final.evidence.context_items} + assert ContextItemType.SYSTEM_PROMPT in component_types + + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/test/sdk/core/agents/test_sdk_langfuse_integration.py b/test/sdk/core/agents/test_sdk_langfuse_integration.py new file mode 100644 index 000000000..cdf3d6519 --- /dev/null +++ b/test/sdk/core/agents/test_sdk_langfuse_integration.py @@ -0,0 +1,222 @@ +""" +SDK Integration Test with Real Model and LangFuse Tracing + +This test verifies: +1. Real LLM model execution through the SDK +2. OpenTelemetry instrumentation captures spans correctly +3. Context management with use_context_items=True works end-to-end +4. LangFuse receives and displays traces properly + +Requirements: +- OPENAI_API_KEY environment variable must be set +- LangFuse configuration in .env (OTEL_EXPORTER_OTLP_ENDPOINT, etc.) +- Network access to OpenAI API and LangFuse + +NOTE: This test is marked as 'local_only' and will be skipped in CI environments. +Run locally with: pytest -m local_only test/sdk/core/agents/test_sdk_langfuse_integration.py +""" + +import os +import pytest +from unittest.mock import MagicMock + +from nexent.core.agents.agent_context import ContextManager +from nexent.core.agents.agent_model import ( + MemoryComponent, + SystemPromptComponent, + ToolsComponent, +) +from nexent.core.agents.summary_config import ContextManagerConfig +from nexent.core.agents.context.handlers import register_all +from nexent.core.context_runtime.managed.runtime import ManagedContextRuntime +from nexent.core.models.openai_llm import OpenAIModel + + +pytestmark = pytest.mark.local_only + + +@pytest.fixture(autouse=True) +def ensure_handlers_registered(): + """Ensure all context handlers are registered.""" + register_all() + + +@pytest.fixture +def real_model(): + """Create a real OpenAI model instance for integration testing.""" + api_key = os.environ.get("OPENAI_API_KEY") + if not api_key: + pytest.skip("OPENAI_API_KEY not set - skipping real model integration test") + + return OpenAIModel( + model_id="gpt-3.5-turbo", + api_key=api_key, + model_name="gpt-3.5-turbo", + url="https://api.openai.com/v1", + temperature=0.1, + top_p=0.95, + ) + + +class TestSDKLangFuseIntegration: + """Integration tests with real model and LangFuse tracing.""" + + def test_context_items_with_real_model(self, real_model): + """Test context management with real model execution.""" + config = ContextManagerConfig( + enabled=True, + token_threshold=10000, + use_context_items=True, + ) + manager = ContextManager(config=config) + + system_prompt = SystemPromptComponent( + content="You are a helpful assistant. Answer questions concisely." + ) + tools = ToolsComponent( + tools=[{"name": "search", "description": "Search the web"}], + formatted_description="Available tools: search", + ) + memory_comp = MemoryComponent( + memories=[{"content": "User prefers Python", "memory_type": "user"}], + formatted_content="User preferences: Python programming language", + ) + + manager.register_component(system_prompt) + manager.register_component(tools) + manager.register_component(memory_comp) + + runtime = ManagedContextRuntime( + manager, + components=[system_prompt, tools, memory_comp] + ) + + memory = MagicMock() + memory.system_prompt = None + memory.steps = [] + + runtime.prepare_run(memory=memory, fallback_system_prompt="You are helpful") + + final = runtime.prepare_step( + model=real_model, + memory=memory, + current_run_start_idx=0, + tools=[], + ) + + assert final is not None + assert len(final.messages) > 0 + assert final.evidence.context_items is not None + assert len(final.evidence.context_items) > 0 + + roles = [msg["role"] for msg in final.messages] + assert "system" in roles + assert "user" in roles + + item_types = [item.item_type for item in final.evidence.context_items] + assert len(item_types) > 0 + + def test_context_compression_with_real_model(self, real_model): + """Test context compression with real model when token threshold is exceeded.""" + config = ContextManagerConfig( + enabled=True, + token_threshold=500, + use_context_items=True, + ) + manager = ContextManager(config=config) + + system_prompt = SystemPromptComponent(content="You are a helpful assistant.") + manager.register_component(system_prompt) + + runtime = ManagedContextRuntime(manager, components=[system_prompt]) + + from smolagents.memory import ActionStep, TaskStep + + memory = MagicMock() + memory.system_prompt = None + + steps = [] + task_step = TaskStep(task="Solve a complex problem") + steps.append(task_step) + + for i in range(10): + action_step = ActionStep( + step_number=i, + timing=MagicMock(), + code_action=f"action_{i}", + observations="This is a very long observation with lots of text. " * 50 + ) + steps.append(action_step) + + memory.steps = steps + + runtime.prepare_run(memory=memory, fallback_system_prompt="You are helpful") + + final = runtime.prepare_step( + model=real_model, + memory=memory, + current_run_start_idx=0, + tools=[], + ) + + assert final is not None + assert len(final.messages) > 0 + + stats = runtime.compression_stats() + assert stats['calls'] > 0 or stats['cache_hits'] > 0 + + def test_history_projector_integration(self, real_model): + """Test HistoryProjector with real model execution.""" + from nexent.core.agents.context.history_projector import HistoryProjector + + config = ContextManagerConfig( + enabled=True, + token_threshold=10000, + use_context_items=True, + ) + manager = ContextManager(config=config) + + def mock_query_fn(conversation_id, message_id=None): + return [ + { + "unit_id": 1, + "unit_type": "user_input", + "unit_content": "What is Python?", + "message_id": 1, + "step_index": 1, + }, + { + "unit_id": 2, + "unit_type": "final_answer", + "unit_content": "Python is a programming language.", + "message_id": 1, + "step_index": 2, + }, + ] + + history_projector = HistoryProjector(query_units_fn=mock_query_fn) + config.history_projector = history_projector + + system_prompt = SystemPromptComponent(content="You are helpful") + manager.register_component(system_prompt) + + runtime = ManagedContextRuntime(manager, components=[system_prompt], conversation_id=123) + + memory = MagicMock() + memory.system_prompt = None + memory.steps = [] + + runtime.prepare_run(memory=memory, fallback_system_prompt="You are helpful") + + final = runtime.prepare_step( + model=real_model, + memory=memory, + current_run_start_idx=0, + tools=[], + ) + + assert final is not None + assert final.evidence.context_items is not None + + item_types = [item.item_type for item in final.evidence.context_items] + assert len(item_types) >= 1