diff --git a/.gitignore b/.gitignore index 3919cf531f..316c6cd418 100644 --- a/.gitignore +++ b/.gitignore @@ -205,3 +205,14 @@ specs/ .codex/ .ttadk/ tests/integration/.tmp_*/ + +# Beads / Dolt files (added by bd init) +.dolt/ +*.db +.beads-credential-key + +# Beads issue tracker local state +.beads/ + +# LoopX local control-plane state +.loopx/ diff --git a/README.md b/README.md index 6040ca1841..79977e3c87 100644 --- a/README.md +++ b/README.md @@ -341,6 +341,10 @@ If you prefer manual configuration, create `~/.openviking/ov.conf`, remove the c > **Note**: For embedding models, supported providers are `volcengine` (Doubao), `openai`, `azure`, `jina`, `ollama`, `voyage`, `dashscope`, `minimax`, `cohere`, `vikingdb`, `gemini` (requires `pip install "google-genai>=1.0.0"`), `litellm`, and `local`. For VLM models, common providers include `volcengine`, `openai`, `openai-codex`, `kimi`, and `glm`. +> **Memory config**: OpenViking always uses the v3 memory extraction pipeline. The legacy `memory.version` setting is deprecated and ignored; existing configs that set it still load without changing behavior. + +> **Memory schema scope**: Memory schema YAML defaults to `stage: "user"` and `peer_enabled: true`. Use `stage: "agent"` for execution-derived schemas such as trajectories/experiences, and set `peer_enabled: false` when a schema must stay in the current user's memory directory instead of peer directories. + ##### Server Configuration Examples 👇 Expand to see the configuration example for your model service: diff --git a/README_CN.md b/README_CN.md index 4a95d4bf1d..70466926cd 100644 --- a/README_CN.md +++ b/README_CN.md @@ -297,6 +297,10 @@ openviking-server doctor > **注意**:对于 embedding 模型,支持 `volcengine`(豆包)、`openai`、`azure`、`jina`、`ollama`、`voyage`、`dashscope`、`minimax`、`cohere`、`vikingdb`、`gemini`(需 `pip install "google-genai>=1.0.0"`)、`litellm` 和 `local`。对于 VLM 模型,常见提供商包括 `volcengine`、`openai`、`openai-codex`、`kimi`、`glm`。 +> **Memory 配置**:OpenViking 始终使用 v3 记忆抽取链路。旧的 `memory.version` 配置项已废弃且会被忽略;已有配置中保留该字段仍可正常加载,但不会改变行为。 + +> **Memory schema 作用域**:Memory schema YAML 默认 `stage: "user"` 且 `peer_enabled: true`。trajectories/experiences 这类执行派生 schema 使用 `stage: "agent"`;如果某类记忆必须保留在当前用户目录而不是 peer 目录,设置 `peer_enabled: false`。 + ##### 服务器配置示例 👇 展开查看您的模型服务的配置示例: diff --git a/benchmark/__init__.py b/benchmark/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/benchmark/locomo/vikingbot/import_to_ov.py b/benchmark/locomo/vikingbot/import_to_ov.py index 95746f9276..4dd1f08ee8 100644 --- a/benchmark/locomo/vikingbot/import_to_ov.py +++ b/benchmark/locomo/vikingbot/import_to_ov.py @@ -13,24 +13,166 @@ import argparse import asyncio +import contextlib import csv import json +import re import sys import time import traceback from datetime import datetime, timedelta from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Awaitable, Callable, Dict, List, Optional, Tuple, TypeVar + +import httpx +from progress_utils import ( + AsyncProgressTracker, + make_three_state_progress, + should_show_progress, +) import openviking as ov +TRACE_ID_RE = re.compile(r"\btrace_?id[=:\s]+([^,\s:)\]]+)") + +_RetryResult = TypeVar("_RetryResult") +_TRANSIENT_HTTP_ERRORS = ( + httpx.ReadError, + httpx.ConnectError, + httpx.RemoteProtocolError, + httpx.TimeoutException, +) + + +def _retry_delay(attempt: int, *, base: float = 0.5, cap: float = 8.0) -> float: + """Deterministic exponential backoff capped for benchmark imports.""" + return min(cap, base * (2 ** max(0, attempt - 1))) + + +def _transport_error_message(exc: BaseException) -> str: + text = str(exc) + return f"{type(exc).__name__}: {text}" if text else repr(exc) + + +async def _retry_transient_http( + label: str, + operation: Callable[[], Awaitable[_RetryResult]], + *, + attempts: int = 5, +) -> _RetryResult: + """Retry transient HTTP transport failures with exponential backoff. + + Used only around operations where retrying is safe enough for import: + create_session has no user-visible payload yet, commit is protected by + session archive semantics, and task polling is read-only. Message writes + are intentionally not wrapped here because the API is append-only and a + lost response could otherwise duplicate messages. + """ + last_exc: BaseException | None = None + for attempt in range(1, max(1, attempts) + 1): + try: + return await operation() + except _TRANSIENT_HTTP_ERRORS as exc: + last_exc = exc + if attempt >= attempts: + break + await asyncio.sleep(_retry_delay(attempt)) + assert last_exc is not None + print( + f" -> [RETRY] {label} failed after {attempts} attempts " + f"with {_transport_error_message(last_exc)}", + file=sys.stderr, + ) + raise last_exc + + +async def _list_session_commit_tasks( + client: Any, + *, + session_id: str, + limit: int = 20, +) -> list[dict[str, Any]]: + """Best-effort raw task listing for recovering from a lost commit response.""" + http = getattr(client, "_http", None) + if http is None: + return [] + response = await http.get( + "/api/v1/tasks", + params={ + "task_type": "session_commit", + "resource_id": session_id, + "limit": limit, + }, + ) + data = client._handle_response(response) + if isinstance(data, list): + return [item for item in data if isinstance(item, dict)] + return [] + + +def _latest_active_commit_task(tasks: list[dict[str, Any]]) -> dict[str, Any] | None: + for task in tasks: + if task.get("status") in {"pending", "running", "completed"}: + return task + return tasks[0] if tasks else None + + +async def _recover_commit_after_transport_error( + client: Any, + *, + session_id: str, + previous_commit_count: int, + attempts: int = 6, +) -> dict[str, Any] | None: + """Recover a commit result when the POST succeeded but response read failed. + + Commit phase 1 clears/archives live messages and creates a session_commit + task. If a ReadError happens while reading the response, retrying POST can + race with that state transition. Prefer detecting the already-created task + and return an equivalent accepted result. + """ + for attempt in range(1, attempts + 1): + try: + session = await _retry_transient_http( + f"get_session_after_commit_read_error session={session_id}", + lambda: client.get_session(session_id), + attempts=2, + ) + commit_count = int(session.get("commit_count") or 0) + if commit_count > previous_commit_count: + tasks = await _retry_transient_http( + f"list_commit_tasks session={session_id}", + lambda: _list_session_commit_tasks(client, session_id=session_id), + attempts=2, + ) + task = _latest_active_commit_task(tasks) + if task is not None: + return { + "status": "accepted", + "session_id": session_id, + "task_id": task.get("task_id"), + "trace_id": "", + "recovered_after_read_error": True, + } + return { + "status": "accepted", + "session_id": session_id, + "task_id": None, + "trace_id": "", + "recovered_after_read_error": True, + } + except _TRANSIENT_HTTP_ERRORS: + pass + await asyncio.sleep(_retry_delay(attempt, base=0.5, cap=5.0)) + return None + def _get_session_number(session_key: str) -> int: """Extract session number from session key.""" return int(session_key.split("_")[1]) -def build_memory_policy(group_chat: bool) -> Dict[str, Dict[str, bool]]: +def build_memory_policy(group_chat: bool) -> Dict[str, Any]: """Build session/commit memory policy for benchmark ingest. LoCoMo eval isolates samples through peer memory. In non-group mode the @@ -38,11 +180,16 @@ def build_memory_policy(group_chat: bool) -> Dict[str, Dict[str, bool]]: speaker. Do not write benchmark memories into the current User self memory, otherwise all samples imported by the same User API key become visible to every question. + + Import only user-facing LoCoMo memories; skip cases/train/runtime memory + types so benchmark ingest does not create training cases. """ del group_chat return { "self": {"enabled": False}, "peer": {"enabled": True}, + "working_memory": {"enabled": False}, + "memory_types": ["entities", "events", "preferences", "profile"], } @@ -192,7 +339,7 @@ def build_session_messages( # --------------------------------------------------------------------------- -def load_success_csv(csv_path: str = "./result/import_success.csv") -> set: +def load_success_csv(csv_path: str = "./result/locomo/import_success.csv") -> set: """加载成功导入的CSV记录,返回已成功的键集合""" success_keys = set() if Path(csv_path).exists(): @@ -205,7 +352,7 @@ def load_success_csv(csv_path: str = "./result/import_success.csv") -> set: def write_success_record( - record: Dict[str, Any], csv_path: str = "./result/import_success.csv" + record: Dict[str, Any], csv_path: str = "./result/locomo/import_success.csv" ) -> None: """写入成功记录到CSV文件""" file_exists = Path(csv_path).exists() @@ -250,7 +397,7 @@ def write_success_record( def write_error_record( - record: Dict[str, Any], error_path: str = "./result/import_errors.log" + record: Dict[str, Any], error_path: str = "./result/locomo/import_errors.log" ) -> None: """写入错误记录到日志文件""" with open(error_path, "a", encoding="utf-8") as f: @@ -261,7 +408,38 @@ def write_error_record( f.write(f"[{timestamp}] ERROR [{sample_id}/{session}]: {error}\n") -def load_ingest_record(record_path: str = "./result/.ingest_record.json") -> Dict[str, Any]: +def extract_trace_id_from_error(error: str) -> str: + match = TRACE_ID_RE.search(error or "") + return match.group(1) if match else "" + + +def record_failed_session( + result: Dict[str, Any], + failed_sessions: list[Dict[str, str]], +) -> None: + trace_id = result.get("trace_id") or extract_trace_id_from_error(result.get("error", "")) + session = result.get("session") or result.get("session_key") or "" + sample_id = result.get("sample_id") or "" + display_id = result.get("display_id") or "" + error_msg = str(result.get("error", ""))[:120] + failed_sessions.append( + { + "session": str(session), + "sample_id": str(sample_id), + "display_id": str(display_id), + "trace_id": str(trace_id) if trace_id else "", + "error": error_msg, + } + ) + + +def record_success_trace_id(result: Dict[str, Any], success_trace_ids: list[str]) -> None: + trace_id = result.get("trace_id") + if trace_id: + success_trace_ids.append(str(trace_id)) + + +def load_ingest_record(record_path: str = "./result/locomo/.ingest_record.json") -> Dict[str, Any]: """Load existing ingest record file, return empty dict if not exists.""" try: with open(record_path, "r", encoding="utf-8") as f: @@ -271,7 +449,7 @@ def load_ingest_record(record_path: str = "./result/.ingest_record.json") -> Dic def save_ingest_record( - record: Dict[str, Any], record_path: str = "./result/.ingest_record.json" + record: Dict[str, Any], record_path: str = "./result/locomo/.ingest_record.json" ) -> None: """Save ingest record to file.""" with open(record_path, "w", encoding="utf-8") as f: @@ -391,11 +569,15 @@ async def viking_ingest( memory_policy = build_memory_policy(group_chat) try: - # Create session - create_res = await client.create_session(memory_policy=memory_policy) + # Create session. Safe to retry: no messages have been written yet. + create_res = await _retry_transient_http( + "create_session", + lambda: client.create_session(memory_policy=memory_policy), + ) session_id = create_res["session_id"] - # Add messages one by one with created_at + # Add messages one by one with created_at. Do not retry append-only + # message writes after a lost response; that can duplicate messages. for idx, msg in enumerate(messages): msg_created_at = None if base_datetime: @@ -411,8 +593,30 @@ async def viking_ingest( peer_id=msg.get("peer_id"), ) - # Commit - result = await client.commit_session(session_id, telemetry=True) + before_commit = await _retry_transient_http( + f"get_session_before_commit session={session_id}", + lambda: client.get_session(session_id), + ) + previous_commit_count = int(before_commit.get("commit_count") or 0) + + # Commit. If the response read fails after the server accepted the + # commit, recover by checking session commit_count and commit tasks + # before retrying the POST. This avoids a blind duplicate commit. + try: + result = await _retry_transient_http( + f"commit_session session={session_id}", + lambda: client.commit_session(session_id, telemetry=True), + attempts=2, + ) + except _TRANSIENT_HTTP_ERRORS as exc: + recovered = await _recover_commit_after_transport_error( + client, + session_id=session_id, + previous_commit_count=previous_commit_count, + ) + if recovered is None: + raise exc + result = recovered # Accept both "committed" and "accepted" as success - accepted means the session was archived if result.get("status") not in ("committed", "accepted"): @@ -425,7 +629,10 @@ async def viking_ingest( # 轮询任务状态直到完成 max_attempts = 2400 # 最多等待40分钟 for _attempt in range(max_attempts): - task = await client.get_task(task_id) + task = await _retry_transient_http( + f"get_task task={task_id}", + lambda: client.get_task(task_id), + ) status = task.get("status") if task else "unknown" if status == "completed": token_usage = _parse_token_usage(task) @@ -467,19 +674,20 @@ async def process_single_session( ) -> Dict[str, Any]: """处理单个会话的导入任务""" meta = meta or {} - ingest_record = ingest_record or {} + if ingest_record is None: + ingest_record = {} csv_id = display_id or str(sample_id) source_sample_id = str(sample_id) try: started_at = time.perf_counter() - if args.api_key: + if args.api_key and args.auth_mode == "api_key": # User API keys already pin account/user on the server side. Passing # account/user headers would be rejected in api_key auth mode. user_id = "" account = "" else: - user_id = str(sample_id) if args.separate_user_by_sample else "" - account = args.account if args.separate_user_by_sample else "" + user_id = args.user or "" + account = args.account if args.auth_mode == "trusted" else "" result = await viking_ingest( messages, args.openviking_url, @@ -498,10 +706,11 @@ async def process_single_session( cache_tokens = token_usage.get("cache", 0) reasoning_tokens = token_usage.get("reasoning", 0) llm_output_tokens = token_usage.get("llm_output", 0) - print( - f" -> [COMPLETED] [{csv_id}/{session_key}] duration={duration_seconds}s, embed={embedding_tokens}, vlm={vlm_tokens}, cache={cache_tokens}, reasoning={reasoning_tokens}, completion={llm_output_tokens}, task_id={task_id}, trace_id={trace_id}", - file=sys.stderr, - ) + if not getattr(args, "_show_progress", False): + print( + f" -> [COMPLETED] [{csv_id}/{session_key}] duration={duration_seconds}s, embed={embedding_tokens}, vlm={vlm_tokens}, cache={cache_tokens}, reasoning={reasoning_tokens}, completion={llm_output_tokens}, task_id={task_id}, trace_id={trace_id}", + file=sys.stderr, + ) # Write success record result = { @@ -525,15 +734,19 @@ async def process_single_session( # 写入成功CSV write_success_record(result, args.success_csv) - # Mark as successfully ingested - mark_ingested(csv_id, session_key, ingest_record, meta) + # Mark as successfully ingested with the canonical source sample id. + # display_id/csv_id is only a human-friendly label (for example, sample_0); + # skip checks and success CSV keys use the real sample_id (for example, conv-26). + mark_ingested(source_sample_id, session_key, ingest_record, meta) save_ingest_record(ingest_record) # Save immediately after success return result except Exception as e: - print(f" -> [ERROR] [{csv_id}/{session_key}] {e}", file=sys.stderr) - traceback.print_exc(file=sys.stderr) + error_message = f"{type(e).__name__}: {e}" if str(e) else repr(e) + if not getattr(args, "_show_progress", False): + print(f" -> [ERROR] [{csv_id}/{session_key}] {error_message}", file=sys.stderr) + traceback.print_exc(file=sys.stderr) # Write error record result = { @@ -542,7 +755,8 @@ async def process_single_session( "display_id": csv_id, "session": session_key, "status": "error", - "error": str(e), + "error": error_message, + "trace_id": extract_trace_id_from_error(error_message), } # 写入错误日志 @@ -667,7 +881,36 @@ async def run_import(args: argparse.Namespace) -> None: total_cache_tokens = 0 total_reasoning_tokens = 0 total_llm_output_tokens = 0 + failed_sessions: list[dict[str, str]] = [] + success_trace_ids: list[str] = [] tasks = [] + progress_tracker: AsyncProgressTracker | None = None + progress = None + show_progress = should_show_progress(args.no_progress) + args._show_progress = show_progress + session_semaphore = ( + asyncio.Semaphore(args.parallel_sessions) if args.parallel_sessions else None + ) + if args.parallel_sessions: + print( + f"[parallel-sessions] global concurrency={args.parallel_sessions}", + file=sys.stderr, + ) + + async def process_single_session_with_progress(**kwargs) -> Dict[str, Any]: + if progress_tracker is not None: + progress_tracker.job_started() + failed = False + try: + result = await process_single_session(**kwargs) + failed = result.get("status") == "error" + return result + except Exception: + failed = True + raise + finally: + if progress_tracker is not None: + progress_tracker.job_finished(failed=failed) if args.input.endswith(".json"): # LoCoMo JSON format @@ -697,16 +940,9 @@ async def run_import(args: argparse.Namespace) -> None: file=sys.stderr, ) - # 为每个 sample 创建独立的处理协程 - async def process_sample(item, sample_index): - nonlocal \ - success_count, \ - error_count, \ - total_embedding_tokens, \ - total_vlm_tokens, \ - total_cache_tokens, \ - total_reasoning_tokens, \ - total_llm_output_tokens + # 预先为每个 sample 构建 session 列表(两条调度路径共用) + sample_info_list: list[tuple[str, str, list[dict[str, any]]]] = [] + for sample_index, item in enumerate(samples): sample_id = item["sample_id"] display_id = f"sample_{sample_index}" @@ -724,82 +960,259 @@ async def process_sample(item, sample_index): print(f"\n=== Sample {display_id} ({sample_id}) ===", file=sys.stderr) print(f" {len(sessions)} session(s) to import", file=sys.stderr) - # 同一 sample 内串行处理所有 sessions - for sess in sessions: + sample_info_list.append((sample_id, display_id, sessions)) + + import_task_count = sum( + 1 + for sample_id, _display_id, sessions in sample_info_list + for sess in sessions + if args.force_ingest + or not is_already_ingested( + sample_id, + sess["meta"]["session_key"], + ingest_record, + success_keys, + ) + ) + if show_progress: + progress, task_id = make_three_state_progress(description="Import") + progress_tracker = AsyncProgressTracker(progress, task_id, total=import_task_count) + + if session_semaphore is not None: + # --- Round-robin 扁平调度:跨 sample 均衡分配并发槽位 --- + # 每轮从每个 sample 各取一个 session,保证所有 sample 齐头并进 + all_sessions_rr: list[tuple[str, str, dict[str, any]]] = [] + max_sessions = max((len(info[2]) for info in sample_info_list), default=0) + for round_i in range(max_sessions): + for sample_id, display_id, sessions in sample_info_list: + if round_i < len(sessions): + all_sessions_rr.append((sample_id, display_id, sessions[round_i])) + + print( + f"[parallel-sessions] global concurrency={args.parallel_sessions} " + f"(round-robin across {len(sample_info_list)} samples)", + file=sys.stderr, + ) + + async def _import_session_rr( + sample_id: str, + display_id: str, + sess: dict[str, any], + ) -> dict[str, any]: + nonlocal \ + success_count, \ + error_count, \ + total_embedding_tokens, \ + total_vlm_tokens, \ + total_cache_tokens, \ + total_reasoning_tokens, \ + total_llm_output_tokens, \ + skipped_count + meta = sess["meta"] messages = sess["messages"] session_key = meta["session_key"] - label = f"{session_key} ({meta['date_time']})" + label = f"{display_id}/{session_key} ({meta['date_time']})" - # Skip already ingested sessions unless force-ingest is enabled + # Skip already ingested sessions if not args.force_ingest and is_already_ingested( sample_id, session_key, ingest_record, success_keys ): - print( - f" [{label}] [SKIP] already imported (use --force-ingest to reprocess)", - file=sys.stderr, - ) - continue + if not show_progress: + print( + f" [{label}] [SKIP] already imported (use --force-ingest to reprocess)", + file=sys.stderr, + ) + return {"status": "skipped"} # Preview messages preview = " | ".join( [f"{msg['role']}: {msg['text'][:30]}..." for msg in messages[:3]] ) - print(f" [{label}] {preview}", file=sys.stderr) - - # 串行执行(等待完成后再处理下一个 session) - res = await process_single_session( - messages=messages, - sample_id=sample_id, - display_id=display_id, - session_key=session_key, - meta=meta, - run_time=run_time, - ingest_record=ingest_record, - args=args, - ) - if res.get("status") == "success": + if not show_progress: + print(f" [{label}] {preview}", file=sys.stderr) + if progress_tracker is not None: + progress_tracker.job_started() + + failed = False + try: + result = await process_single_session( + messages=messages, + sample_id=sample_id, + display_id=display_id, + session_key=session_key, + meta=meta, + run_time=run_time, + ingest_record=ingest_record, + args=args, + ) + failed = result.get("status") == "error" + except Exception: + failed = True + raise + finally: + if progress_tracker is not None: + progress_tracker.job_finished(failed=failed) + + if result.get("status") == "success": success_count += 1 - total_embedding_tokens += res.get("embedding_tokens", 0) - total_vlm_tokens += res.get("vlm_tokens", 0) - total_cache_tokens += res.get("cache_tokens", 0) - total_reasoning_tokens += res.get("reasoning_tokens", 0) - total_llm_output_tokens += res.get("llm_output_tokens", 0) - elif res.get("status") == "error": + record_success_trace_id(result, success_trace_ids) + total_embedding_tokens += result.get("embedding_tokens", 0) + total_vlm_tokens += result.get("vlm_tokens", 0) + total_cache_tokens += result.get("cache_tokens", 0) + total_reasoning_tokens += result.get("reasoning_tokens", 0) + total_llm_output_tokens += result.get("llm_output_tokens", 0) + elif result.get("status") == "error": error_count += 1 + record_failed_session(result, failed_sessions) + elif result.get("status") == "skipped": + skipped_count += 1 - if args.parallel_samples: - semaphore = asyncio.Semaphore(args.parallel_samples) + return result - async def process_sample_with_limit(item, sample_index): - async with semaphore: - await process_sample(item, sample_index) + async def _import_session_limited( + sample_id: str, display_id: str, sess: dict[str, any] + ) -> dict[str, any]: + async with session_semaphore: + return await _import_session_rr(sample_id, display_id, sess) + # 按 round-robin 顺序创建所有 task;semaphore 的 FIFO 队列保证执行顺序 + # 也基本是 round-robin 的,从而实现跨 sample 均衡 tasks = [ - asyncio.create_task(process_sample_with_limit(item, idx)) - for idx, item in enumerate(samples) + asyncio.create_task(_import_session_limited(sid, did, sess)) + for sid, did, sess in all_sessions_rr ] + else: - tasks = [ - asyncio.create_task(process_sample(item, idx)) for idx, item in enumerate(samples) - ] + # --- Per-sample 串行调度:每个 sample 内 session 顺序执行 --- + # sample 间并发由 --parallel-samples 控制(默认无限制) + async def process_sample(sample_id, display_id, sessions): + nonlocal \ + success_count, \ + error_count, \ + total_embedding_tokens, \ + total_vlm_tokens, \ + total_cache_tokens, \ + total_reasoning_tokens, \ + total_llm_output_tokens, \ + skipped_count + + async def import_one_session(sess): + meta = sess["meta"] + messages = sess["messages"] + session_key = meta["session_key"] + label = f"{session_key} ({meta['date_time']})" + + # Skip already ingested sessions + if not args.force_ingest and is_already_ingested( + sample_id, session_key, ingest_record, success_keys + ): + if not show_progress: + print( + f" [{label}] [SKIP] already imported (use --force-ingest to reprocess)", + file=sys.stderr, + ) + return {"status": "skipped"} + + # Preview messages + preview = " | ".join( + [f"{msg['role']}: {msg['text'][:30]}..." for msg in messages[:3]] + ) + if not show_progress: + print(f" [{label}] {preview}", file=sys.stderr) + if progress_tracker is not None: + progress_tracker.job_started() + + failed = False + try: + result = await process_single_session( + messages=messages, + sample_id=sample_id, + display_id=display_id, + session_key=session_key, + meta=meta, + run_time=run_time, + ingest_record=ingest_record, + args=args, + ) + failed = result.get("status") == "error" + return result + except Exception: + failed = True + raise + finally: + if progress_tracker is not None: + progress_tracker.job_finished(failed=failed) + + session_results = [] + for sess in sessions: + session_results.append(await import_one_session(sess)) + + for res in session_results: + if isinstance(res, Exception): + error_count += 1 + print(f" [ERROR] session task failed: {res}", file=sys.stderr) + continue + if res.get("status") == "success": + success_count += 1 + record_success_trace_id(res, success_trace_ids) + total_embedding_tokens += res.get("embedding_tokens", 0) + total_vlm_tokens += res.get("vlm_tokens", 0) + total_cache_tokens += res.get("cache_tokens", 0) + total_reasoning_tokens += res.get("reasoning_tokens", 0) + total_llm_output_tokens += res.get("llm_output_tokens", 0) + elif res.get("status") == "error": + error_count += 1 + record_failed_session(res, failed_sessions) + elif res.get("status") == "skipped": + skipped_count += 1 + + if args.parallel_samples: + semaphore = asyncio.Semaphore(args.parallel_samples) + + async def process_sample_with_limit(sample_id, display_id, sessions): + async with semaphore: + await process_sample(sample_id, display_id, sessions) + + tasks = [ + asyncio.create_task(process_sample_with_limit(sid, did, sessions)) + for sid, did, sessions in sample_info_list + ] + else: + tasks = [ + asyncio.create_task(process_sample(sid, did, sessions)) + for sid, did, sessions in sample_info_list + ] else: # Plain text format sessions = parse_test_file(args.input) print(f"Found {len(sessions)} session(s) in text file", file=sys.stderr) + text_session_keys = [f"txt-session-{idx}" for idx in range(1, len(sessions) + 1)] + import_task_count = sum( + 1 + for session_key in text_session_keys + if args.force_ingest + or not is_already_ingested("txt", session_key, ingest_record, success_keys) + ) + if show_progress: + progress, task_id = make_three_state_progress(description="Import") + progress_tracker = AsyncProgressTracker(progress, task_id, total=import_task_count) for idx, session in enumerate(sessions, start=1): - session_key = f"txt-session-{idx}" - print(f"\n=== Text Session {idx} ===", file=sys.stderr) + session_key = text_session_keys[idx - 1] + if not show_progress: + print(f"\n=== Text Session {idx} ===", file=sys.stderr) # Skip already ingested sessions unless force-ingest is enabled if not args.force_ingest and is_already_ingested( "txt", session_key, ingest_record, success_keys ): - print( - " [SKIP] already imported (use --force-ingest to reprocess)", file=sys.stderr - ) + if not show_progress: + print( + " [SKIP] already imported (use --force-ingest to reprocess)", + file=sys.stderr, + ) skipped_count += 1 continue @@ -817,11 +1230,12 @@ async def process_sample_with_limit(item, sample_index): ) preview = " | ".join([f"{msg['role']}: {msg['text'][:30]}..." for msg in messages[:3]]) - print(f" {preview}", file=sys.stderr) + if not show_progress: + print(f" {preview}", file=sys.stderr) # 创建异步任务 task = asyncio.create_task( - process_single_session( + process_single_session_with_progress( messages=messages, sample_id="txt", display_id=f"txt_{idx}", @@ -839,7 +1253,9 @@ async def process_sample_with_limit(item, sample_index): f"\n[INFO] Starting import with {len(tasks)} tasks to process", file=sys.stderr, ) - task_results = await asyncio.gather(*tasks, return_exceptions=True) + ctx = progress if show_progress and progress is not None else contextlib.nullcontext() + with ctx: + task_results = await asyncio.gather(*tasks, return_exceptions=True) # 统计纯文本路径的结果(JSON 路径已在 process_sample 内统计) if not args.input.endswith(".json"): @@ -850,11 +1266,13 @@ async def process_sample_with_limit(item, sample_index): if isinstance(r, dict): if r.get("status") == "success": success_count += 1 + record_success_trace_id(r, success_trace_ids) total_embedding_tokens += r.get("embedding_tokens", 0) total_vlm_tokens += r.get("vlm_tokens", 0) total_llm_output_tokens += r.get("llm_output_tokens", 0) elif r.get("status") == "error": error_count += 1 + record_failed_session(r, failed_sessions) # Final summary total_processed = success_count + error_count + skipped_count @@ -863,6 +1281,23 @@ async def process_sample_with_limit(item, sample_index): print(f"Successfully imported: {success_count}", file=sys.stderr) print(f"Failed: {error_count}", file=sys.stderr) print(f"Skipped (already imported): {skipped_count}", file=sys.stderr) + if success_trace_ids: + preview = success_trace_ids[:10] + suffix = " ..." if len(success_trace_ids) > 10 else "" + print( + f"Success trace IDs ({len(success_trace_ids)}): {' '.join(preview)}{suffix}", + file=sys.stderr, + ) + if failed_sessions: + print(f"\nFailed sessions ({len(failed_sessions)}):", file=sys.stderr) + for idx, s in enumerate(failed_sessions, 1): + display_id = s.get("display_id") or s.get("sample_id") or "unknown" + session = s.get("session") or "unknown" + session_label = f"{display_id}/{session}" + trace = s.get("trace_id", "") + trace_part = f", trace_id={trace}" if trace else ", trace_id=(none)" + error_part = s.get("error", "") + print(f" [{idx}] {session_label}{trace_part} — {error_part}", file=sys.stderr) print("\n=== Token usage summary ===", file=sys.stderr) print(f"Total Embedding tokens: {total_embedding_tokens}", file=sys.stderr) print(f"Total VLM tokens: {total_vlm_tokens}", file=sys.stderr) @@ -909,13 +1344,13 @@ def main(): ) parser.add_argument( "--success-csv", - default="./result/import_success.csv", - help="Path to success records CSV file (default: import_success.csv)", + default="./result/locomo/import_success.csv", + help="Path to success records CSV file (default: ./result/locomo/import_success.csv)", ) parser.add_argument( "--error-log", - default="./result/import_errors.log", - help="Path to error log file (default: import_errors.log)", + default="./result/locomo/import_errors.log", + help="Path to error log file (default: ./result/locomo/import_errors.log)", ) parser.add_argument( "--openviking-url", @@ -932,6 +1367,17 @@ def main(): default="default", help="OpenViking account identifier (default: default)", ) + parser.add_argument( + "--user", + default="default", + help="OpenViking user identifier for trusted mode when --no-separate-user-by-sample is used (default: default)", + ) + parser.add_argument( + "--auth-mode", + choices=["api_key", "trusted"], + default="api_key", + help="OpenViking server auth mode for request identity wiring (default: api_key)", + ) parser.add_argument( "--sample", type=int, @@ -961,6 +1407,15 @@ def main(): default=None, help="Max number of samples to import concurrently. Default: no limit; create one task per sample.", ) + parser.add_argument( + "--parallel-sessions", + type=int, + default=0, + help=( + "Max number of sessions to import concurrently inside each sample. " + "Default 0 means serial per sample." + ), + ) parser.add_argument( "--force-ingest", action="store_true", @@ -984,6 +1439,11 @@ def main(): default=None, help="Path to a judged result CSV. Only import sessions needed by valid wrong questions.", ) + parser.add_argument( + "--no-progress", + action="store_true", + help="Disable the live progress bar (fall back to line-by-line logs). Auto-disabled when stderr is not a TTY.", + ) args = parser.parse_args() # 确保输出目录存在 diff --git a/benchmark/locomo/vikingbot/judge.py b/benchmark/locomo/vikingbot/judge.py index 65a510fc29..d540ecdb25 100644 --- a/benchmark/locomo/vikingbot/judge.py +++ b/benchmark/locomo/vikingbot/judge.py @@ -1,12 +1,21 @@ import argparse +import asyncio import csv import json import os -import asyncio -from openai import AsyncOpenAI -from dotenv import load_dotenv +import sys +import time from pathlib import Path +from dotenv import load_dotenv +from openai import AsyncOpenAI +from progress_utils import ( + AsyncProgressTracker, + format_duration, + make_three_state_progress, + should_show_progress, +) + # 加载本地环境变量文件 env_file = Path.home() / ".openviking_benchmark_env" load_dotenv(env_file) @@ -91,8 +100,8 @@ async def main(): ) parser.add_argument( "--input", - default="./result/locomo_qa_result_only_sys_memory.csv", - help="Path to QA result csv file, default: ./result/locomo_qa_result.csv", + default="./result/locomo/locomo_qa_result_only_sys_memory.csv", + help="Path to QA result csv file, default: ./result/locomo/locomo_qa_result.csv", ) parser.add_argument( "--base-url", @@ -110,9 +119,15 @@ async def main(): help="Judge model name, default: doubao-seed-2-0-pro-260215", ) parser.add_argument( - "--parallel", type=int, default=5, help="Parallel request count, default: 5" + "--parallel", type=int, default=100, help="Parallel request count, default: 100" + ) + parser.add_argument( + "--no-progress", + action="store_true", + help="Disable the live progress bar (fall back to line-by-line logs). Auto-disabled when stderr is not a TTY.", ) args = parser.parse_args() + started_at = time.perf_counter() if not args.token: print("Error: API token is required") @@ -128,7 +143,7 @@ async def main(): total = len(rows) # 筛选未评分的行 ungraded = [i for i, row in enumerate(rows) if not row.get("result")] - print(f"Total answers: {total}, ungraded: {len(ungraded)}") + print(f"Total answers: {total}, ungraded: {len(ungraded)}", file=sys.stderr) if not ungraded: print("All answers already graded, exit") @@ -141,6 +156,15 @@ async def main(): semaphore = asyncio.Semaphore(args.parallel) file_lock = asyncio.Lock() # 用于同步文件写入 + show_progress = should_show_progress(args.no_progress) + + if show_progress: + progress, task_id = make_three_state_progress(description="Judge") + progress_tracker = AsyncProgressTracker(progress, task_id, total=len(ungraded)) + else: + progress = None + progress_tracker = None + async def save_results(): """保存当前所有结果到CSV文件,使用临时文件+原子替换避免文件损坏""" async with file_lock: @@ -153,29 +177,55 @@ async def save_results(): async def process_row(idx): async with semaphore: - row = rows[idx] - question = row["question"] - gold = row["answer"] - response = row["response"] - print(f"Grading {idx + 1}/{total}: {question[:60]}...") - is_correct, reasoning = await grade_answer(client, args.model, question, gold, response) - row["result"] = "CORRECT" if is_correct else "WRONG" - row["reasoning"] = reasoning - - # 处理完一条就立即保存结果 - await save_results() - print(f"Saved result for {idx + 1}/{total}: {row['result']}") - - return idx, row + if progress_tracker is not None: + progress_tracker.job_started() + + failed = False + try: + row = rows[idx] + question = row["question"] + gold = row["answer"] + response = row["response"] + if not show_progress: + print(f"Grading {idx + 1}/{total}: {question[:60]}...") + + is_correct, reasoning = await grade_answer( + client, args.model, question, gold, response + ) + + row["result"] = "CORRECT" if is_correct else "WRONG" + row["reasoning"] = reasoning + + # 处理完一条就立即保存结果 + await save_results() + if not show_progress: + print(f"Saved result for {idx + 1}/{total}: {row['result']}") + + return idx, row + except Exception: + failed = True + raise + finally: + if progress_tracker is not None: + progress_tracker.job_finished(failed=failed) tasks = [process_row(idx) for idx in ungraded] - await asyncio.gather(*tasks) + + if show_progress: + with progress: + await asyncio.gather(*tasks) + else: + await asyncio.gather(*tasks) # 统计结果 correct = sum(1 for row in rows if row.get("result") == "CORRECT") total_graded = sum(1 for row in rows if row.get("result")) accuracy = correct / total_graded if total_graded > 0 else 0.0 - print(f"\nGrading completed: {correct}/{total_graded} correct, accuracy: {accuracy:.2%}") + elapsed = format_duration(time.perf_counter() - started_at) + print( + f"\nGrading completed: {correct}/{total_graded} correct, " + f"accuracy: {accuracy:.2%}, elapsed: {elapsed}" + ) print(f"All results saved to {args.input}") diff --git a/benchmark/locomo/vikingbot/preflight_eval_config.py b/benchmark/locomo/vikingbot/preflight_eval_config.py index 8ef8a703f8..90c805d6f5 100644 --- a/benchmark/locomo/vikingbot/preflight_eval_config.py +++ b/benchmark/locomo/vikingbot/preflight_eval_config.py @@ -78,10 +78,11 @@ def _resolve_ov_conf_path() -> Path: def _warn_deprecated_or_conflicting_fields(ov_data: dict) -> None: ov_server = (ov_data.get("bot") or {}).get("ov_server") or {} + server_auth_mode = str((ov_data.get("server") or {}).get("auth_mode") or "").strip().lower() if str(ov_server.get("root_api_key") or "").strip(): _log_warn("bot.ov_server.root_api_key 已废弃,评测不会再把它当作认证 key 使用。") api_key_type = str(ov_server.get("api_key_type") or "").strip().lower() - if api_key_type and api_key_type != "user": + if api_key_type and api_key_type != "user" and server_auth_mode != "trusted": _log_warn("bot.ov_server.api_key_type 不是 user;后续会在 User key 校验通过后同步为 user。") @@ -99,7 +100,7 @@ def main() -> int: return 1 _warn_deprecated_or_conflicting_fields(ov_data) - _log_ok("本地配置可读取;将继续连接 OpenViking 校验 User API key。") + _log_ok("本地配置可读取;将继续连接 OpenViking 校验 API key。") return 0 except KeyboardInterrupt: _log_error("用户取消。") diff --git a/benchmark/locomo/vikingbot/preflight_eval_runtime.py b/benchmark/locomo/vikingbot/preflight_eval_runtime.py index a1fcece5de..31c255545c 100644 --- a/benchmark/locomo/vikingbot/preflight_eval_runtime.py +++ b/benchmark/locomo/vikingbot/preflight_eval_runtime.py @@ -28,7 +28,7 @@ def _error(message: str) -> None: class UserKeyValidationError(RuntimeError): - """Raised when the configured OpenViking key is not a usable User key.""" + """Raised when the configured OpenViking key is not usable for this auth mode.""" def _load_json(path: Path) -> dict: @@ -72,16 +72,28 @@ def _prompt_yes_no(prompt: str, default: bool = False) -> bool: return answer in {"y", "yes", "1", "true"} -def _resolve_configured_account_hint() -> str: +def _resolve_configured_identity_hint() -> tuple[str, str]: + account = "" + user_id = "" + path = resolve_config_path(None, OPENVIKING_CLI_CONFIG_ENV, DEFAULT_OVCLI_CONF) - if path is None: - return "default" + if path is not None: + try: + data = _load_json(Path(path)) + account = str(data.get("account") or "").strip() + user_id = str(data.get("user") or "").strip() + except Exception: + pass + try: - data = _load_json(Path(path)) + ov_data = _load_ov_conf() + ov_server = (ov_data.get("bot") or {}).get("ov_server") or {} + account = account or str(ov_server.get("account_id") or "").strip() + user_id = user_id or str(ov_server.get("admin_user_id") or "").strip() except Exception: - return "default" - account = str(data.get("account") or "").strip() - return account or "default" + pass + + return account or "default", user_id or "default" def _resolve_openviking_url() -> str: @@ -138,7 +150,7 @@ def _write_json_with_backup(path: Path, data: dict) -> None: f.write("\n") -def _sync_bot_identity(account_id: str, user_id: str, api_key: str) -> None: +def _sync_bot_identity(account_id: str, user_id: str, api_key: str, *, auth_mode: str) -> None: ov_conf_path = resolve_config_path(None, OPENVIKING_CONFIG_ENV, DEFAULT_OV_CONF) if ov_conf_path is None: return @@ -147,7 +159,7 @@ def _sync_bot_identity(account_id: str, user_id: str, api_key: str) -> None: ov_server = ov_data.setdefault("bot", {}).setdefault("ov_server", {}) changed = False desired = { - "api_key_type": "user", + "api_key_type": "root" if auth_mode == "trusted" else "user", "api_key": api_key, "account_id": account_id, "admin_user_id": user_id, @@ -160,7 +172,7 @@ def _sync_bot_identity(account_id: str, user_id: str, api_key: str) -> None: _write_json_with_backup(path, ov_data) _log( "已同步 bot.ov_server: " - f"api_key_type=user, account_id={account_id}, admin_user_id={user_id}" + f"api_key_type={desired['api_key_type']}, account_id={account_id}, admin_user_id={user_id}" ) @@ -208,12 +220,50 @@ def _resolve_root_api_key() -> tuple[str, str]: return "", "not configured" -def _parse_health_identity(body: str) -> tuple[str, str, str]: +def _request_health( + url: str, + api_key: str, + *, + account_id: str | None = None, + user_id: str | None = None, +) -> dict: + headers = { + "X-API-Key": api_key, + "Content-Type": "application/json", + } + if account_id: + headers["X-OpenViking-Account"] = account_id + if user_id: + headers["X-OpenViking-User"] = user_id + req = urllib.request.Request( + f"{url}/health", + headers=headers, + method="GET", + ) + + try: + with urllib.request.urlopen(req, timeout=10) as resp: + body = resp.read().decode("utf-8", errors="replace") + except urllib.error.HTTPError as e: + detail = _read_http_error(e) + raise UserKeyValidationError( + f"OpenViking server 检查失败(HTTP {e.code}): {detail}" + ) from e + except Exception as exc: + raise UserKeyValidationError(f"OpenViking server 不可用: {exc}") from exc + try: payload = json.loads(body) except Exception as exc: raise UserKeyValidationError(f"/health 返回非 JSON: {exc}") from exc + return payload + + +def _health_auth_mode(payload: dict) -> str: + return str(payload.get("auth_mode") or "").strip().lower() + +def _parse_health_identity_payload(payload: dict) -> tuple[str, str, str]: account_id = str(payload.get("account_id") or "").strip() user_id = str(payload.get("user_id") or "").strip() role = str(payload.get("role") or "").strip().lower() @@ -224,6 +274,14 @@ def _parse_health_identity(body: str) -> tuple[str, str, str]: return account_id, user_id, role +def _parse_health_identity(body: str) -> tuple[str, str, str]: + try: + payload = json.loads(body) + except Exception as exc: + raise UserKeyValidationError(f"/health 返回非 JSON: {exc}") from exc + return _parse_health_identity_payload(payload) + + def _read_http_error(exc: urllib.error.HTTPError) -> str: return exc.read().decode("utf-8", errors="replace") @@ -339,27 +397,8 @@ def _ensure_default_user_key(url: str, root_api_key: str) -> tuple[str, str, str def _ensure_server_and_user_key_ready( url: str, selected_account: str, api_key: str ) -> tuple[str, str]: - req = urllib.request.Request( - f"{url}/health", - headers={ - "X-API-Key": api_key, - "Content-Type": "application/json", - }, - method="GET", - ) - - try: - with urllib.request.urlopen(req, timeout=10) as resp: - body = resp.read().decode("utf-8", errors="replace") - except urllib.error.HTTPError as e: - detail = _read_http_error(e) - raise UserKeyValidationError( - f"OpenViking server 检查失败(HTTP {e.code}): {detail}" - ) from e - except Exception as exc: - raise UserKeyValidationError(f"OpenViking server 不可用: {exc}") from exc - - account_id, user_id, role = _parse_health_identity(body) + payload = _request_health(url, api_key) + account_id, user_id, role = _parse_health_identity_payload(payload) if role != "user": raise UserKeyValidationError(f"当前 API key 解析为 role={role}。评测需要普通 User key。") @@ -377,26 +416,60 @@ def _ensure_server_and_user_key_ready( return account_id, user_id +def _ensure_trusted_server_ready( + url: str, account_id: str, user_id: str, api_key: str +) -> tuple[str, str]: + payload = _request_health(url, api_key, account_id=account_id, user_id=user_id) + health_account, health_user, role = _parse_health_identity_payload(payload) + if role != "user": + raise UserKeyValidationError(f"当前 trusted 身份解析为 role={role}。评测需要普通 User 身份。") + + _log( + "OpenViking server 可用,trusted 身份: " + f"account={health_account}, user={health_user}, role={role}。" + ) + return health_account, health_user + + def _resolve_ready_user_identity( openviking_url: str, selected_account: str, + selected_user: str, api_key: str, key_source: str, -) -> tuple[str, str, str]: +) -> tuple[str, str, str, str]: if api_key: + try: + probe = _request_health(openviking_url, api_key) + except UserKeyValidationError as exc: + _error(str(exc)) + raise SystemExit(1) from exc + + auth_mode = _health_auth_mode(probe) + if auth_mode == "trusted": + _log(f"使用 {key_source} 校验 OpenViking trusted key") + try: + account, user_id = _ensure_trusted_server_ready( + openviking_url, selected_account, selected_user, api_key + ) + except UserKeyValidationError as exc: + _error(str(exc)) + raise SystemExit(1) from exc + return account, user_id, api_key, "trusted" + _log(f"使用 {key_source} 校验 OpenViking User key") try: account, user_id = _ensure_server_and_user_key_ready( openviking_url, selected_account, api_key ) - return account, user_id, api_key + return account, user_id, api_key, "api_key" except UserKeyValidationError as exc: _error(str(exc)) prompt = ( "[preflight] 当前 User key 不可用,是否使用 Root key 自动生成 default User API key" ) else: - _error("未配置 OpenViking User API key。") + _error("未配置 OpenViking API key。") prompt = "[preflight] 是否使用 Root key 自动生成 default User API key" if not _prompt_yes_no(prompt, default=False): @@ -421,17 +494,18 @@ def _resolve_ready_user_identity( _error(str(exc)) raise SystemExit(1) from exc _log("已生成可用的 default User API key。") - return checked_account or account, checked_user_id or user_id, user_key + return checked_account or account, checked_user_id or user_id, user_key, "api_key" def _write_env_file( - path: Path, account: str, openviking_url: str, api_key: str, user_id: str + path: Path, account: str, openviking_url: str, api_key: str, user_id: str, auth_mode: str ) -> None: with open(path, "w", encoding="utf-8") as f: f.write(f"ACCOUNT={shlex.quote(account)}\n") f.write(f"OPENVIKING_URL={shlex.quote(openviking_url)}\n") f.write(f"OPENVIKING_API_KEY={shlex.quote(api_key)}\n") f.write(f"OPENVIKING_USER={shlex.quote(user_id)}\n") + f.write(f"OPENVIKING_AUTH_MODE={shlex.quote(auth_mode)}\n") def main() -> int: @@ -445,17 +519,17 @@ def main() -> int: ) args = parser.parse_args() - selected_account = _resolve_configured_account_hint() + selected_account, selected_user = _resolve_configured_identity_hint() openviking_url = _resolve_openviking_url() api_key, key_source = _resolve_openviking_api_key() _log(f"本次导入使用 OpenViking URL: {openviking_url}") - account, user_id, api_key = _resolve_ready_user_identity( - openviking_url, selected_account, api_key, key_source + account, user_id, api_key, auth_mode = _resolve_ready_user_identity( + openviking_url, selected_account, selected_user, api_key, key_source ) - _sync_bot_identity(account, user_id, api_key) - _write_env_file(Path(args.output_env_file), account, openviking_url, api_key, user_id) + _sync_bot_identity(account, user_id, api_key, auth_mode=auth_mode) + _write_env_file(Path(args.output_env_file), account, openviking_url, api_key, user_id, auth_mode) return 0 diff --git a/benchmark/locomo/vikingbot/progress_utils.py b/benchmark/locomo/vikingbot/progress_utils.py new file mode 100644 index 0000000000..04debb50e0 --- /dev/null +++ b/benchmark/locomo/vikingbot/progress_utils.py @@ -0,0 +1,297 @@ +"""Shared progress bar utilities for LoCoMo benchmark scripts. + +Provides a four-state progress bar (successful / failed / running / pending) built on +top of ``rich.progress``, plus helpers for both threaded and asyncio +scenarios. +""" + +from __future__ import annotations + +import sys +import time +from typing import Optional + +from rich.console import Console +from rich.progress import ( + Progress, + ProgressColumn, + Task, + TaskID, + TextColumn, +) +from rich.table import Column +from rich.text import Text + +from openviking.session.train.components.progress import ProgressSummaryColumn + + +def format_duration(seconds: float) -> str: + """Format a duration as compact h/m/s text for progress displays.""" + total_seconds = max(0, int(round(seconds))) + hours, remainder = divmod(total_seconds, 3600) + minutes, secs = divmod(remainder, 60) + if hours: + return f"{hours}h{minutes:02d}m{secs:02d}s" + if minutes: + return f"{minutes}m{secs:02d}s" + return f"{secs}s" + + +class ElapsedTimeColumn(ProgressColumn): + """Render wall-clock elapsed time for a progress task.""" + + def render(self, task: Task) -> Text: + started_at = task.fields.get("started_at") + if started_at is None: + return Text("elapsed: 0s", style="dim") + elapsed = format_duration(time.monotonic() - float(started_at)) + return Text(f"elapsed: {elapsed}", style="dim") + + +# --------------------------------------------------------------------------- +# Multi-state bar column +# --------------------------------------------------------------------------- + + +class ThreeStateBarColumn(ProgressColumn): + """A progress bar that renders four states: successful / failed / running / pending. + + - **successful** = solid green + - **failed** = solid red + - **running** = shaded yellow + - **pending** = hollow / background + + The number of in-flight items is read from + ``task.fields.get("running", 0)``. The total bar width maps to + ``task.total``; ``task.completed`` is the processed count and + ``task.fields["failed"]`` splits failures out of that processed count. + """ + + def __init__( + self, + bar_width: int = 28, + style: str = "bar.complete", + running_style: str = "bar.finished", + complete_style: Optional[str] = None, + finished_style: Optional[str] = None, + pulse_style: str = "bar.pulse", + table_column: Optional[Column] = None, + ) -> None: + super().__init__(table_column=table_column or Column(ratio=1, no_wrap=True)) + self.bar_width = bar_width + self.style = style + # "complete" = done portion, "finished" = running portion. + # We accept both naming conventions so callers can override freely. + self.complete_style = complete_style or style + self.finished_style = finished_style or running_style + self.pulse_style = pulse_style + + def render(self, task: Task) -> Text: + """Render the four-state bar.""" + bar_width = self.bar_width or 40 + + total = max(task.total or 0, 0) + processed = max(task.completed or 0, 0) + failed = max(int(task.fields.get("failed", 0) or 0), 0) + done = max(int(task.fields.get("succeeded", processed - failed) or 0), 0) + running = max(int(task.fields.get("running", 0) or 0), 0) + + if total <= 0: + # Degenerate case: show empty bar + return Text("─" * bar_width, style="bar.back") + + # Clamp so we don't exceed 100% visually + done = min(done, total) + failed = min(failed, total - done) + running = min(running, total - done - failed) + + done_width = int(bar_width * done / total) + failed_width = int(bar_width * (done + failed) / total) - done_width + running_width = ( + int(bar_width * (done + failed + running) / total) - done_width - failed_width + ) + pending_width = bar_width - done_width - failed_width - running_width + + bar = Text() + if done_width > 0: + bar.append("█" * done_width, style="green") + if failed_width > 0: + bar.append("█" * failed_width, style="red") + if running_width > 0: + bar.append("▓" * running_width, style="yellow") + if pending_width > 0: + bar.append("░" * pending_width, style="dim") + + return bar + + +# --------------------------------------------------------------------------- +# Factory +# --------------------------------------------------------------------------- + + +def make_three_state_progress( + description: str = "Progress", + console: Optional[Console] = None, + transient: bool = False, +) -> tuple[Progress, TaskID]: + """Create a :class:`Progress` instance with a four-state bar. + + Returns the ``(progress, task_id)`` pair. The task starts with + ``completed=0``, ``total=0``, ``running=0``, ``failed=0``; callers + should call ``progress.update(task_id, total=N)`` to set the total and + use a tracker to keep success/failure/running fields in sync. + """ + console = console or Console(stderr=True, soft_wrap=False) + progress = Progress( + ThreeStateBarColumn(), + TextColumn("[progress.percentage]{task.percentage:>3.0f}%"), + ProgressSummaryColumn(), + ElapsedTimeColumn(), + console=console, + transient=transient, + ) + task_id = progress.add_task( + description, + total=0, + running=0, + succeeded=0, + failed=0, + started_at=time.monotonic(), + ) + return progress, task_id + + +def should_show_progress(no_progress_flag: bool) -> bool: + """Decide whether to render a progress bar. + + Disabled when the user explicitly passed ``--no-progress`` or when + stderr is not a TTY (e.g. redirected to a file / CI logs). + """ + if no_progress_flag: + return False + return sys.stderr.isatty() + + +# --------------------------------------------------------------------------- +# Thread-safe counter (for run_eval.py's ThreadPoolExecutor) +# --------------------------------------------------------------------------- + + +class ThreadSafeProgressTracker: + """Thin wrapper around a rich Progress that keeps ``running`` count + correct when tasks start/finish from multiple threads. + + Usage:: + + tracker = ThreadSafeProgressTracker(progress, task_id, total=N) + # For each job: + tracker.job_started() + try: + ... do work ... + finally: + tracker.job_finished() + """ + + def __init__(self, progress: Progress, task_id: TaskID, total: int) -> None: + import threading + + self._progress = progress + self._task_id = task_id + self._lock = threading.Lock() + self._running = 0 + self._done = 0 + self._failed = 0 + self._total = total + self._progress.update(task_id, total=total, completed=0, running=0, succeeded=0, failed=0) + + def job_started(self) -> None: + """Call when a worker thread picks up a job.""" + with self._lock: + self._running += 1 + self._progress.update( + self._task_id, + running=self._running, + ) + + def job_finished(self, *, failed: bool = False) -> None: + """Call when a worker thread finishes a job.""" + with self._lock: + self._running = max(0, self._running - 1) + self._done += 1 + if failed: + self._failed += 1 + self._progress.update( + self._task_id, + completed=self._done, + running=self._running, + succeeded=self._done - self._failed, + failed=self._failed, + ) + + @property + def done(self) -> int: + with self._lock: + return self._done + + @property + def running(self) -> int: + with self._lock: + return self._running + + @property + def failed(self) -> int: + with self._lock: + return self._failed + + +# --------------------------------------------------------------------------- +# Async-safe counter (for judge.py's asyncio semaphore pattern) +# --------------------------------------------------------------------------- + + +class AsyncProgressTracker: + """Same idea as :class:`ThreadSafeProgressTracker` but for asyncio. + + Since asyncio tasks all run on the same event loop, we don't strictly + need a lock; we keep one anyway for clarity and in case someone later + mixes threads in. + """ + + def __init__(self, progress: Progress, task_id: TaskID, total: int) -> None: + self._progress = progress + self._task_id = task_id + self._running = 0 + self._done = 0 + self._failed = 0 + self._total = total + self._progress.update(task_id, total=total, completed=0, running=0, succeeded=0, failed=0) + + def job_started(self) -> None: + self._running += 1 + self._progress.update(self._task_id, running=self._running) + + def job_finished(self, *, failed: bool = False) -> None: + self._running = max(0, self._running - 1) + self._done += 1 + if failed: + self._failed += 1 + self._progress.update( + self._task_id, + completed=self._done, + running=self._running, + succeeded=self._done - self._failed, + failed=self._failed, + ) + + @property + def done(self) -> int: + return self._done + + @property + def running(self) -> int: + return self._running + + @property + def failed(self) -> int: + return self._failed diff --git a/benchmark/locomo/vikingbot/run_eval.py b/benchmark/locomo/vikingbot/run_eval.py index 1fe40d3cce..1046a10b14 100644 --- a/benchmark/locomo/vikingbot/run_eval.py +++ b/benchmark/locomo/vikingbot/run_eval.py @@ -1,15 +1,23 @@ import argparse +import contextlib import csv import json import os import re import subprocess +import sys import threading import time from concurrent.futures import ThreadPoolExecutor, as_completed from datetime import datetime from pathlib import Path +from progress_utils import ( + ThreadSafeProgressTracker, + make_three_state_progress, + should_show_progress, +) + def get_evidence_text(evidence_list: list, sample: dict) -> list[str]: """根据 evidence 列表获取原始对话文本 @@ -301,7 +309,7 @@ def run_vikingbot_chat( time_cost = end_time - start_time output = result.stdout.strip() - # 解析返回的json结果,处理换行、多余前缀等特殊情况 + # 解析返回的 JSON 结果 try: resp_json = json.loads(output, strict=False) response = resp_json.get("text", "") @@ -407,8 +415,8 @@ def main(): ) parser.add_argument( "--output", - default="./result/locomo_qa_result.csv", - help="Path to output csv file, default: ./result/locomo_qa_result.csv", + default="./result/locomo/locomo_qa_result.csv", + help="Path to output csv file, default: ./result/locomo/locomo_qa_result.csv", ) parser.add_argument( "--errors", @@ -436,7 +444,7 @@ def main(): "--count", type=int, default=None, help="Number of QA questions to run, default all" ) parser.add_argument( - "--threads", type=int, default=40, help="Number of concurrent threads, default: 40" + "--threads", type=int, default=100, help="Number of concurrent threads, default: 100" ) parser.add_argument( "--update-mode", @@ -459,6 +467,11 @@ def main(): action="store_true", help="Skip questions already present in the output file", ) + parser.add_argument( + "--no-progress", + action="store_true", + help="Disable the live progress bar (fall back to line-by-line logs). Auto-disabled when stderr is not a TTY.", + ) args = parser.parse_args() # 如果指定了 question-index,自动设置 count=1 @@ -562,115 +575,142 @@ def main(): remaining_qa = [qa for qa in qa_list if qa["question"] not in processed_questions] remaining_count = len(remaining_qa) print( - f"Starting evaluation with {args.threads} concurrent threads, {remaining_count} questions to process" + f"Starting evaluation with {args.threads} concurrent threads, {remaining_count} questions to process", + file=sys.stderr, ) + show_progress = should_show_progress(args.no_progress) + + if show_progress: + progress, task_id = make_three_state_progress(description="Eval") + progress_tracker = ThreadSafeProgressTracker(progress, task_id, total=remaining_count) + else: + progress = None + progress_tracker = None + def process_qa(qa_item, idx, total_count): """单个QA处理函数,供多线程调用""" - question = qa_item["question"] - answer = qa_item["answer"] - question_time = qa_item.get("question_time") - # 使用 question_id 作为 session_id,实现完全独立并行 - question_id = qa_item.get("question_id") - speakers = qa_item.get("speakers", []) - source_sample_id = qa_item.get("original_sample_id") - sender_peer_id = source_sample_id - memory_peer_ids = None - if args.group_chat: - sender_peer_id = speakers[0] if speakers else source_sample_id - memory_peer_ids = speakers[1:] if len(speakers) > 1 else None - print(f"Processing {idx}/{total_count}: {question[:60]}...") - if question_time: - print(f" [time context: {question_time}]") - if source_sample_id: - print(f" [sample peer: {source_sample_id}]") - if speakers: - print(f" [speakers: {speakers}]") - if sender_peer_id: - print(f" [sender peer: {sender_peer_id}]") - if memory_peer_ids: - print(f" [memory peers: {memory_peer_ids}]") + if progress_tracker is not None: + progress_tracker.job_started() - response, token_usage, time_cost, iteration, tools_used_names, bot_log_file = ( - run_vikingbot_chat( - question, - question_time, - sender_peer_id, - question_id, - args.config, - memory_peer_ids, + failed = False + try: + question = qa_item["question"] + answer = qa_item["answer"] + question_time = qa_item.get("question_time") + # 使用 question_id 作为 session_id,实现完全独立并行 + question_id = qa_item.get("question_id") + speakers = qa_item.get("speakers", []) + source_sample_id = qa_item.get("original_sample_id") + sender_peer_id = source_sample_id + memory_peer_ids = None + if args.group_chat: + sender_peer_id = speakers[0] if speakers else source_sample_id + memory_peer_ids = speakers[1:] if len(speakers) > 1 else None + + if not show_progress: + print(f"Processing {idx}/{total_count}: {question[:60]}...") + if question_time: + print(f" [time context: {question_time}]") + if source_sample_id: + print(f" [sample peer: {source_sample_id}]") + if speakers: + print(f" [speakers: {speakers}]") + if sender_peer_id: + print(f" [sender peer: {sender_peer_id}]") + if memory_peer_ids: + print(f" [memory peers: {memory_peer_ids}]") + + response, token_usage, time_cost, iteration, tools_used_names, bot_log_file = ( + run_vikingbot_chat( + question, + question_time, + sender_peer_id, + question_id, + args.config, + memory_peer_ids, + ) ) - ) - row = { - "sample_id": qa_item["sample_id"], - "question_index": qa_item.get("question_index", ""), - "result": "", - "question": question, - "answer": answer, - "category": qa_item.get("category", ""), - "question_time": question_time or "", - "evidence": json.dumps(qa_item.get("evidence", [])), - "evidence_text": json.dumps(qa_item.get("evidence_text", [])), - "response": response, - "token_usage": json.dumps(token_usage, ensure_ascii=False), - "time_cost": round(time_cost, 2), - "iteration": iteration, - "tools_used_names": json.dumps(tools_used_names, ensure_ascii=False), - "bot_log_file": bot_log_file, - "is_invalid": qa_item.get("is_invalid", False), - } - - # 线程安全的结果收集 - with write_lock: - nonlocal processed_count - if args.update_mode: - if os.path.exists(args.output): - with open(args.output, "r", encoding="utf-8", newline="") as f: - reader = csv.DictReader(f) - existing_rows = list(reader) - existing_fieldnames = reader.fieldnames or fieldnames - if "bot_log_file" not in existing_fieldnames: - existing_fieldnames.append("bot_log_file") - - q_idx = str(row.get("question_index", "")) - found = False - for existing_row in existing_rows: - if str(existing_row.get("question_index", "")) == q_idx: - existing_row.update(row) - found = True - break - if not found: - existing_rows.append(row) - - with open(args.output, "w", encoding="utf-8", newline="") as f: - writer = csv.DictWriter(f, fieldnames=existing_fieldnames) - writer.writeheader() - writer.writerows(existing_rows) + row = { + "sample_id": qa_item["sample_id"], + "question_index": qa_item.get("question_index", ""), + "result": "", + "question": question, + "answer": answer, + "category": qa_item.get("category", ""), + "question_time": question_time or "", + "evidence": json.dumps(qa_item.get("evidence", [])), + "evidence_text": json.dumps(qa_item.get("evidence_text", [])), + "response": response, + "token_usage": json.dumps(token_usage, ensure_ascii=False), + "time_cost": round(time_cost, 2), + "iteration": iteration, + "tools_used_names": json.dumps(tools_used_names, ensure_ascii=False), + "bot_log_file": bot_log_file, + "is_invalid": qa_item.get("is_invalid", False), + } + + # 线程安全的结果收集 + with write_lock: + nonlocal processed_count + if args.update_mode: + if os.path.exists(args.output): + with open(args.output, "r", encoding="utf-8", newline="") as f: + reader = csv.DictReader(f) + existing_rows = list(reader) + existing_fieldnames = reader.fieldnames or fieldnames + if "bot_log_file" not in existing_fieldnames: + existing_fieldnames.append("bot_log_file") + + q_idx = str(row.get("question_index", "")) + found = False + for existing_row in existing_rows: + if str(existing_row.get("question_index", "")) == q_idx: + existing_row.update(row) + found = True + break + if not found: + existing_rows.append(row) + + with open(args.output, "w", encoding="utf-8", newline="") as f: + writer = csv.DictWriter(f, fieldnames=existing_fieldnames) + writer.writeheader() + writer.writerows(existing_rows) + else: + append_row_to_csv(args.output, fieldnames, row) else: append_row_to_csv(args.output, fieldnames, row) - else: - append_row_to_csv(args.output, fieldnames, row) - new_rows.append(row) - processed_questions.add(question) - processed_count += 1 - print(f"Completed {processed_count}/{total_count}, time cost: {round(time_cost, 2)}s") - return True + new_rows.append(row) + processed_questions.add(question) + processed_count += 1 + if not show_progress: + print(f"Completed {processed_count}/{total_count}, time cost: {round(time_cost, 2)}s") + + return True + except Exception: + failed = True + raise + finally: + if progress_tracker is not None: + progress_tracker.job_finished(failed=failed) # 使用线程池处理:全局并行,每个 question 独立 session - with ThreadPoolExecutor(max_workers=args.threads) as executor: - # 提交所有任务 - futures = [] - for idx, qa_item in enumerate(remaining_qa, 1): - futures.append(executor.submit(process_qa, qa_item, idx, remaining_count)) - - # 等待所有任务完成 - for future in as_completed(futures): - try: - future.result() - except Exception as e: - print(f"Error processing QA item: {str(e)}") + ctx = progress if show_progress else contextlib.nullcontext() + with ctx: + with ThreadPoolExecutor(max_workers=args.threads) as executor: + # 提交所有任务 + futures = [] + for idx, qa_item in enumerate(remaining_qa, 1): + futures.append(executor.submit(process_qa, qa_item, idx, remaining_count)) + + # 等待所有任务完成 + for future in as_completed(futures): + try: + future.result() + except Exception as e: + print(f"Error processing QA item: {str(e)}", file=sys.stderr) print(f"Evaluation completed, results saved to {args.output}") diff --git a/benchmark/locomo/vikingbot/run_full_eval.sh b/benchmark/locomo/vikingbot/run_full_eval.sh index 010c0350a4..0f5191c4a4 100755 --- a/benchmark/locomo/vikingbot/run_full_eval.sh +++ b/benchmark/locomo/vikingbot/run_full_eval.sh @@ -11,6 +11,7 @@ # ./run_full_eval.sh 0 2 --group-chat # 单题群聊模式 # ./run_full_eval.sh --skip-import --auto-commit # 评测全部,跳过导入,自动提交 # ./run_full_eval.sh --retry-wrong result/locomo_result_xxx.csv # 只重跑错题 +# ./run_full_eval.sh --parallel-import-sessions 20 0 1 # 覆盖默认 session 导入并发数 set -e @@ -30,6 +31,7 @@ for arg in "$@"; do echo " --no-group-chat 非群聊模式(默认),使用 sample_id 作为 Peer" echo " --auto-commit 自动提交未提交的代码变更,结果文件名带 commit id 和时间戳" echo " --retry-wrong CSV 只重跑指定结果文件中的有效错题(导入相关对话+重新问答)" + echo " --parallel-import-sessions N 单 sample 内并发导入 sessions(默认 200)" exit 0 fi done @@ -39,6 +41,7 @@ SKIP_IMPORT=false GROUP_CHAT=false AUTO_COMMIT=false RETRY_WRONG="" +PARALLEL_IMPORT_SESSIONS="200" if command -v python3 >/dev/null 2>&1; then PYTHON_BIN="python3" @@ -113,6 +116,11 @@ for arg in "$@"; do PREV_ARG="" continue fi + if [ "$PREV_ARG" = "--parallel-import-sessions" ]; then + PARALLEL_IMPORT_SESSIONS="$arg" + PREV_ARG="" + continue + fi if [ "$arg" = "--skip-import" ]; then SKIP_IMPORT=true elif [ "$arg" = "--group-chat" ]; then @@ -124,9 +132,16 @@ for arg in "$@"; do elif [ "$arg" = "--retry-wrong" ]; then PREV_ARG="$arg" continue + elif [ "$arg" = "--parallel-import-sessions" ]; then + PREV_ARG="$arg" + continue fi PREV_ARG="" done +if [ -n "$PREV_ARG" ]; then + echo "Error: $PREV_ARG requires a value" >&2 + exit 1 +fi # 过滤掉开关参数和 --retry-wrong 的值,获取位置参数 ARGS=() @@ -136,7 +151,7 @@ for arg in "$@"; do SKIP_NEXT=false continue fi - if [ "$arg" = "--retry-wrong" ]; then + if [ "$arg" = "--retry-wrong" ] || [ "$arg" = "--parallel-import-sessions" ]; then SKIP_NEXT=true continue fi @@ -154,7 +169,19 @@ else fi IMPORT_OPTS=() if [ -n "${OPENVIKING_API_KEY:-}" ]; then - IMPORT_OPTS+=("--api-key" "$OPENVIKING_API_KEY" "--no-separate-user-by-sample") + IMPORT_OPTS+=("--api-key" "$OPENVIKING_API_KEY" "--auth-mode" "${OPENVIKING_AUTH_MODE:-api_key}") + if [ "${OPENVIKING_AUTH_MODE:-api_key}" = "trusted" ]; then + IMPORT_OPTS+=("--user" "${OPENVIKING_USER:-default}") + fi + IMPORT_OPTS+=("--no-separate-user-by-sample") +fi +if [ -n "${PARALLEL_IMPORT_SESSIONS:-}" ]; then + if ! [[ "$PARALLEL_IMPORT_SESSIONS" =~ ^[1-9][0-9]*$ ]]; then + echo "Error: --parallel-import-sessions requires a positive integer" >&2 + exit 1 + fi + IMPORT_OPTS+=("--parallel-sessions" "$PARALLEL_IMPORT_SESSIONS") + echo "[import] 单 sample 内 session 并发已开启: $PARALLEL_IMPORT_SESSIONS" fi SAMPLE=${ARGS[0]} @@ -162,7 +189,7 @@ QUESTION_INDEX=${ARGS[1]} INPUT_FILE="$SCRIPT_DIR/../data/locomo10.json" # Export for inline Python usage -export SCRIPT_DIR INPUT_FILE RETRY_WRONG ACCOUNT OPENVIKING_URL OPENVIKING_API_KEY GROUP_CHAT +export SCRIPT_DIR INPUT_FILE RETRY_WRONG PARALLEL_IMPORT_SESSIONS ACCOUNT OPENVIKING_URL OPENVIKING_API_KEY OPENVIKING_USER OPENVIKING_AUTH_MODE GROUP_CHAT # auto-commit 逻辑 if [ "$AUTO_COMMIT" = "true" ]; then @@ -176,7 +203,7 @@ if [ "$AUTO_COMMIT" = "true" ]; then fi GIT_COMMIT_ID=$(git rev-parse --short HEAD) TIMESTAMP=$(date +%Y%m%d%H%M%S) -IMPORT_SUCCESS_CSV="./result/import_success.csv" +IMPORT_SUCCESS_CSV="./result/locomo/import_success.csv" IMPORT_ROW_START=0 IMPORT_PERFORMED=false @@ -316,9 +343,9 @@ if [ -n "$RETRY_WRONG" ]; then echo "源文件: $RETRY_WRONG" if [ "$AUTO_COMMIT" = "true" ]; then - RESULT_FILE="./result/locomo_retry_${TIMESTAMP}_${GIT_COMMIT_ID}.csv" + RESULT_FILE="./result/locomo/locomo_retry_${TIMESTAMP}_${GIT_COMMIT_ID}.csv" else - RESULT_FILE="./result/locomo_retry_${TIMESTAMP}.csv" + RESULT_FILE="./result/locomo/locomo_retry_${TIMESTAMP}.csv" fi # 从错题 CSV 中提取需要导入的对话(复用 import_to_ov.py 的并行逻辑) @@ -344,13 +371,13 @@ if [ -n "$RETRY_WRONG" ]; then "$INPUT_FILE" \ --output "$RESULT_FILE" \ --retry-wrong "$RETRY_WRONG" \ - --threads 20 \ + --threads 100 \ --config "$OPENVIKING_CONFIG_FILE" \ "${COMMON_OPTS[@]}" # 裁判打分 echo "[3/3] 裁判打分..." - "$PYTHON_BIN" "$SCRIPT_DIR/judge.py" --input "$RESULT_FILE" --parallel 20 + "$PYTHON_BIN" "$SCRIPT_DIR/judge.py" --input "$RESULT_FILE" --parallel 100 # 统计结果 "$PYTHON_BIN" "$SCRIPT_DIR/stat_judge_result.py" --input "$RESULT_FILE" @@ -367,9 +394,9 @@ if [ -z "$SAMPLE" ]; then echo "=== 全量评测模式 ===" if [ "$AUTO_COMMIT" = "true" ]; then - RESULT_FILE="./result/locomo_result_${TIMESTAMP}_${GIT_COMMIT_ID}.csv" + RESULT_FILE="./result/locomo/locomo_result_${TIMESTAMP}_${GIT_COMMIT_ID}.csv" else - RESULT_FILE="./result/locomo_result_${TIMESTAMP}.csv" + RESULT_FILE="./result/locomo/locomo_result_${TIMESTAMP}.csv" fi # 导入数据 @@ -391,7 +418,7 @@ if [ -z "$SAMPLE" ]; then # 裁判打分 echo "[3/4] 裁判打分..." - "$PYTHON_BIN" "$SCRIPT_DIR/judge.py" --input "$RESULT_FILE" --parallel 40 + "$PYTHON_BIN" "$SCRIPT_DIR/judge.py" --input "$RESULT_FILE" --parallel 100 # 计算结果 echo "[4/4] 计算结果..." @@ -470,9 +497,9 @@ if [ -n "$QUESTION_INDEX" ]; then echo "[2/3] Running evaluation..." fi if [ "$AUTO_COMMIT" = "true" ]; then - OUTPUT_FILE=./result/locomo_${SAMPLE}_${QUESTION_INDEX}_result_${TIMESTAMP}_${GIT_COMMIT_ID}.csv + OUTPUT_FILE=./result/locomo/locomo_${SAMPLE}_${QUESTION_INDEX}_result_${TIMESTAMP}_${GIT_COMMIT_ID}.csv else - OUTPUT_FILE=./result/locomo_${SAMPLE}_${QUESTION_INDEX}_result_${TIMESTAMP}.csv + OUTPUT_FILE=./result/locomo/locomo_${SAMPLE}_${QUESTION_INDEX}_result_${TIMESTAMP}.csv fi prepare_bot_log_dir "$OUTPUT_FILE" "$PYTHON_BIN" "$SCRIPT_DIR/run_eval.py" \ @@ -490,7 +517,7 @@ if [ -n "$QUESTION_INDEX" ]; then else echo "[3/3] Running judge..." fi - "$PYTHON_BIN" "$SCRIPT_DIR/judge.py" --input "$OUTPUT_FILE" --parallel 1 + "$PYTHON_BIN" "$SCRIPT_DIR/judge.py" --input "$OUTPUT_FILE" --parallel 100 # 输出结果 echo "" @@ -576,16 +603,16 @@ PY echo "[2/4] Running evaluation for all questions..." fi if [ "$AUTO_COMMIT" = "true" ]; then - OUTPUT_FILE=./result/locomo_${SAMPLE}_result_${TIMESTAMP}_${GIT_COMMIT_ID}.csv + OUTPUT_FILE=./result/locomo/locomo_${SAMPLE}_result_${TIMESTAMP}_${GIT_COMMIT_ID}.csv else - OUTPUT_FILE=./result/locomo_${SAMPLE}_result_${TIMESTAMP}.csv + OUTPUT_FILE=./result/locomo/locomo_${SAMPLE}_result_${TIMESTAMP}.csv fi prepare_bot_log_dir "$OUTPUT_FILE" "$PYTHON_BIN" "$SCRIPT_DIR/run_eval.py" \ "$INPUT_FILE" \ --sample "$SAMPLE_ID_FOR_CMD" \ --output "$OUTPUT_FILE" \ - --threads 10 \ + --threads 100 \ --config "$OPENVIKING_CONFIG_FILE" \ "${COMMON_OPTS[@]}" @@ -595,7 +622,7 @@ PY else echo "[3/4] Running judge..." fi - "$PYTHON_BIN" "$SCRIPT_DIR/judge.py" --input "$OUTPUT_FILE" --parallel 40 + "$PYTHON_BIN" "$SCRIPT_DIR/judge.py" --input "$OUTPUT_FILE" --parallel 100 # 输出统计结果 if [ "$SKIP_IMPORT" = "true" ]; then diff --git a/benchmark/locomo/vikingbot/stat_judge_result.py b/benchmark/locomo/vikingbot/stat_judge_result.py index 19577a4978..872b857aac 100644 --- a/benchmark/locomo/vikingbot/stat_judge_result.py +++ b/benchmark/locomo/vikingbot/stat_judge_result.py @@ -26,8 +26,8 @@ def main(): parser = argparse.ArgumentParser(description="Statistics for judge result csv") parser.add_argument( "--input", - default="./result/locomo_qa_result_only_sys_memory.csv", - help="Path to judge result csv file, default: ./result/judge_result.csv", + default="./result/locomo/locomo_qa_result_only_sys_memory.csv", + help="Path to judge result csv file, default: ./result/locomo/judge_result.csv", ) args = parser.parse_args() diff --git a/benchmark/tau2/__init__.py b/benchmark/tau2/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/benchmark/tau2/common/__init__.py b/benchmark/tau2/common/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/benchmark/tau2/common/tau2_env/__init__.py b/benchmark/tau2/common/tau2_env/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/benchmark/tau2/common/tau2_env/tau2_environment.py b/benchmark/tau2/common/tau2_env/tau2_environment.py new file mode 100644 index 0000000000..a5877a3f54 --- /dev/null +++ b/benchmark/tau2/common/tau2_env/tau2_environment.py @@ -0,0 +1,413 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import importlib +import json +import os +import time +from collections.abc import Callable +from functools import wraps +from typing import Any +from uuid import uuid4 + +from openviking.utils.model_retry import is_retryable_rate_limit_error, rate_limit_retry_delay +from openviking_cli.utils import get_logger + +logger = get_logger(__name__) + +DEFAULT_TAU2_USER_LLM = "openai/doubao-seed-2-0-code-preview-260215" + +_TAU2_GENERATE_REFERENCE_MODULES = ( + "tau2.agent.llm_agent", + "tau2.user.user_simulator", + "tau2.evaluator.evaluator_nl_assertions", + "tau2.environment.utils.interface_agent", +) + + +def _is_tau2_retryable_rate_limit_error(exc: BaseException) -> bool: + return is_retryable_rate_limit_error(exc) + + +def _tau2_rate_limit_retry_delay(attempt: int) -> float: + return rate_limit_retry_delay(attempt) + + +def _wrap_tau2_generate_with_rate_limit_retry(generate: Callable[..., Any]) -> Callable[..., Any]: + if getattr(generate, "_openviking_tau2_rate_limit_retry", False): + return generate + + @wraps(generate) + def generate_with_rate_limit_retry(*args: Any, **kwargs: Any) -> Any: + attempt = 1 + while True: + try: + return generate(*args, **kwargs) + except Exception as exc: + if not _is_tau2_retryable_rate_limit_error(exc): + raise + delay = _tau2_rate_limit_retry_delay(attempt) + logger.warning( + "tau2 LiteLLM generate rate limited; retrying attempt=%d delay=%.1fs error=%s", + attempt, + delay, + exc, + ) + time.sleep(delay) + attempt += 1 + + generate_with_rate_limit_retry._openviking_tau2_rate_limit_retry = True + generate_with_rate_limit_retry._openviking_original_generate = generate + return generate_with_rate_limit_retry + + +def _install_tau2_litellm_rate_limit_retry() -> None: + """Patch tau2-bench's sync LiteLLM generate path with rate-limit retry. + + AgentGymEnv's user simulator and orchestrator call tau2.utils.llm_utils.generate + through synchronous module globals imported with ``from ... import generate``. + Those calls run in tau2's own worker thread, so a sync sleep-based retry is + safe and does not block the OpenViking service event loop. + """ + try: + llm_utils = importlib.import_module("tau2.utils.llm_utils") + except Exception as exc: + logger.debug("tau2 llm_utils unavailable for rate-limit retry patch: %s", exc) + return + + original = getattr(llm_utils, "_openviking_original_generate", None) + current = getattr(llm_utils, "generate", None) + if not callable(current): + return + if getattr(current, "_openviking_tau2_rate_limit_retry", False): + wrapped = current + original = getattr(current, "_openviking_original_generate", original) + else: + original = current + wrapped = _wrap_tau2_generate_with_rate_limit_retry(original) + llm_utils.generate = wrapped + llm_utils._openviking_original_generate = original + + for module_name in _TAU2_GENERATE_REFERENCE_MODULES: + try: + module = importlib.import_module(module_name) + except Exception: + continue + module_generate = getattr(module, "generate", None) + if ( + module_generate is original + or module_generate is current + or getattr(module_generate, "_openviking_tau2_rate_limit_retry", False) + ): + module.generate = wrapped + + +def _install_tau2_litellm_unknown_cost_suppression() -> None: + """Suppress noisy LiteLLM cost lookup errors for private/proxy model names. + + tau2-bench logs an ERROR whenever LiteLLM cannot find a public price entry + for models such as Doubao private gateway names. Cost accounting is not used + by our rollout evaluator, and at high concurrency those repeated ERROR logs + add significant IO noise. Replace cost lookup failures with a zero-cost + result while preserving normal behavior for mapped models. + """ + + try: + llm_utils = importlib.import_module("tau2.utils.llm_utils") + except Exception as exc: + logger.debug("tau2 llm_utils unavailable for cost suppression patch: %s", exc) + return + + for name in ("get_response_cost", "get_cost"): + current = getattr(llm_utils, name, None) + if not callable(current) or getattr(current, "_openviking_tau2_cost_suppressed", False): + continue + + @wraps(current) + def suppressed_cost(*args: Any, __fn: Callable[..., Any] = current, **kwargs: Any) -> Any: + try: + return __fn(*args, **kwargs) + except Exception as exc: + text = str(exc) + if "model isn't mapped" not in text and "not mapped" not in text: + raise + logger.debug("suppressed tau2 LiteLLM cost lookup failure: %s", exc) + return 0.0 + + suppressed_cost._openviking_tau2_cost_suppressed = True + suppressed_cost._openviking_original_cost_fn = current + setattr(llm_utils, name, suppressed_cost) + + +try: + from tau2.gym.gym_agent import AgentGymEnv +except ModuleNotFoundError: + AgentGymEnv = None + + +class CommunicateWithUser: + """The agent's only channel for talking to the user. + + tau2's environment has no native "speak to the user" action, so we add this + tool: whatever ``content`` the agent passes is delivered to the tau2 user + simulator, and the simulator's reply comes back as the tool result. This class + owns both the tool's schema (``openai_schema``, consumed by + ``Tau2BenchToolProvider``) and its execution (``forward``, invoked by + ``Tau2BenchEnv.tool_call``). + """ + + name = "communicate_with_user" + description = ( + "say something to the user. Note that the customer cannot see the answer " + "returned in `final_answer`. You must communicate with the customer " + "exclusively through this tool." + ) + parameters = { + "type": "object", + "properties": { + "content": { + "type": "string", + "description": "the content to say to the user", + } + }, + "required": ["content"], + } + + def __init__(self, env): + # ``env`` is the underlying tau2 AgentGymEnv. + self.env = env + + def forward(self, content: str): + """Deliver ``content`` to the tau2 user simulator. + + Returns the raw gym step tuple ``(obs, reward, terminated, truncated, info)``; + ``Tau2BenchEnv.tool_call`` cleans the observation and tracks termination. + """ + response = self.env.tool_call(self.name, {"content": content}) + return response + + @classmethod + def openai_schema(cls) -> dict: + return { + "type": "function", + "function": { + "name": cls.name, + "description": cls.description, + "parameters": cls.parameters, + }, + } + + +class Tau2BenchEnv: + def __init__(self, domain: str, task_id: str): + if AgentGymEnv is not None: + self._impl = _GymTau2BenchEnv(domain, task_id) + else: + self._impl = _NativeTau2BenchEnv(domain, task_id) + + def reset(self): + self._impl.reset() + self.env = self._impl.env + self.terminated = self._impl.terminated + self.user_query = self._impl.user_query + self.task = self._impl.task + self.simulation_run = self._impl.simulation_run + self.policy = self._impl.policy + self.tool_schemas = self._impl.tool_schemas + self.ground_truth = self._impl.ground_truth + self.user_scenario = self._impl.user_scenario + + def tool_call(self, tool_name: str, arguments: dict) -> str: + response = self._impl.tool_call(tool_name, arguments) + self.terminated = self._impl.terminated + return response + + def append_agent_message(self, content: str) -> None: + append_message = getattr(self._impl, "append_agent_message", None) + if callable(append_message): + append_message(content) + + def _get_reward(self): + return self._impl._get_reward() + + +class _GymTau2BenchEnv: + def __init__(self, domain: str, task_id: str): + _install_tau2_litellm_rate_limit_retry() + _install_tau2_litellm_unknown_cost_suppression() + self.env = AgentGymEnv( + domain=domain, + task_id=task_id, + user_llm=os.getenv("TAU2_USER_LLM") or DEFAULT_TAU2_USER_LLM, + ) + self.terminated = False + + def reset(self): + user_query, info_dict = self.env.reset() + self.user_query = user_query.lstrip("user: ") + self.task = info_dict["task"] + self.simulation_run = info_dict["simulation_run"] + self.policy = info_dict["policy"] + self.tool_schemas = [tool.openai_schema for tool in info_dict["tools"]] + self.tool_schemas.append(CommunicateWithUser.openai_schema()) + self.ground_truth = str(self.task.evaluation_criteria) + self.user_scenario = self.task.user_scenario + + def tool_call(self, tool_name: str, arguments: dict) -> str: + if self.terminated: + return "Task Terminated" + + if tool_name == CommunicateWithUser.name: + obs, reward, terminated, truncated, info = self.env.step(arguments["content"]) + else: + action = {"name": tool_name, "arguments": arguments} + obs, reward, terminated, truncated, info = self.env.step(json.dumps(action)) + + self.terminated = terminated + return _clean_obs(obs) + + def append_agent_message(self, content: str) -> None: + if not content.strip(): + return + simulation = self._simulation_run_from_env_info() + if simulation is None: + return + + from tau2.data_model.message import AssistantMessage + + simulation.messages.append(AssistantMessage(role="assistant", content=content)) + self.env._simulation_run = simulation + self.simulation_run = simulation.model_dump_json(indent=2) + + def _get_reward(self): + reward, reward_info = self.env._get_reward() + try: + return reward, json.loads(reward_info) + except (TypeError, json.JSONDecodeError): + return reward, reward_info + + def _simulation_run_from_env_info(self): + simulation_run_json = self.env._get_info().get("simulation_run") + if not simulation_run_json: + return None + try: + from tau2.data_model.simulation import SimulationRun + + return SimulationRun.model_validate_json(simulation_run_json) + except Exception: + return None + + +class _NativeTau2BenchEnv: + def __init__(self, domain: str, task_id: str): + self.domain = domain + self.task_id = task_id + self.env = None + self.terminated = False + self.simulation_run = None + + def reset(self): + from tau2.evaluator.evaluator import EvaluationType, evaluate_simulation + from tau2.registry import registry + + self._evaluate_simulation = evaluate_simulation + self._evaluation_type = EvaluationType.ALL + self.env = registry.get_env_constructor(self.domain)() + tasks = registry.get_tasks_loader(self.domain)() + task_by_id = {str(task.id): task for task in tasks} + self.task = task_by_id[self.task_id] + self.env.set_state( + initialization_data=( + self.task.initial_state.initialization_data + if self.task.initial_state is not None + else None + ), + initialization_actions=( + self.task.initial_state.initialization_actions + if self.task.initial_state is not None + else None + ), + message_history=( + self.task.initial_state.message_history + if self.task.initial_state is not None + and self.task.initial_state.message_history is not None + else [] + ), + ) + self.policy = self.env.get_policy() + self.tool_schemas = [tool.openai_schema for tool in self.env.get_tools()] + self.tool_schemas.append(CommunicateWithUser.openai_schema()) + self.user_query = str(self.task.user_scenario) + self.ground_truth = str(self.task.evaluation_criteria) + self.user_scenario = self.task.user_scenario + self._messages = [] + + def tool_call(self, tool_name: str, arguments: dict) -> str: + from tau2.data_model.message import AssistantMessage, ToolCall + + if self.terminated: + return "Task Terminated" + + if tool_name == CommunicateWithUser.name: + # tau2 evaluates required customer-facing information by scanning + # AssistantMessage text content. Record this synthetic communication + # as assistant text so the native fallback matches gym trajectories. + self._messages.append( + AssistantMessage(role="assistant", content=str(arguments["content"])) + ) + return ( + "User simulator is unavailable in this tau2 version; " + "continue using tools and final answer." + ) + + tool_call = ToolCall( + id=f"call_{uuid4().hex}", + name=tool_name, + arguments=arguments, + requestor="assistant", + ) + assistant_message = AssistantMessage(role="assistant", tool_calls=[tool_call]) + tool_message = self.env.get_response(tool_call) + self._messages.extend([assistant_message, tool_message]) + return _clean_obs(tool_message.content or "") + + def append_agent_message(self, content: str) -> None: + from tau2.data_model.message import AssistantMessage + + if content.strip(): + self._messages.append(AssistantMessage(role="assistant", content=content)) + + def _get_reward(self): + from tau2.data_model.simulation import SimulationRun + from tau2.utils.utils import get_now + + now = get_now() + simulation = SimulationRun( + id=f"native_tau2_{self.domain}_{self.task_id}_{uuid4().hex}", + task_id=self.task.id, + start_time=now, + end_time=now, + duration=0.0, + termination_reason="agent_stop", + reward_info=None, + messages=self._messages, + ) + reward_info = self._evaluate_simulation( + domain=self.domain, + task=self.task, + simulation=simulation, + evaluation_type=self._evaluation_type, + solo_mode=False, + ) + simulation.reward_info = reward_info + self.simulation_run = simulation + return reward_info.reward, reward_info + + +def _clean_obs(obs: str) -> str: + if "tool: " in obs: + obs = obs.removeprefix("tool: ") + if "user: " in obs: + obs = obs.removeprefix("user: ") + return obs diff --git a/benchmark/tau2/vikingbot/tau2_env/tau2_tool_provider.py b/benchmark/tau2/common/tau2_env/tau2_tool_provider.py similarity index 70% rename from benchmark/tau2/vikingbot/tau2_env/tau2_tool_provider.py rename to benchmark/tau2/common/tau2_env/tau2_tool_provider.py index fec8e5a26d..f13abdc789 100644 --- a/benchmark/tau2/vikingbot/tau2_env/tau2_tool_provider.py +++ b/benchmark/tau2/common/tau2_env/tau2_tool_provider.py @@ -65,9 +65,25 @@ def load_task_id(data_split: str, task_no: int) -> tuple[str, str]: "TAU2_DATA_ROOT is not set. Point it at your tau2-bench data dir, e.g. " "export TAU2_DATA_ROOT=/data/tau2 (see setup_env.sh)." ) - split_path = os.path.join(data_root, "domains", domain, "split_tasks.json") - with open(split_path, "r", encoding="utf-8") as f: - data = json.load(f) - task_ids = data[split] + domain_dir = os.path.join(data_root, "domains", domain) + split_path = os.path.join(domain_dir, "split_tasks.json") + if os.path.exists(split_path): + with open(split_path, "r", encoding="utf-8") as f: + data = json.load(f) + task_ids = data[split] + else: + tasks_path = os.path.join(domain_dir, "tasks.json") + if not os.path.exists(tasks_path): + raise FileNotFoundError( + f"Neither split_tasks.json nor tasks.json found under: {domain_dir}" + ) + with open(tasks_path, "r", encoding="utf-8") as f: + tasks = json.load(f) + task_ids = [str(task["id"]) for task in tasks] + split_at = max(1, len(task_ids) // 2) + if split == "train": + task_ids = task_ids[:split_at] + elif split == "test": + task_ids = task_ids[split_at:] if split_at < len(task_ids) else [] task_id = task_ids[task_no] return domain, task_id diff --git a/benchmark/tau2/llm/config/baseline.yaml b/benchmark/tau2/llm/config/baseline.yaml index 97d74185d2..4941d31d10 100644 --- a/benchmark/tau2/llm/config/baseline.yaml +++ b/benchmark/tau2/llm/config/baseline.yaml @@ -40,8 +40,8 @@ eval: airline: ${TAU2_AIRLINE_FIXED_FIRST_USER_FILE:-} model: - agent_llm: ${TAU2_AGENT_LLM:-openai/doubao-seed-2-0-pro-260215} - user_llm: ${TAU2_USER_LLM:-openai/doubao-seed-2-0-pro-260215} + agent_llm: ${TAU2_AGENT_LLM:-openai/doubao-seed-2-0-code-preview-260215} + user_llm: ${TAU2_USER_LLM:-openai/doubao-seed-2-0-code-preview-260215} temperature: 0.0 openviking: diff --git a/benchmark/tau2/train/README.md b/benchmark/tau2/train/README.md new file mode 100644 index 0000000000..8fe6f44f92 --- /dev/null +++ b/benchmark/tau2/train/README.md @@ -0,0 +1,274 @@ +# Tau2 Train/Eval Pipeline + +Tau2 training/evaluation uses the generic OpenViking session/train batch +pipeline. In day-to-day runs, use +`benchmark/tau2/train/restart_vikingbot_train_eval.sh` as the main entrypoint: +it restarts the required services, points them at the same slot/config, waits for +health checks, and then launches the batch runner. + +`benchmark/tau2/train/run_batch_train_eval.sh` is only the lower-level Tau2 +wrapper. Use it when you have already started OpenViking and the Tau2 rollout +service yourself. + +## 1. Main entrypoint: restart VikingBot train/eval + +```bash +bash benchmark/tau2/train/restart_vikingbot_train_eval.sh +``` + +What the launcher does: + +1. prepares the OpenViking config/data directory for the selected slot; +2. restarts OpenViking and the VikingBot API; +3. waits for `http://127.0.0.1:/bot/v1/health`; +4. restarts the Tau2 rollout service with `--rollout-backend vikingbot`; +5. waits for `http://127.0.0.1:/health`; +6. runs `benchmark/tau2/train/run_batch_train_eval.sh` with the matching + `--config`, `--server-url`, `--benchmark-service-url`, and result directory. + +Default train/eval arguments, when no custom train/eval args are passed, are: + +```bash +--commit-concurrency 200 --epochs 2 --trials 8 --train-trials 1 --skip-final-eval +``` + +If you pass any train/eval arguments to `restart_vikingbot_train_eval.sh`, that +custom argument list replaces the launcher's default list, so include the options +you still want, such as `--skip-final-eval`. + +Example: train one task and evaluate the same train task for 8 trials after each +epoch, without pre-training baseline or extra final eval: + +```bash +bash benchmark/tau2/train/restart_vikingbot_train_eval.sh \ + --epochs 2 \ + --train-index 14 \ + --eval-split train \ + --eval-index 14 \ + --trials 8 \ + --train-trials 1 \ + --skip-baseline-eval \ + --skip-final-eval +``` + +Example: reuse the cached epoch-0/no-memory train rollout if it already exists; +on cache miss, run the rollout normally and write the cache: + +```bash +bash benchmark/tau2/train/restart_vikingbot_train_eval.sh \ + --epochs 3 \ + --train-index 5 \ + --eval-split train \ + --eval-index 5 \ + --trials 8 \ + --train-trials 4 \ + --skip-baseline-eval \ + --skip-final-eval \ + --reuse-train-rollout-cache +``` + +`--reuse-train-rollout-cache` is off by default and only affects training +rollouts for epoch `0`, before memory training has changed the policy. Later +training epochs and eval rollouts are always executed normally. + +## 2. Evaluation modes from the restart launcher + +The restart launcher always runs the full VikingBot path. Evaluation behavior is +controlled by the train/eval args passed after any launcher-only options. + +### Eval-only score + +Use `--epochs 0` to restart services and run evaluation without training: + +```bash +bash benchmark/tau2/train/restart_vikingbot_train_eval.sh \ + --epochs 0 \ + --eval-index 24 \ + --trials 8 +``` + +By default, eval uses the `test` split. Use `--eval-split train` to evaluate on +train tasks, or `--eval-split none` to disable eval. + +### Training with baseline and per-epoch eval + +The Tau2 wrapper enables `--eval-each-epoch`, so a normal training run evaluates +after every epoch using `--eval-split` and `--eval-index`. + +Before training, the runner also computes a baseline eval unless +`--skip-baseline-eval` is set. For the same dataset/domain, eval indices, trials, +and rollout options, the baseline is cached under +`result/tau2//cache/baseline/` and reused by later runs. Use +`--force-baseline-recompute` only when you intentionally want to refresh it. + +For quick train-split iteration, the common pattern is: + +```bash +bash benchmark/tau2/train/restart_vikingbot_train_eval.sh \ + --epochs 3 \ + --train-index 5 \ + --eval-split train \ + --eval-index 5 \ + --trials 8 \ + --train-trials 4 \ + --skip-baseline-eval \ + --skip-final-eval \ + --reuse-train-rollout-cache +``` + +Use `--skip-final-eval` to avoid the extra final eval pass. This is common with +Tau2 because per-epoch eval is already enabled. + +## 3. Multiple isolated slots + +The restart launcher accepts a launcher-only `--slot N` before the normal +train/eval arguments. Slot `0` is the default legacy setup. Slot `N > 0` uses +independent ports, OpenViking config/data, logs, and result directory so multiple +experiments can run at the same time: + +| Slot value | OpenViking port | VikingBot port | Tau2 service port | OpenViking root | Result directory | +|------------|-----------------|----------------|-------------------|-----------------|------------------| +| `0` | `1933` | `18790` | `1944` | `~/.openviking` | `result/tau2/train` | +| `1` | `1934` | `18791` | `1945` | `~/.openviking_1` | `result/tau2/train_1` | +| `N` | `1933 + N` | `18790 + N` | `1944 + N` | `~/.openviking_N` | `result/tau2/train_N` | + +Example: run slot 1 without touching slot 0 services or data: + +```bash +bash benchmark/tau2/train/restart_vikingbot_train_eval.sh \ + --slot 1 \ + --epochs 2 \ + --train-index 14 \ + --eval-split train \ + --eval-index 14 \ + --trials 8 \ + --train-trials 1 \ + --skip-baseline-eval \ + --skip-final-eval +``` + +Environment variables such as `OPENVIKING_PORT`, `OPENVIKING_BOT_PORT`, +`TAU2_SERVICE_PORT`, `OPENVIKING_CONFIG_FILE`, `OPENVIKING_DATA_DIR`, +`RESULT_DIR_NAME`, and `LOG_DIR` can still override the slot-derived defaults. +For non-zero slots, the launcher copies base `~/.openviking/*.conf*` config files +when needed and rewrites the slot config's `storage.workspace`, `server.port`, +`server.bot_api_url`, and `bot.ov_server.server_url`. + +The launcher writes service logs and pid files under: + +```text +result/tau2//service_logs/ +``` + +## 4. Options + +### Launcher-only options + +| Option | Default | Description | +|--------|---------|-------------| +| `--slot N` | `0` | Run an isolated experiment slot. Must appear before train/eval args. | + +### Common train/eval options + +| Option | Default | Description | +|--------|---------|-------------| +| `--domain` | `airline` | Benchmark domain to run | +| `--epochs` | `1`; restart default `2` | Number of training epochs. Use `0` for eval-only. | +| `--batch-size` | whole split | Train/eval batch size (cases per batch) | +| `--concurrency` | `200` in Tau2 wrapper | Max concurrent rollout executions | +| `--commit-concurrency` | `200` in Tau2 wrapper | Max concurrent `session.commit` submissions during training | +| `--trials` | `8` | Run each eval case N times and aggregate scores | +| `--train-trials` | `1` | Run each train case N times per epoch | +| `--train-index` | all | Run train sample(s) at 0-based split index/indices, e.g. `7` or `1,5,6` | +| `--eval-split` | `test` | Split used for baseline/per-epoch/final eval: `test`, `train`, or `none` | +| `--eval-index` | all | Run eval sample(s) at 0-based split index/indices within `--eval-split`, e.g. `14` or `1,5,6` | +| `--max-iterations` | `30` | Max steps per rollout | +| `--force-baseline-recompute` | off | Recompute cached pre-training baseline instead of reusing it | +| `--skip-baseline-eval` | off | Skip pre-training baseline eval/cache entirely | +| `--eval-each-epoch` | on in Tau2 wrapper | Run eval after every training epoch using `--eval-split` | +| `--skip-final-eval` | off; restart default on | Skip the extra final eval pass | +| `--reuse-train-rollout-cache` | off | Reuse cached epoch-0/no-memory train rollouts when present; write cache on miss | +| `--clean-result` / `--no-clean-result` | clean | Whether to prune previous result artifacts | +| `--keep-recent-results` | `5` | Number of recent default `run_` directories to keep when cleaning; cache and non-`run_` directories are preserved | +| `--output` | auto | JSON report output path | +| `--events-output` | auto | Streaming JSONL event output path | +| `--result-dir-name` | `train`; slots use `train_N` | Result subdirectory under `result//` | +| `--benchmark-service-url` | set by restart launcher | Benchmark runtime service URL | +| `--config` | set by restart launcher | ov.conf path | +| `--server-url` | set by restart launcher | OpenViking server URL | +| `--api-key` | from config | OpenViking API key | +| `--account-id` | `default` | OpenViking trusted account id | +| `--user-id` | `default` | OpenViking trusted user id | + +## 5. Manual service mode + +Use this only when you want to manage services yourself instead of using +`restart_vikingbot_train_eval.sh`. + +Start the Tau2 service manually: + +```bash +bash benchmark/tau2/train/run_service.sh --host 127.0.0.1 --port 1944 +``` + +Service options: + +| Option | Default | Description | +|--------|---------|-------------| +| `--host` | `127.0.0.1` | Service listen address | +| `--port` | `1944` | Service listen port | +| `--data-root` | auto-detect / `$TAU2_DATA_ROOT` | Path to `tau2-bench/data/tau2` | +| `--config` | `~/.openviking/ov.conf` | ov.conf for VikingBot / OpenViking access | +| `--rollout-language` | `default` | Rollout response language. Use `zh` for Chinese user-facing replies. | +| `--rollout-backend` | `vikingbot` | Rollout implementation backend. `native` for fast Python executor, `vikingbot` for full VikingBot AgentLoop. | +| `--native-thread-workers` | `128` | Thread pool size for native rollout executor. | +| `--rollout-thread-workers` | `200` | Worker threads used to host rollout executions off the uvicorn event loop. Use `0` to disable threaded hosting. | +| `--max-rollout-concurrency` | `200` | Maximum concurrent rollout executions accepted by the service. | +| `--no-kill-existing` | off | Don't kill existing process on the same port. | + +Then run the lower-level Tau2 wrapper: + +```bash +bash benchmark/tau2/train/run_batch_train_eval.sh \ + --epochs 4 \ + --trials 8 +``` + +The wrapper expands to the generic runner with Tau2 defaults: + +```bash +bash openviking/session/train/run_batch_train_eval.sh \ + --dataset tau2 \ + --domain airline \ + --eval-each-epoch \ + --concurrency 200 \ + --commit-concurrency 200 \ + --benchmark-service-url http://127.0.0.1:1944 +``` + +The batch runner does **not** send a backend choice — it always uses whatever the +Tau2 service is configured with. + +## 6. Result and rollout artifacts + +By default each run writes artifacts under the repository-level result directory: + +```text +result/tau2//run__/ + report.json + rollouts_index.json + rollouts/ +``` + +`result/tau2//latest_rollouts` points to the most recent +rollouts directory. Each rollout artifact group is one original task; each +rollout has its own subdirectory with `memory_context.md`, `messages.json`, +`tool_calls.json`, `evaluation.json`, and `commit_messages.json`. These files, +plus `rollouts_index.json`, are written as soon as each remote rollout finishes. +Train rollouts are enriched later with `commit_result.json` and +`memory_diff.json` as commit progress becomes available. + +Streaming JSONL events are written to +`result/tau2//run__/events.jsonl`; train +commit events include `trace_id` for live `tail -f` debugging. Use +`--events-output` to override the path. diff --git a/benchmark/tau2/train/__init__.py b/benchmark/tau2/train/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/benchmark/tau2/train/_rollout_helpers.py b/benchmark/tau2/train/_rollout_helpers.py new file mode 100644 index 0000000000..b655ff27f5 --- /dev/null +++ b/benchmark/tau2/train/_rollout_helpers.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Shared private helpers for the Tau2 vikingbot and native rollout executors. + +These are intentionally underscore-prefixed: they remain an internal surface +between the two executor implementations and the tests that reach through +``rollout_executor.py`` re-exports. Do not import them from outside the +``benchmark.tau2.train`` package. +""" + +from __future__ import annotations + +import json +from typing import Any + +from fastapi.encoders import jsonable_encoder + +from openviking.message import Message, TextPart +from openviking.session.train import Case, CriterionResult, RubricEvaluation + + +def _message( + message_id: str, + role: str, + text: str, + *, + created_at: str | None = None, +) -> Message: + return Message(id=message_id, role=role, parts=[TextPart(text=text)], created_at=created_at) + + +def _metadata_message( + message_id: str, + text: str, + *, + created_at: str | None = None, +) -> Message: + return _message(message_id, "user", text, created_at=created_at) + + +def _is_communicate_with_user(tool_name: str) -> bool: + return tool_name == "communicate_with_user" + + +def _communicate_text_from_tool_input(tool_input: dict[str, Any] | None) -> str: + if not isinstance(tool_input, dict): + return "" + content = tool_input.get("content") + if content is None: + return "" + return str(content) + + +def _case_trial(case: Case) -> Any: + return case.input.get("eval_trial", case.input.get("train_trial")) + + +def _tau2_evaluation(*, reward: Any, evaluation_result: Any, source: str = "tau2") -> RubricEvaluation: + score = _safe_float(reward, default=0.0) + passed = score >= 1.0 + feedback = [] if passed else ["tau2 environment reward is below 1.0."] + evaluation_jsonable = _to_jsonable(evaluation_result) + if evaluation_jsonable is not None: + feedback.append(_stringify(evaluation_jsonable)) + return RubricEvaluation( + passed=passed, + score=score, + criterion_results=[ + CriterionResult( + criterion_name="tau2_reward", + passed=passed, + score=score, + feedback=feedback, + evidence=[_stringify(evaluation_jsonable)] if evaluation_jsonable is not None else [], + metadata={"reward": score}, + ) + ], + feedback=feedback, + metadata={ + "source": source, + "reward": score, + "evaluation_result": evaluation_jsonable, + }, + ) + + +def _as_tool_input(args: Any) -> dict[str, Any]: + if isinstance(args, dict): + return args + if isinstance(args, str): + try: + parsed = json.loads(args) + except json.JSONDecodeError: + return {"arguments": args} + if isinstance(parsed, dict): + return parsed + return {"arguments": parsed} + return {"arguments": args} + + +def _safe_float(value: Any, *, default: float) -> float: + try: + return float(value) + except (TypeError, ValueError): + return default + + +def _to_jsonable(value: Any) -> Any: + return jsonable_encoder(value) + + +def _stringify(value: Any) -> str: + if isinstance(value, str): + return value + return json.dumps(_to_jsonable(value), ensure_ascii=False, sort_keys=True) diff --git a/benchmark/tau2/train/case_loader.py b/benchmark/tau2/train/case_loader.py new file mode 100644 index 0000000000..5d827755fe --- /dev/null +++ b/benchmark/tau2/train/case_loader.py @@ -0,0 +1,143 @@ +#!/usr/bin/env python3 +"""Tau2 task CaseLoader for OpenViking batch policy training.""" + +from __future__ import annotations + +import json +import os +from collections.abc import AsyncIterator +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from openviking.session.train import Case, Rubric, RubricCriterion + + +def _load_tau2_task(domain: str, task_id: str): + from tau2.registry import registry + + tasks = registry.get_tasks_loader(domain)() + task_by_id = {str(task.id): task for task in tasks} + try: + return task_by_id[task_id] + except KeyError as exc: + raise ValueError(f"tau2 task not found domain={domain} task_id={task_id}") from exc + + +@dataclass(slots=True) +class Tau2CaseLoader: + """Load tau2 split tasks as train-domain Cases.""" + + domain: str + split: str + batch_size: int | None = None + data_root: str | None = None + task_indices: list[int] | None = None + + async def batches(self, context: Any = None) -> AsyncIterator[list[Case]]: + del context + task_ids = self.load_task_ids() + size = self.batch_size or 1 + if size <= 0: + raise ValueError("batch_size must be > 0") + for start in range(0, len(task_ids), size): + yield [ + self._case_from_task(task_no, task_id) + for task_no, task_id in task_ids[start : start + size] + ] + + def load_cases(self) -> list[Case]: + task_ids = self.load_task_ids() + return [self._case_from_task(task_no, task_id) for task_no, task_id in task_ids] + + def load_task_ids(self) -> list[tuple[int, str]]: + data = _load_split_tasks(self.domain, self.data_root) + values = data.get(self.split) + if not isinstance(values, list): + return [] + task_ids = [(task_no, str(item)) for task_no, item in enumerate(values)] + if self.task_indices is None: + return task_ids + selected: list[tuple[int, str]] = [] + for index in self.task_indices: + if index < 0: + raise ValueError("task_indices must be >= 0") + try: + selected.append(task_ids[index]) + except IndexError as exc: + raise ValueError( + f"task index out of range for split {self.split!r}: {index} " + f"(size={len(task_ids)})" + ) from exc + return selected + + def split_exists(self) -> bool: + data = _load_split_tasks(self.domain, self.data_root) + values = data.get(self.split) + return isinstance(values, list) and bool(values) + + def _case_from_task(self, task_no: int, task_id: str) -> Case: + task = _load_tau2_task(self.domain, task_id) + policy = "" + user_query = str(task.user_scenario) + data_split = f"{self.domain}_{self.split}" + return Case( + name=f"tau2_{data_split}_{task_no}", + task_signature=f"tau2:{self.domain}:{self.split}:{task_id}", + input={ + "domain": self.domain, + "split": self.split, + "data_split": data_split, + "task_no": task_no, + "task_id": task_id, + "data_root": self.data_root, + "user_query": user_query, + "policy": policy, + }, + rubric=Rubric( + name=f"tau2_{data_split}_{task_no}_rubric", + description="Tau2 task reward must reach 1.0.", + criteria=[ + RubricCriterion( + name="tau2_reward", + description="The tau2 environment reward is 1.0.", + required=True, + weight=1.0, + ) + ], + ), + metadata={"source": "tau2", "domain": self.domain, "split": self.split}, + ) + + +def _load_split_tasks(domain: str, data_root: str | None = None) -> dict[str, Any]: + root = data_root or os.getenv("TAU2_DATA_ROOT") + if not root: + raise RuntimeError( + "TAU2_DATA_ROOT is not set. Point it at your tau2-bench data dir, e.g. " + "export TAU2_DATA_ROOT=/data/tau2." + ) + domain_dir = Path(root).expanduser() / "domains" / domain + split_path = domain_dir / "split_tasks.json" + if split_path.exists(): + return json.loads(split_path.read_text(encoding="utf-8")) + + tasks_path = domain_dir / "tasks.json" + if not tasks_path.exists(): + raise FileNotFoundError( + f"Neither split_tasks.json nor tasks.json found under: {domain_dir}" + ) + return _derive_split_tasks_from_tasks_json(tasks_path) + + +def _derive_split_tasks_from_tasks_json(tasks_path: Path) -> dict[str, list[str]]: + tasks = json.loads(tasks_path.read_text(encoding="utf-8")) + if not isinstance(tasks, list): + raise ValueError(f"tasks.json must be a list: {tasks_path}") + task_ids = [str(task.get("id")) for task in tasks if isinstance(task, dict) and "id" in task] + if not task_ids: + raise ValueError(f"tasks.json contains no task ids: {tasks_path}") + split_at = max(1, len(task_ids) // 2) + if split_at >= len(task_ids): + return {"train": task_ids, "test": []} + return {"train": task_ids[:split_at], "test": task_ids[split_at:]} diff --git a/benchmark/tau2/train/experience_loader_template/SKILL.md b/benchmark/tau2/train/experience_loader_template/SKILL.md new file mode 100644 index 0000000000..1f96eda43a --- /dev/null +++ b/benchmark/tau2/train/experience_loader_template/SKILL.md @@ -0,0 +1,35 @@ +--- +name: experience_loader +description: Load relevant OpenViking experience memories via case-linked experience candidates before solving a task. +--- + +# experience_loader + +Use this skill before taking task actions when reusable execution experience may help. + +## Required workflow + +1. Before taking task actions, call `search_experience` with a natural-language query that describes the current task. +2. Build the query from the current domain, user intent, target object, requested operation, policy keywords, and likely tool/action family. Avoid vague queries such as "help user". +3. Review the returned candidates. Each candidate is a matched case plus linked experience entries; each experience entry includes its `name`, `uri`, and a short `situation` snippet describing its applicability and exclusions. +4. **Gate before reading.** For each linked experience, read its `situation` snippet and check whether the current task matches the experience's applicability AND does NOT match any of its exclusions / "不适用于" / "does not apply to" items. Skip experiences whose situation explicitly excludes your case (e.g. wrong cabin class, flights already flown, different action family, or different change type). Only call `read_experience` on experiences that plausibly apply after this check. If no experience passes the gate, continue without experience guidance. +5. You may call `search_experience` multiple times with refined keywords, and you may call `read_experience` multiple times for the experiences that pass the gate. +6. Treat loaded experiences as reusable guidance, not as current-task truth. Current policy, current tool results, and current user facts override prior experience. +7. **Re-verify after reading.** Even after `read_experience`, before acting on the experience, check its full `## Situation` against current facts you have obtained from tools (cabin class, reservation status, flight dates, segment state, etc.). If any "不适用于" / exclusion condition matches the current task now that you have concrete facts, DISCARD the experience and proceed from policy and tool results instead — do NOT apply its Approach or Reflect. +8. Multi-intent tasks (e.g. "cancel, then book", "upgrade then change flight", "refuse a modification then offer a fallback") may legitimately require more than one experience; gate and apply each segment's experience independently. Do not end the task (`done` / `transfer_to_human_agents`) just because one segment's experience reaches a local return marker — check whether the user has a remaining intent. +9. If no linked experience is plausibly relevant after gating, continue without experience guidance. + +## Local return markers in loaded experiences + +Experience return markers are **local to the covered intent/subtask**. They are not whole-task success/failure labels and are not automatic permission to call `done`. + +- `RETURN_COMPLETED`: the specific intent/subtask covered by this experience has been completed, usually after the required business read/write tool calls and required customer communication. If the user has another independent intent, continue with that next intent instead of ending the conversation. +- `RETURN_BLOCKED(reason="...")`: the covered intent/subtask cannot proceed under the current facts, policy, missing input, refusal boundary, or escalation boundary. Perform any required communication/escalation from the experience, then continue other remaining user intents if they are still actionable. +- `RETURN_NOT_APPLICABLE`: the experience does not match the current facts; discard it and use another applicable experience or current policy/tool facts. + +Refusal, no-option, policy-ineligible, missing-input, and `transfer_to_human_agents` branches should be interpreted as `RETURN_BLOCKED(...)` for that local intent, not as whole-task completion. Before ending globally, verify that every user intent is completed, blocked, not applicable, or explicitly transferred/stopped by the user/environment. + +## Tools + +- `search_experience(query, limit=10)`: searches OpenViking `memories/cases` under the current user, reads each matched case's `## Linked Experiences` section, and returns JSON candidates with case score, case URI, task signature, input summary, and linked experience entries (each with `name`, `uri`, and a `situation` snippet from the experience's `## Situation` section). +- `read_experience(experience_uri)`: reads one OpenViking experience memory by full URI and returns Markdown. diff --git a/benchmark/tau2/train/restart_vikingbot_train_eval.sh b/benchmark/tau2/train/restart_vikingbot_train_eval.sh new file mode 100755 index 0000000000..84aefabfd9 --- /dev/null +++ b/benchmark/tau2/train/restart_vikingbot_train_eval.sh @@ -0,0 +1,371 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Restart the OpenViking bot server and tau2 rollout service, wait until both +# are healthy, then start tau2 vikingbot batch train/eval. +# +# Default training args match the common vikingbot run: +# --commit-concurrency 200 --epochs 2 --trials 8 --train-trials 1 --skip-final-eval +# Pass any non-launcher arguments to override/extend the batch train/eval invocation. +# +# Launcher-only options: +# --slot N Run an isolated slot. Slot 0 is the default legacy setup. Slot N>0 +# uses separate ports, OpenViking config/data, logs, and result dir. + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +TAU2_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)" +REPO_ROOT="$(cd "${TAU2_DIR}/../.." && pwd)" + +SLOT="${TAU2_TRAIN_SLOT:-0}" +declare -a TRAIN_CLI_ARGS=() + +usage() { + cat <<'USAGE' +Usage: + bash benchmark/tau2/train/restart_vikingbot_train_eval.sh [--slot N] [train/eval args...] + +Launcher options: + --slot N Isolated experiment slot. Slot 0 is default/legacy. Slot N>0 uses: + OV port = 1933 + N + OV bot port = 18790 + N + tau2 port = 1944 + N + OV config = ~/.openviking_N/ov.conf + OV data = ~/.openviking_N/data + result dir = result/tau2/train_N + +All remaining args are passed to benchmark/tau2/train/run_batch_train_eval.sh. +USAGE +} + +parse_launcher_args() { + while [[ $# -gt 0 ]]; do + case "$1" in + --slot) + if [[ $# -lt 2 ]]; then + echo "[restart-vikingbot-train] ERROR: --slot requires a value" >&2 + exit 1 + fi + SLOT="$2" + shift 2 + ;; + --slot=*) + SLOT="${1#--slot=}" + shift 1 + ;; + -h|--help) + usage + exit 0 + ;; + --) + shift + TRAIN_CLI_ARGS+=("$@") + break + ;; + *) + TRAIN_CLI_ARGS+=("$1") + shift 1 + ;; + esac + done +} + +validate_slot() { + if ! [[ "${SLOT}" =~ ^[0-9]+$ ]]; then + echo "[restart-vikingbot-train] ERROR: --slot must be a non-negative integer, got: ${SLOT}" >&2 + exit 1 + fi +} + +parse_launcher_args "$@" +validate_slot + +if [[ "${SLOT}" == "0" ]]; then + DEFAULT_OPENVIKING_PORT="1933" + DEFAULT_OPENVIKING_BOT_PORT="18790" + DEFAULT_TAU2_SERVICE_PORT="1944" + DEFAULT_RESULT_DIR_NAME="train" + DEFAULT_LOG_DIR="${REPO_ROOT}/result/tau2/train/service_logs" + DEFAULT_OPENVIKING_CONFIG_FILE="${HOME}/.openviking/ov.conf" + DEFAULT_OPENVIKING_DATA_DIR="${HOME}/.openviking/data" + DEFAULT_SLOT_ROOT="${HOME}/.openviking" +else + DEFAULT_OPENVIKING_PORT="$((1933 + SLOT))" + DEFAULT_OPENVIKING_BOT_PORT="$((18790 + SLOT))" + DEFAULT_TAU2_SERVICE_PORT="$((1944 + SLOT))" + DEFAULT_RESULT_DIR_NAME="train_${SLOT}" + DEFAULT_LOG_DIR="${REPO_ROOT}/result/tau2/${DEFAULT_RESULT_DIR_NAME}/service_logs" + DEFAULT_SLOT_ROOT="${HOME}/.openviking_${SLOT}" + DEFAULT_OPENVIKING_CONFIG_FILE="${DEFAULT_SLOT_ROOT}/ov.conf" + DEFAULT_OPENVIKING_DATA_DIR="${DEFAULT_SLOT_ROOT}/data" +fi + +OPENVIKING_PORT="${OPENVIKING_PORT:-${DEFAULT_OPENVIKING_PORT}}" +OPENVIKING_BOT_PORT="${OPENVIKING_BOT_PORT:-${DEFAULT_OPENVIKING_BOT_PORT}}" +TAU2_SERVICE_HOST="${TAU2_SERVICE_HOST:-127.0.0.1}" +TAU2_SERVICE_PORT="${TAU2_SERVICE_PORT:-${DEFAULT_TAU2_SERVICE_PORT}}" +TAU2_ROLLOUT_BACKEND="${TAU2_ROLLOUT_BACKEND:-vikingbot}" +TAU2_MAX_ROLLOUT_CONCURRENCY="${TAU2_MAX_ROLLOUT_CONCURRENCY:-150}" +TAU2_ROLLOUT_THREAD_WORKERS="${TAU2_ROLLOUT_THREAD_WORKERS:-${TAU2_MAX_ROLLOUT_CONCURRENCY}}" +WAIT_TIMEOUT_SECONDS="${WAIT_TIMEOUT_SECONDS:-180}" +RESULT_DIR_NAME="${RESULT_DIR_NAME:-${DEFAULT_RESULT_DIR_NAME}}" +LOG_DIR="${LOG_DIR:-${DEFAULT_LOG_DIR}}" +OPENVIKING_CONFIG_FILE="${OPENVIKING_CONFIG_FILE:-${DEFAULT_OPENVIKING_CONFIG_FILE}}" +OPENVIKING_DATA_DIR="${OPENVIKING_DATA_DIR:-${DEFAULT_OPENVIKING_DATA_DIR}}" +SLOT_ROOT="${SLOT_ROOT:-${DEFAULT_SLOT_ROOT}}" + +OPENVIKING_LOG="${LOG_DIR}/openviking-server.log" +TAU2_SERVICE_LOG="${LOG_DIR}/tau2-service.log" + +mkdir -p "${LOG_DIR}" + +log() { + printf '[restart-vikingbot-train] %s\n' "$*" +} + +fail() { + printf '[restart-vikingbot-train] ERROR: %s\n' "$*" >&2 + exit 1 +} + +stop_existing_listener() { + local name="$1" + local port="$2" + local pids + pids="$(lsof -tiTCP:"${port}" -sTCP:LISTEN 2>/dev/null || true)" + if [[ -z "${pids}" ]]; then + log "no existing ${name} listener on port ${port}" + return 0 + fi + + log "stopping existing ${name} listener(s) on port ${port}: ${pids}" + kill ${pids} 2>/dev/null || true + for _ in {1..20}; do + sleep 0.2 + if ! lsof -tiTCP:"${port}" -sTCP:LISTEN >/dev/null 2>&1; then + log "✓ stopped existing ${name} listener(s) on port ${port}" + return 0 + fi + done + + pids="$(lsof -tiTCP:"${port}" -sTCP:LISTEN 2>/dev/null || true)" + if [[ -n "${pids}" ]]; then + log "force stopping existing ${name} listener(s) on port ${port}: ${pids}" + kill -9 ${pids} 2>/dev/null || true + fi +} + +json_string_escape() { + local value="$1" + value="${value//\\/\\\\}" + value="${value//\"/\\\"}" + value="${value//$'\n'/\\n}" + value="${value//$'\r'/\\r}" + value="${value//$'\t'/\\t}" + printf '%s' "${value}" +} + +prepare_slot_config() { + if [[ "${SLOT}" == "0" && "${OPENVIKING_CONFIG_FILE}" == "${HOME}/.openviking/ov.conf" ]]; then + return 0 + fi + + local escaped_workspace + local config_dir + escaped_workspace="$(json_string_escape "${OPENVIKING_DATA_DIR}")" + config_dir="$(dirname "${OPENVIKING_CONFIG_FILE}")" + mkdir -p "${config_dir}" "${OPENVIKING_DATA_DIR}" + + if [[ -d "${HOME}/.openviking" && "${config_dir}" != "${HOME}/.openviking" ]]; then + local config_name + for config_name in ov.conf ovcli.conf ovcli.settings.conf; do + if [[ -f "${HOME}/.openviking/${config_name}" ]]; then + cp -f "${HOME}/.openviking/${config_name}" "${config_dir}/${config_name}" + fi + done + fi + + if [[ ! -f "${OPENVIKING_CONFIG_FILE}" ]]; then + cat > "${OPENVIKING_CONFIG_FILE}" < experience links. Diagnostic trajectories are +# not injected into runtime recall. +ov_server["case_recall_limit"] = max(int(ov_server.get("case_recall_limit", 0) or 0), 3) +ov_server["trajectory_recall_limit"] = 0 +ov_server["exp_recall_limit"] = max(int(ov_server.get("exp_recall_limit", 0) or 0), 6) +ov_server["exp_recall_max_chars"] = max( + int(ov_server.get("exp_recall_max_chars", 0) or 0), 14000 +) +config_path.parent.mkdir(parents=True, exist_ok=True) +config_path.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") +PY +} + +wait_for_http_json_ok() { + local name="$1" + local url="$2" + local required_pattern="$3" + local log_file="$4" + local deadline=$((SECONDS + WAIT_TIMEOUT_SECONDS)) + local response="" + + log "waiting for ${name}: ${url}" + while (( SECONDS < deadline )); do + response="$(curl -fsS "${url}" 2>/dev/null || true)" + if [[ -n "${response}" && "${response//[[:space:]]/}" == *"${required_pattern}"* ]]; then + log "✓ ${name} is ready" + return 0 + fi + sleep 2 + done + + log "last ${name} response: ${response:-}" + if [[ -f "${log_file}" ]]; then + log "recent ${name} logs:" + tail -80 "${log_file}" >&2 || true + fi + fail "${name} did not become ready within ${WAIT_TIMEOUT_SECONDS}s" +} + +start_openviking_server() { + prepare_slot_config + log "slot=${SLOT} result_dir=result/tau2/${RESULT_DIR_NAME}" + log "slot root: ${SLOT_ROOT}" + log "OpenViking config: ${OPENVIKING_CONFIG_FILE}" + log "OpenViking data: ${OPENVIKING_DATA_DIR}" + log "restarting OpenViking server on port ${OPENVIKING_PORT}, bot port ${OPENVIKING_BOT_PORT}" + log "OpenViking log: ${OPENVIKING_LOG}" + : > "${OPENVIKING_LOG}" + stop_existing_listener "OpenViking server" "${OPENVIKING_PORT}" + stop_existing_listener "OpenViking bot" "${OPENVIKING_BOT_PORT}" + + ( + cd "${REPO_ROOT}" + export OPENVIKING_CONFIG_FILE + exec bot/scripts/restart_openviking_server.sh \ + --port "${OPENVIKING_PORT}" \ + --bot-port "${OPENVIKING_BOT_PORT}" \ + --config "${OPENVIKING_CONFIG_FILE}" \ + --data-dir "${OPENVIKING_DATA_DIR}" \ + --no-kill-all-vikingbot + ) >"${OPENVIKING_LOG}" 2>&1 & + + echo "$!" > "${LOG_DIR}/openviking-server.pid" + log "OpenViking restart wrapper pid: $(cat "${LOG_DIR}/openviking-server.pid")" + + wait_for_http_json_ok \ + "OpenViking bot API" \ + "http://127.0.0.1:${OPENVIKING_PORT}/bot/v1/health" \ + '"status":"healthy"' \ + "${OPENVIKING_LOG}" +} + +start_tau2_service() { + log "restarting tau2 service on ${TAU2_SERVICE_HOST}:${TAU2_SERVICE_PORT} backend=${TAU2_ROLLOUT_BACKEND}" + log "tau2 service concurrency=${TAU2_MAX_ROLLOUT_CONCURRENCY} rollout_thread_workers=${TAU2_ROLLOUT_THREAD_WORKERS}" + log "tau2 service log: ${TAU2_SERVICE_LOG}" + : > "${TAU2_SERVICE_LOG}" + stop_existing_listener "tau2 rollout service" "${TAU2_SERVICE_PORT}" + + ( + cd "${REPO_ROOT}" + export OPENVIKING_CONFIG_FILE + exec benchmark/tau2/train/run_service.sh \ + --host "${TAU2_SERVICE_HOST}" \ + --port "${TAU2_SERVICE_PORT}" \ + --config "${OPENVIKING_CONFIG_FILE}" \ + --rollout-backend "${TAU2_ROLLOUT_BACKEND}" \ + --max-rollout-concurrency "${TAU2_MAX_ROLLOUT_CONCURRENCY}" \ + --rollout-thread-workers "${TAU2_ROLLOUT_THREAD_WORKERS}" + ) >"${TAU2_SERVICE_LOG}" 2>&1 & + + echo "$!" > "${LOG_DIR}/tau2-service.pid" + log "tau2 service pid: $(cat "${LOG_DIR}/tau2-service.pid")" + + wait_for_http_json_ok \ + "tau2 rollout service" \ + "http://${TAU2_SERVICE_HOST}:${TAU2_SERVICE_PORT}/health" \ + '"status":"ok"' \ + "${TAU2_SERVICE_LOG}" +} + +run_train_eval() { + local -a train_args=("$@") + if [[ ${#train_args[@]} -eq 0 ]]; then + train_args=( + --commit-concurrency 200 + --epochs 2 + --trials 8 + --train-trials 1 + --skip-final-eval + ) + fi + + export OPENVIKING_CONFIG_FILE + export BENCHMARK_SERVICE_URL="http://${TAU2_SERVICE_HOST}:${TAU2_SERVICE_PORT}" + log "starting batch train/eval with BENCHMARK_SERVICE_URL=${BENCHMARK_SERVICE_URL}" + log "command: benchmark/tau2/train/run_batch_train_eval.sh --config ${OPENVIKING_CONFIG_FILE} --server-url http://127.0.0.1:${OPENVIKING_PORT} --result-dir-name ${RESULT_DIR_NAME} ${train_args[*]}" + cd "${REPO_ROOT}" + exec benchmark/tau2/train/run_batch_train_eval.sh \ + --config "${OPENVIKING_CONFIG_FILE}" \ + --server-url "http://127.0.0.1:${OPENVIKING_PORT}" \ + --result-dir-name "${RESULT_DIR_NAME}" \ + "${train_args[@]}" +} + +main() { + start_openviking_server + start_tau2_service + run_train_eval "${TRAIN_CLI_ARGS[@]}" +} + +main diff --git a/benchmark/tau2/train/rollout_executor.py b/benchmark/tau2/train/rollout_executor.py new file mode 100644 index 0000000000..a3b10f696d --- /dev/null +++ b/benchmark/tau2/train/rollout_executor.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python3 +"""Switchable Tau2 RolloutExecutor implementations.""" + +from __future__ import annotations + +from typing import Any, Literal + +from benchmark.tau2.train._rollout_helpers import ( + _as_tool_input, + _safe_float, + _stringify, + _tau2_evaluation, + _to_jsonable, +) +from benchmark.tau2.train.rollout_executor_native import NativeTau2RolloutExecutor +from benchmark.tau2.train.rollout_executor_vikingbot import ( + Tau2RolloutExecutor as VikingBotTau2RolloutExecutor, +) +from benchmark.tau2.train.rollout_executor_vikingbot import ( # re-export vikingbot-only helpers for tests + _append_final_answer_for_tau2_evaluation, + _build_rollout_messages, + _configure_tools, +) + +Tau2RolloutBackend = Literal["native", "vikingbot"] +DEFAULT_TAU2_ROLLOUT_BACKEND: Tau2RolloutBackend = "native" + + +def normalize_tau2_rollout_backend(value: Any) -> Tau2RolloutBackend: + backend = str(value or DEFAULT_TAU2_ROLLOUT_BACKEND).strip().lower() + if backend not in {"native", "vikingbot"}: + raise ValueError("rollout_backend must be 'native' or 'vikingbot'") + return backend # type: ignore[return-value] + + +def make_tau2_rollout_executor( + *, + backend: Any = DEFAULT_TAU2_ROLLOUT_BACKEND, + options: dict[str, Any] | None = None, + config_path: str | None = None, + concurrency: int = 1, + rollout_language: str = "default", +): + """Create a tau2 rollout executor for the selected backend.""" + + selected = normalize_tau2_rollout_backend(backend) + opts = dict(options or {}) + if selected == "vikingbot": + return VikingBotTau2RolloutExecutor( + config_path=opts.get("config_path") or config_path, + concurrency=concurrency, + keep_default_tools=_bool_option(opts.get("keep_default_tools"), default=True), + max_iterations=int(opts.get("max_iterations") or 30), + rollout_language=str(opts.get("rollout_language") or rollout_language), + ) + return NativeTau2RolloutExecutor( + concurrency=concurrency, + agent_llm=_optional_str(opts.get("agent_llm")), + user_llm=_optional_str(opts.get("user_llm")), + agent_llm_args=_dict_option(opts.get("agent_llm_args")), + user_llm_args=_dict_option(opts.get("user_llm_args")), + base_agent=str(opts.get("base_agent") or "llm_agent"), + user=str(opts.get("user") or "user_simulator"), + max_steps=int(opts.get("max_steps") or opts.get("max_iterations") or 200), + max_errors=int(opts.get("max_errors") or 10), + seed=int(opts.get("seed") or 300), + memory_enabled=_bool_option( + opts.get("memory_enabled", opts.get("keep_default_tools")), + default=True, + ), + retrieval_mode=str(opts.get("retrieval_mode") or "first_user_prewrite"), + search_uri=str(opts.get("search_uri") or "viking://user/memories/experiences"), + retrieval_top_k=int(opts.get("retrieval_top_k") or 4), + first_user_retrieval_top_k=_optional_int(opts.get("first_user_retrieval_top_k")), + first_user_inject_top_k=_optional_int(opts.get("first_user_inject_top_k")), + prewrite_retrieval_top_k=_optional_int(opts.get("prewrite_retrieval_top_k")), + prewrite_inject_top_k=_optional_int(opts.get("prewrite_inject_top_k")), + memory_inject_max_chars=_optional_int(opts.get("memory_inject_max_chars")), + first_user_memory_inject_max_chars=_optional_int( + opts.get("first_user_memory_inject_max_chars") + ), + prewrite_memory_inject_max_chars=_optional_int( + opts.get("prewrite_memory_inject_max_chars") + ), + openviking_url=_optional_str(opts.get("openviking_url")), + openviking_api_key=_optional_str(opts.get("openviking_api_key")), + openviking_account=_optional_str(opts.get("openviking_account")), + openviking_user=_optional_str(opts.get("openviking_user")), + openviking_timeout=float(opts.get("openviking_timeout") or 600.0), + scope_prompt=str(opts.get("scope_prompt") or ""), + rollout_language=str(opts.get("rollout_language") or rollout_language), + show_progress=_bool_option(opts.get("show_progress"), default=False), + progress_label=str(opts.get("progress_label") or "tau2"), + ) + + +# Historical name now points at the default backend for new construction. +Tau2RolloutExecutor = NativeTau2RolloutExecutor + + + +def _bool_option(value: Any, *, default: bool) -> bool: + if value is None: + return default + if isinstance(value, bool): + return value + if isinstance(value, str): + text = value.strip().lower() + if text in {"1", "true", "yes", "on"}: + return True + if text in {"0", "false", "no", "off"}: + return False + raise ValueError(f"Invalid boolean option: {value!r}") + return bool(value) + +def _dict_option(value: Any) -> dict[str, Any]: + if value is None: + return {} + if isinstance(value, dict): + return dict(value) + if isinstance(value, str) and value.strip(): + import json + + parsed = json.loads(value) + if not isinstance(parsed, dict): + raise ValueError("LLM args options must decode to an object") + return parsed + raise ValueError("LLM args options must be dict or JSON object string") + + +def _optional_int(value: Any) -> int | None: + if value is None or value == "": + return None + return int(value) + + +def _optional_str(value: Any) -> str | None: + if value is None: + return None + text = str(value) + return text if text.strip() else None + + +__all__ = [ + "DEFAULT_TAU2_ROLLOUT_BACKEND", + "NativeTau2RolloutExecutor", + "Tau2RolloutBackend", + "Tau2RolloutExecutor", + "VikingBotTau2RolloutExecutor", + "make_tau2_rollout_executor", + "normalize_tau2_rollout_backend", + "_append_final_answer_for_tau2_evaluation", + "_as_tool_input", + "_build_rollout_messages", + "_configure_tools", + "_safe_float", + "_stringify", + "_tau2_evaluation", + "_to_jsonable", +] diff --git a/benchmark/tau2/train/rollout_executor_native.py b/benchmark/tau2/train/rollout_executor_native.py new file mode 100644 index 0000000000..38109ca23a --- /dev/null +++ b/benchmark/tau2/train/rollout_executor_native.py @@ -0,0 +1,943 @@ +#!/usr/bin/env python3 +"""TAU-2 native RolloutExecutor implementation for batch policy training.""" + +from __future__ import annotations + +import asyncio +import json +import time +from dataclasses import dataclass, field +from typing import Any + +from benchmark.tau2.train._rollout_helpers import ( + _as_tool_input, + _case_trial, + _communicate_text_from_tool_input, + _is_communicate_with_user, + _message, + _metadata_message, + _stringify, + _to_jsonable, +) +from benchmark.tau2.train._rollout_helpers import ( + _tau2_evaluation as _tau2_evaluation_helper, +) +from openviking.message import Message, TextPart, ToolPart +from openviking.session.train import ( + Case, + ExecutionContext, + ExperienceSet, + Rollout, + RubricEvaluation, +) +from openviking.session.train.components.progress import ProgressPrinter +from openviking_cli.utils import get_logger + +logger = get_logger(__name__) + + +def _progress_stage_label(stage: Any, *, default: str) -> str: + stage_text = str(stage or "") + stage_name = stage_text.split(maxsplit=1)[0] + if stage_name.endswith("_rollout"): + return f"{stage_name}_start" + if stage_name.endswith("_rollout_start"): + return stage_name + return default + + +AGENT_NAME_PREFIX = "openviking_native_memory_agent" +_NATIVE_AGENT_CONFIGS: dict[str, "NativeTau2RolloutExecutor"] = {} +WRITE_TOOL_PREFIXES = ( + "toggle_", + "enable_", + "disable_", + "set_", + "reset_", + "update_", + "modify_", + "cancel_", + "book_", + "exchange_", + "return_", + "grant_", + "reboot_", +) + + +@dataclass(slots=True) +class NativeTau2RolloutExecutor: + """Execute tau2 cases through TAU-2's native orchestrator and agent APIs.""" + + agent_llm: str | None = None + user_llm: str | None = None + agent_llm_args: dict[str, Any] = field(default_factory=dict) + user_llm_args: dict[str, Any] = field(default_factory=dict) + base_agent: str = "llm_agent" + user: str = "user_simulator" + max_steps: int = 200 + max_errors: int = 10 + concurrency: int = 20 + seed: int = 300 + retrieval_mode: str = "first_user_prewrite" + search_uri: str = "viking://user/memories/experiences" + retrieval_top_k: int = 4 + first_user_retrieval_top_k: int | None = None + first_user_inject_top_k: int | None = None + prewrite_retrieval_top_k: int | None = None + prewrite_inject_top_k: int | None = None + memory_inject_max_chars: int | None = None + first_user_memory_inject_max_chars: int | None = None + prewrite_memory_inject_max_chars: int | None = None + openviking_url: str | None = None + openviking_api_key: str | None = None + openviking_account: str | None = None + openviking_user: str | None = None + openviking_timeout: float = 600.0 + memory_enabled: bool = True + scope_prompt: str = "" + rollout_language: str = "default" + log_timings: bool = True + show_progress: bool = False + progress_label: str = "tau2" + + def __post_init__(self) -> None: + if self.concurrency <= 0: + raise ValueError("concurrency must be > 0") + if self.max_steps <= 0: + raise ValueError("max_steps must be > 0") + if self.max_errors <= 0: + raise ValueError("max_errors must be > 0") + if self.retrieval_top_k <= 0: + raise ValueError("retrieval_top_k must be > 0") + if self.retrieval_mode not in {"first_user", "prewrite", "first_user_prewrite"}: + raise ValueError("retrieval_mode must be first_user, prewrite, or first_user_prewrite") + if self.rollout_language not in {"default", "zh"}: + raise ValueError("rollout_language must be 'default' or 'zh'") + for name in ( + "memory_inject_max_chars", + "first_user_memory_inject_max_chars", + "prewrite_memory_inject_max_chars", + ): + value = getattr(self, name) + if value is not None and value < 0: + raise ValueError(f"{name} must be non-negative") + self.first_user_retrieval_top_k = self.first_user_retrieval_top_k or self.retrieval_top_k + self.first_user_inject_top_k = ( + self.first_user_inject_top_k or self.first_user_retrieval_top_k + ) + self.prewrite_retrieval_top_k = self.prewrite_retrieval_top_k or self.retrieval_top_k + self.prewrite_inject_top_k = self.prewrite_inject_top_k or self.prewrite_retrieval_top_k + if self.first_user_memory_inject_max_chars is None: + self.first_user_memory_inject_max_chars = self.memory_inject_max_chars + if self.prewrite_memory_inject_max_chars is None: + self.prewrite_memory_inject_max_chars = self.memory_inject_max_chars + + async def execute( + self, + cases: list[Case], + policy_set: ExperienceSet, + context: ExecutionContext, + ) -> list[Rollout]: + self._sync_openviking_options(policy_set) + progress = ProgressPrinter( + total=len(cases), + label=_progress_stage_label(context.metadata.get("stage"), default=self.progress_label), + enabled=self.show_progress, + description=f"Running {len(cases)} tau2 native rollouts, concurrency={self.concurrency}", + ) + progress.render() + semaphore = asyncio.Semaphore(self.concurrency) + + async def run_one(index: int, case: Case) -> Rollout: + async with semaphore: + progress.start_one() + try: + rollout = await asyncio.to_thread(self._execute_one_sync, case, context, index) + progress.complete_one() + return rollout + except Exception: + progress.fail_one() + raise + + try: + return list( + await asyncio.gather(*(run_one(index, case) for index, case in enumerate(cases))) + ) + finally: + progress.finish() + + def _sync_openviking_options(self, policy_set: ExperienceSet) -> None: + metadata = dict(policy_set.metadata or {}) + self.openviking_url = self.openviking_url or _optional_metadata_str( + metadata, "openviking_url", "server_url" + ) + self.openviking_api_key = self.openviking_api_key or _optional_metadata_str( + metadata, "openviking_api_key", "api_key" + ) + self.openviking_account = self.openviking_account or _optional_metadata_str( + metadata, "openviking_account", "account_id", "account" + ) + self.openviking_user = self.openviking_user or _optional_metadata_str( + metadata, "openviking_user", "user_id", "user" + ) + + def _execute_one_sync(self, case: Case, context: ExecutionContext, case_index: int) -> Rollout: + started_at = time.perf_counter() + domain = str(case.input["domain"]) + task_id = str(case.input["task_id"]) + task_no = int(case.input["task_no"]) + data_split = str(case.input["data_split"]) + trial = _case_trial(case) + seed = _case_seed(self.seed, case_index=case_index, eval_trial=trial) + + _ensure_tau2_llm_api_bases() + + from tau2.evaluator.evaluator import EvaluationType, evaluate_simulation + from tau2.registry import registry + from tau2.run import run_task + + llm_agent, llm_args_agent, llm_user, llm_args_user = _resolve_llm_runtime_config(self) + tasks = registry.get_tasks_loader(domain)() + task_by_id = {str(task.id): task for task in tasks} + try: + task = task_by_id[task_id] + except KeyError as exc: + raise ValueError(f"tau2 task not found domain={domain} task_id={task_id}") from exc + + agent_name = self.base_agent + if self.memory_enabled: + agent_name = _register_native_memory_agent(self) + + simulation = run_task( + domain=domain, + task=task, + agent=agent_name, + user=self.user, + llm_agent=llm_agent, + llm_args_agent=llm_args_agent, + llm_user=llm_user, + llm_args_user=llm_args_user, + max_steps=self.max_steps, + max_errors=self.max_errors, + evaluation_type=EvaluationType.ALL, + seed=seed, + ) + reward_info = simulation.reward_info + if reward_info is None: + reward_info = evaluate_simulation( + domain=domain, + task=task, + simulation=simulation, + evaluation_type=EvaluationType.ALL, + solo_mode=False, + ) + simulation.reward_info = reward_info + reward = _to_jsonable(getattr(reward_info, "reward", 0.0)) + evaluation_result = _to_jsonable(reward_info) + messages = _build_rollout_messages_from_simulation( + simulation=simulation, + reward=reward, + evaluation_result=evaluation_result, + ) + memory_context = _memory_context_from_simulation(simulation) + rollout = Rollout( + case=case, + messages=messages, + policy_snapshot_id=context.policy_snapshot_id, + evaluation=_tau2_evaluation(reward=reward, evaluation_result=evaluation_result), + metadata={ + "rollout_backend": "native", + "domain": domain, + "data_split": data_split, + "task_no": task_no, + "task_id": task_id, + "eval_trial": case.input.get("eval_trial"), + "eval_trial_count": case.input.get("eval_trial_count"), + "train_trial": case.input.get("train_trial"), + "train_trial_count": case.input.get("train_trial_count"), + "original_case_name": case.input.get("original_case_name"), + "seed": seed, + "reward": reward, + "evaluation_result": evaluation_result, + "termination_reason": getattr(simulation, "termination_reason", None), + "duration": getattr(simulation, "duration", None), + "agent_cost": getattr(simulation, "agent_cost", None), + "user_cost": getattr(simulation, "user_cost", None), + "tools_used": _tool_usage_from_simulation(simulation), + "memory": memory_context, + "memory_enabled": self.memory_enabled, + "retrieval_mode": self.retrieval_mode if self.memory_enabled else None, + "search_uri": self.search_uri if self.memory_enabled else None, + "execution_metadata": dict(context.metadata), + }, + ) + if self.log_timings: + logger.info( + "tau2 native rollout timing case=%s total_ms=%.1f task_id=%s task_no=%s " + "split=%s reward=%s message_count=%s", + case.name, + (time.perf_counter() - started_at) * 1000.0, + task_id, + task_no, + data_split, + reward, + len(rollout.messages), + ) + return rollout + + +def _resolve_llm_runtime_config( + executor: NativeTau2RolloutExecutor, +) -> tuple[str, dict[str, Any], str, dict[str, Any]]: + """Resolve tau2 native LLM settings for direct ``run_task`` calls. + + ``tau2.run.run_task`` does not apply ``RunConfig`` defaults. Passing + ``None`` through leaves ``LLMAgent.llm`` unset, and the orchestrator then + fails during ``set_seed`` before the first model call. Mirror tau2's + RunConfig defaults here while still letting request options and env vars + override the model names. + """ + + import os + from copy import deepcopy + + from tau2.config import ( + DEFAULT_LLM_AGENT, + DEFAULT_LLM_ARGS_AGENT, + DEFAULT_LLM_ARGS_USER, + DEFAULT_LLM_USER, + ) + + llm_agent = _first_non_empty( + executor.agent_llm, + os.getenv("TAU2_AGENT_LLM"), + DEFAULT_LLM_AGENT, + name="agent_llm", + ) + llm_user = _first_non_empty( + executor.user_llm, + os.getenv("TAU2_USER_LLM"), + DEFAULT_LLM_USER, + name="user_llm", + ) + llm_args_agent = deepcopy(DEFAULT_LLM_ARGS_AGENT or {}) + llm_args_agent.update(dict(executor.agent_llm_args or {})) + llm_args_user = deepcopy(DEFAULT_LLM_ARGS_USER or {}) + llm_args_user.update(dict(executor.user_llm_args or {})) + return llm_agent, llm_args_agent, llm_user, llm_args_user + + +def _first_non_empty(*values: Any, name: str) -> str: + for value in values: + if value is None: + continue + text = str(value).strip() + if text: + return text + raise ValueError(f"{name} must be set for tau2 native rollout") + + +def _ensure_tau2_llm_api_bases() -> None: + import os + + base_url = ( + os.environ.get("OPENAI_API_BASE") + or os.environ.get("OPENAI_BASE_URL") + or os.environ.get("ARK_BASE_URL") + ) + if not base_url: + return + os.environ.setdefault("OPENAI_API_BASE", base_url) + os.environ.setdefault("OPENAI_BASE_URL", base_url) + os.environ.setdefault("AGENT_API_BASE", base_url) + os.environ.setdefault("USER_API_BASE", base_url) + + +def _optional_metadata_str(metadata: dict[str, Any], *keys: str) -> str | None: + for key in keys: + value = metadata.get(key) + if value is None: + continue + text = str(value).strip() + if text: + return text + return None + + +def _register_native_memory_agent(executor: NativeTau2RolloutExecutor) -> str: + from tau2.agent.llm_agent import LLMAgent, LLMAgentState + from tau2.data_model.message import ( + AssistantMessage, + MultiToolMessage, + SystemMessage, + UserMessage, + ) + from tau2.registry import registry + from tau2.utils.llm_utils import generate + + agent_name = _native_agent_name(executor) + _NATIVE_AGENT_CONFIGS[agent_name] = executor + if agent_name in registry.get_agents(): + return agent_name + + class OpenVikingNativeMemoryAgent(LLMAgent): + @property + def _executor(self) -> NativeTau2RolloutExecutor: + return _NATIVE_AGENT_CONFIGS[agent_name] + + def get_init_state(self, message_history=None): + executor_config = self._executor + state = super().get_init_state(message_history) + self._openviking_memory_contexts: list[str] = [] + if executor_config.scope_prompt: + state.system_messages.append( + SystemMessage(role="system", content=executor_config.scope_prompt) + ) + if executor_config.rollout_language == "zh": + state.system_messages.append( + SystemMessage( + role="system", + content=( + "Communicate with the user and write final responses in Chinese. " + "Do not translate tool names, identifiers, JSON field names, " + "reservation IDs, flight numbers, or other structured values." + ), + ) + ) + return state + + def _retrieve( + self, + query: str, + *, + search_limit: int, + inject_limit: int, + inject_max_chars: int | None = None, + exclude_uris: set[str] | None = None, + ) -> tuple[str, list[dict[str, Any]], set[str]]: + """Retrieve and rank memories. + + Returns: + block: joined text of injected memories + rows: detail rows for each match + injected_uris: set of URIs that were actually injected + """ + executor_config = self._executor + client = _client(executor_config) + rows: list[dict[str, Any]] = [] + injected_uris: set[str] = set() + try: + result = client.search( + query=query, + target_uri=executor_config.search_uri, + limit=search_limit, + ) + memories = list(getattr(result, "memories", []) or []) + # URI deduplication: keep the highest-scoring match per URI + deduped: dict[str, Any] = {} + for match in memories[:search_limit]: + uri = getattr(match, "uri", "") + if not uri: + continue + if exclude_uris and uri in exclude_uris: + continue + if uri in deduped: + prev_score = getattr(deduped[uri], "score", 0) or 0 + curr_score = getattr(match, "score", 0) or 0 + if curr_score <= prev_score: + continue + deduped[uri] = match + deduped_memories = sorted( + deduped.values(), + key=lambda m: getattr(m, "score", 0) or 0, + reverse=True, + ) + + blocks: list[str] = [] + injected_chars_used = 0 + for index, match in enumerate(deduped_memories, 1): + uri = getattr(match, "uri", "") + text, read_error = _read_memory_text(client, match) + clean_text = text.strip() + block_text = ( + f"Memory {index} ({uri}):\n{clean_text}" if clean_text else "" + ) + block_chars = len(block_text) + budget_used_before = injected_chars_used + budget_dropped = False + truncated = False + injected = index <= inject_limit and bool(block_text) + if injected and inject_max_chars is not None: + remaining = inject_max_chars - injected_chars_used + if remaining <= 0: + injected = False + budget_dropped = True + elif block_chars > remaining: + if not blocks: + block_text = block_text[:remaining] + block_chars = len(block_text) + truncated = True + else: + injected = False + budget_dropped = True + if injected: + injected_chars_used += block_chars + injected_uris.add(uri) + row = { + "uri": uri, + "score": getattr(match, "score", None), + "level": getattr(match, "level", None), + "text_chars": len(text), + "block_chars": block_chars, + "injected": injected, + "inject_max_chars": inject_max_chars, + "inject_budget_used_before": budget_used_before, + "inject_budget_used_after": injected_chars_used, + "inject_budget_dropped": budget_dropped, + "inject_budget_truncated": truncated, + } + if read_error: + row["read_error"] = read_error + rows.append(row) + if injected: + blocks.append(block_text) + return "\n\n".join(blocks), rows, injected_uris + finally: + client.close() + + def _generate(self, messages): + def _is_empty_assistant(response: Any) -> bool: + content = str(getattr(response, "content", "") or "") + tool_calls = getattr(response, "tool_calls", None) or [] + return not content.strip() and not tool_calls + + try: + response = generate( + model=self.llm, + tools=self.tools, + messages=messages, + is_agent=True, + **self.llm_args, + ) + if not _is_empty_assistant(response): + return response + except json.JSONDecodeError: + retry_prompt = ( + "Retry the last assistant step once. If you call a tool, " + "the tool arguments must be syntactically valid JSON." + ) + else: + retry_prompt = ( + "Retry the last assistant step once. Return either a useful natural " + "language response or a valid tool call; do not return an empty assistant message." + ) + try: + response = generate( + model=self.llm, + tools=self.tools, + messages=messages + [SystemMessage(role="system", content=retry_prompt)], + is_agent=True, + **self.llm_args, + ) + if not _is_empty_assistant(response): + return response + return AssistantMessage( + role="assistant", + content="I need to continue with the available task information.", + raw_data={"openviking_memory_agent_error": "empty_assistant_message"}, + ) + except json.JSONDecodeError as exc: + return AssistantMessage( + role="assistant", + content="I need to continue with the available task information.", + raw_data={ + "openviking_memory_agent_error": "invalid_tool_call_json", + "error": str(exc), + }, + ) + + def generate_next_message(self, message, state: LLMAgentState): + executor_config = self._executor + if isinstance(message, MultiToolMessage): + state.messages.extend(message.tool_messages) + else: + state.messages.append(message) + + role = getattr(message, "role", "") + role_value = getattr(role, "value", role) + is_first_user = ( + executor_config.retrieval_mode in {"first_user", "first_user_prewrite"} + and str(role_value) == "user" + and not getattr(self, "_first_user_memory_injected", False) + ) + injected_uris: set[str] = getattr(self, "_injected_memory_uris", set()) + + if is_first_user: + query = str(getattr(message, "content", "") or "") + block, matches, new_injected = self._retrieve( + query, + search_limit=int( + executor_config.first_user_retrieval_top_k + or executor_config.retrieval_top_k + ), + inject_limit=int( + executor_config.first_user_inject_top_k or executor_config.retrieval_top_k + ), + inject_max_chars=executor_config.first_user_memory_inject_max_chars, + exclude_uris=injected_uris, + ) + self._first_user_memory_injected = True + if block: + injected_uris.update(new_injected) + self._injected_memory_uris = injected_uris + # Prepend experience reminder to the user message content + # so it becomes part of the conversation history (visible in messages.json) + reminder_prefix = ( + "[Experience Reminder]\n" + "## Relevant Agent Experience\n\n" + + block + + "\n\n---\n\n" + ) + message.content = reminder_prefix + str(message.content or "") + self._openviking_memory_contexts.append(block) + + assistant_message = self._generate(state.system_messages + state.messages) + if executor_config.retrieval_mode in {"prewrite", "first_user_prewrite"}: + tool_calls = list(getattr(assistant_message, "tool_calls", None) or []) + write_calls = [call for call in tool_calls if _is_write_tool_call(call)] + if write_calls: + query = _tool_call_query(write_calls, state.messages) + block, matches, new_injected = self._retrieve( + query, + search_limit=int( + executor_config.prewrite_retrieval_top_k + or executor_config.retrieval_top_k + ), + inject_limit=int( + executor_config.prewrite_inject_top_k or executor_config.retrieval_top_k + ), + inject_max_chars=executor_config.prewrite_memory_inject_max_chars, + exclude_uris=injected_uris, + ) + if block: + injected_uris.update(new_injected) + self._injected_memory_uris = injected_uris + self._openviking_memory_contexts.append(block) + reminder_content = ( + "[Experience Reminder]\n" + "## Relevant Agent Experience (before write action)\n\n" + + block + ) + # Inject as a user message so it's part of the conversation history + state.messages.append(UserMessage(role="user", content=reminder_content)) + assistant_message = self._generate( + state.system_messages + state.messages + ) + contexts = list(getattr(self, "_openviking_memory_contexts", []) or []) + if contexts: + raw_data = dict(getattr(assistant_message, "raw_data", None) or {}) + raw_data["openviking_memory_context"] = "\n\n".join(contexts) + assistant_message.raw_data = raw_data + state.messages.append(assistant_message) + return assistant_message, state + + registry.register_agent(OpenVikingNativeMemoryAgent, agent_name) + return agent_name + + +def _native_agent_name(executor: NativeTau2RolloutExecutor) -> str: + import hashlib + + payload = { + "retrieval_mode": executor.retrieval_mode, + "search_uri": executor.search_uri, + "first_user_retrieval_top_k": executor.first_user_retrieval_top_k, + "first_user_inject_top_k": executor.first_user_inject_top_k, + "prewrite_retrieval_top_k": executor.prewrite_retrieval_top_k, + "prewrite_inject_top_k": executor.prewrite_inject_top_k, + "memory_inject_max_chars": executor.memory_inject_max_chars, + "first_user_memory_inject_max_chars": executor.first_user_memory_inject_max_chars, + "prewrite_memory_inject_max_chars": executor.prewrite_memory_inject_max_chars, + "openviking_url_set": bool(executor.openviking_url), + "openviking_api_key_set": bool(executor.openviking_api_key), + "openviking_account": executor.openviking_account, + "openviking_user": executor.openviking_user, + "openviking_timeout": executor.openviking_timeout, + "scope_prompt": executor.scope_prompt, + "rollout_language": executor.rollout_language, + } + digest = hashlib.sha1( + json.dumps(payload, ensure_ascii=False, sort_keys=True, default=str).encode("utf-8") + ).hexdigest()[:12] + return f"{AGENT_NAME_PREFIX}_{digest}" + + +def _client(executor: NativeTau2RolloutExecutor): + import openviking as ov + + client = ov.SyncHTTPClient( + url=executor.openviking_url, + api_key=executor.openviking_api_key, + account=executor.openviking_account, + user=executor.openviking_user, + timeout=executor.openviking_timeout, + extra_headers={}, + profile_enabled=False, + ) + client.initialize() + return client + + +def _read_memory_text(client: Any, match: Any) -> tuple[str, str | None]: + try: + return client.read(getattr(match, "uri", "")), None + except Exception as exc: + fallback = getattr(match, "abstract", "") or getattr(match, "overview", "") or "" + return fallback, f"{type(exc).__name__}: {exc}" + + +def _tool_call_name(tool_call: Any) -> str: + if isinstance(tool_call, dict): + return str(tool_call.get("name") or tool_call.get("function", {}).get("name") or "") + return str(getattr(tool_call, "name", "") or "") + + +def _tool_call_arguments(tool_call: Any) -> Any: + if isinstance(tool_call, dict): + return tool_call.get("arguments") or tool_call.get("function", {}).get("arguments") or {} + return getattr(tool_call, "arguments", {}) or {} + + +def _is_write_tool_call(tool_call: Any) -> bool: + name = _tool_call_name(tool_call) + return bool(name) and name.startswith(WRITE_TOOL_PREFIXES) + + +def _tool_call_query(tool_calls: list[Any], state_messages: list[Any]) -> str: + rendered = [] + for call in tool_calls: + rendered.append( + f"{_tool_call_name(call) or 'unknown_tool'}(" + f"{json.dumps(_tool_call_arguments(call), ensure_ascii=False, sort_keys=True, default=str)}" + ")" + ) + recent_user = [ + str(getattr(message, "content", "") or "") + for message in state_messages[-8:] + if str(getattr(message, "role", "")) == "user" + and str(getattr(message, "content", "") or "").strip() + ] + recent_observations = [ + str(getattr(message, "content", "") or "")[:600] + for message in state_messages[-12:] + if str(getattr(message, "role", "")) == "tool" + and str(getattr(message, "content", "") or "").strip() + ] + parts = [ + "Before executing write-like tool call(s): " + "; ".join(rendered), + "Recent user context: " + " | ".join(recent_user[-3:]), + ] + if recent_observations: + parts.append("Recent tool observations: " + " | ".join(recent_observations[-4:])) + return "\n".join(parts) + + +def _case_seed(base_seed: int, *, case_index: int, eval_trial: Any) -> int: + trial = 0 + try: + if eval_trial is not None: + trial = int(eval_trial) + except (TypeError, ValueError): + trial = 0 + return int(base_seed) + case_index + trial * 100_000 + + +def _build_rollout_messages_from_simulation( + *, + simulation: Any, + reward: Any, + evaluation_result: Any, +) -> list[Message]: + messages: list[Message] = [] + pending_tool_calls: dict[str, tuple[str, dict[str, Any]]] = {} + for index, message in enumerate(getattr(simulation, "messages", []) or []): + converted = _simulation_message_to_rollout_messages( + message, + index, + pending_tool_calls=pending_tool_calls, + ) + messages.extend(converted) + for call_id, (tool_name, tool_input) in pending_tool_calls.items(): + messages.append( + Message( + id=f"tau2-tool-pending-{index if 'index' in locals() else 0}-{len(messages)}", + role="assistant", + parts=[ + ToolPart( + tool_id=call_id, + tool_name=tool_name, + tool_input=tool_input, + tool_status="running", + ) + ], + ) + ) + reward_jsonable = _to_jsonable(reward) + evaluation_jsonable = _to_jsonable(evaluation_result) + success = reward_jsonable == 1 or reward_jsonable == 1.0 + messages.append( + _message( + "tau2-reward", + "user", + f"task_success: {success}\ntask_reward: {reward_jsonable}\n" + f"evaluation report: {_stringify(evaluation_jsonable)}", + ) + ) + return messages + + +def _simulation_message_to_rollout_messages( + message: Any, + index: int, + *, + pending_tool_calls: dict[str, tuple[str, dict[str, Any]]] | None = None, +) -> list[Message]: + role = _role_value(getattr(message, "role", "assistant")) + if role == "system": + content = str(getattr(message, "content", "") or "") + return [ + _metadata_message( + f"tau2-system-{index}", + f"system:\n{content}", + ) + ] + if role in {"user", "assistant"}: + content = str(getattr(message, "content", "") or "") + tool_calls = list(getattr(message, "tool_calls", None) or []) + if tool_calls: + rows = [] + for call_idx, call in enumerate(tool_calls): + call_id = str(getattr(call, "id", "") or f"tau2-tool-{index}-{call_idx}") + tool_name = _tool_call_name(call) + tool_input = _as_tool_input(_tool_call_arguments(call)) + if _is_communicate_with_user(tool_name): + assistant_text = _communicate_text_from_tool_input(tool_input) + if assistant_text.strip(): + rows.append( + Message( + id=f"tau2-communicate-assistant-{index}-{call_idx}", + role="assistant", + parts=[TextPart(text=assistant_text)], + created_at=getattr(message, "timestamp", None), + ) + ) + if pending_tool_calls is not None: + pending_tool_calls[call_id] = (tool_name, tool_input) + continue + if pending_tool_calls is not None: + pending_tool_calls[call_id] = (tool_name, tool_input) + continue + rows.append( + Message( + id=f"tau2-tool-call-{index}-{call_idx}", + role="assistant" if role == "assistant" else "user", + parts=[ + ToolPart( + tool_id=call_id, + tool_name=tool_name, + tool_input=tool_input, + tool_status="running", + ) + ], + created_at=getattr(message, "timestamp", None), + ) + ) + return rows + return [ + Message( + id=f"tau2-{role}-{index}", + role="user" if role == "user" else "assistant", + parts=[TextPart(text=content)], + created_at=getattr(message, "timestamp", None), + ) + ] + if role == "tool": + call_id = str(getattr(message, "id", "") or f"tau2-tool-{index}") + pending = pending_tool_calls.pop(call_id, None) if pending_tool_calls is not None else None + tool_name, tool_input = pending if pending is not None else ("unknown", None) + output = str(getattr(message, "content", "") or "") + if _is_communicate_with_user(tool_name): + if not output.strip(): + return [] + return [ + Message( + id=f"tau2-communicate-user-{index}", + role="user", + parts=[TextPart(text=output)], + created_at=getattr(message, "timestamp", None), + ) + ] + return [ + Message( + id=f"tau2-tool-result-{index}", + role="user", + parts=[ + ToolPart( + tool_id=call_id, + tool_name=tool_name, + tool_input=tool_input, + tool_output=output, + tool_status="error" + if bool(getattr(message, "error", False)) + else "completed", + ) + ], + created_at=getattr(message, "timestamp", None), + ) + ] + content = str(getattr(message, "content", "") or "") + return [_message(f"tau2-message-{index}", "assistant", content)] + + +def _tool_usage_from_simulation(simulation: Any) -> list[dict[str, Any]]: + usages: list[dict[str, Any]] = [] + pending: dict[str, dict[str, Any]] = {} + for message in getattr(simulation, "messages", []) or []: + role = _role_value(getattr(message, "role", "")) + if role in {"user", "assistant"}: + for call in list(getattr(message, "tool_calls", None) or []): + call_id = str(getattr(call, "id", "") or f"call_{len(usages)}") + row = { + "tool_name": _tool_call_name(call), + "args": _tool_call_arguments(call), + "requestor": getattr(call, "requestor", role), + } + pending[call_id] = row + usages.append(row) + elif role == "tool": + call_id = str(getattr(message, "id", "") or "") + row = pending.get(call_id) + if row is not None: + row["result"] = getattr(message, "content", None) + row["error"] = bool(getattr(message, "error", False)) + return usages + + +def _memory_context_from_simulation(simulation: Any) -> str | None: + blocks = [] + for message in getattr(simulation, "messages", []) or []: + raw_data = getattr(message, "raw_data", None) + if isinstance(raw_data, dict) and raw_data.get("openviking_memory_context"): + blocks.append(str(raw_data["openviking_memory_context"])) + return "\n\n".join(blocks) if blocks else None + + +def _role_value(role: Any) -> str: + return str(getattr(role, "value", role)) + + +def _tau2_evaluation(*, reward: Any, evaluation_result: Any) -> RubricEvaluation: + return _tau2_evaluation_helper( + reward=reward, evaluation_result=evaluation_result, source="tau2_native_executor" + ) diff --git a/benchmark/tau2/train/rollout_executor_vikingbot.py b/benchmark/tau2/train/rollout_executor_vikingbot.py new file mode 100644 index 0000000000..9145ae1309 --- /dev/null +++ b/benchmark/tau2/train/rollout_executor_vikingbot.py @@ -0,0 +1,1592 @@ +#!/usr/bin/env python3 +"""Tau2 RolloutExecutor implementation for batch policy training.""" + +from __future__ import annotations + +import asyncio +import json +import posixpath +import re +import time +from collections.abc import Callable +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from benchmark.tau2.train._rollout_helpers import ( + _as_tool_input, + _case_trial, + _communicate_text_from_tool_input, + _is_communicate_with_user, + _message, + _metadata_message, + _stringify, + _to_jsonable, +) +from benchmark.tau2.train._rollout_helpers import ( + _tau2_evaluation as _tau2_evaluation_helper, +) +from openviking.message import Message, ToolPart +from openviking.session.train import ( + Case, + ExecutionContext, + ExperienceSet, + Rollout, + RubricEvaluation, +) +from openviking_cli.utils import get_logger + +logger = get_logger(__name__) + + +def _tau2_policy_current_time_match(policy: str) -> re.Match[str] | None: + return re.search( + r"(?im)\bcurrent\s+time\s+is\s+" + r"(\d{4}-\d{2}-\d{2})[ T](\d{2}:\d{2}:\d{2})\s*([A-Z]{2,5})?", + policy or "", + ) + + +def _tau2_policy_current_time_display(policy: str) -> str | None: + """Return tau2's authoritative business clock for prompt display.""" + match = _tau2_policy_current_time_match(policy) + if not match: + return None + date_part, time_part, tz_name = match.groups() + suffix = f" ({tz_name}; from tau2 policy)" if tz_name else " (from tau2 policy)" + return f"{date_part} {time_part}{suffix}" + + +def _tau2_policy_current_time_iso(policy: str) -> str | None: + """Return tau2's authoritative business clock as an ISO timestamp. + + Tau2 airline embeds the authoritative business clock in the policy, e.g. + ``The current time is 2024-05-15 15:00:00 EST.`` Rollout artifacts should + use that clock for message ``created_at`` so downstream trajectory/experience + extraction does not treat the wall-clock run timestamp as business time. + """ + match = _tau2_policy_current_time_match(policy) + if not match: + return None + + date_part, time_part, tz_name = match.groups() + tz_offsets = { + "UTC": "+00:00", + "GMT": "+00:00", + "EST": "-05:00", + "EDT": "-04:00", + "CST": "-06:00", + "CDT": "-05:00", + "MST": "-07:00", + "MDT": "-06:00", + "PST": "-08:00", + "PDT": "-07:00", + } + offset = tz_offsets.get((tz_name or "").upper()) + if offset is not None: + return f"{date_part}T{time_part}{offset}" + return f"{date_part}T{time_part}" + + +def _viking_is_tool_result_success(result: Any) -> bool: + # Mirror vikingbot.agent.loop._is_tool_result_success locally to avoid importing + # private names from the bot package. + if result is None or isinstance(result, Exception): + return False + text = str(result).lstrip() + return bool(text) and not text.startswith("Error:") + + +def _tool_provider_cls(): + from benchmark.tau2.common.tau2_env.tau2_tool_provider import Tau2BenchToolProvider + + return Tau2BenchToolProvider + + +def _vikingbot_imports() -> dict[str, Any]: + try: + from vikingbot.agent.context import ContextBuilder + from vikingbot.agent.loop import ( + AgentLoop, + _PlainTextContext, + _PlainTextDelivered, + _PlainTextFinal, + ) + from vikingbot.agent.tools.base import Tool + from vikingbot.bus.queue import MessageBus + from vikingbot.cli.commands import _init_bot_data, _make_provider + from vikingbot.config.loader import ensure_config + from vikingbot.config.schema import SessionKey + from vikingbot.sandbox.manager import SandboxManager + from vikingbot.session.manager import SessionManager + from vikingbot.utils.helpers import get_source_workspace_path + except ImportError as exc: # pragma: no cover - benchmark environment dependency + raise RuntimeError( + "Failed to import vikingbot. Source benchmark/tau2/vikingbot/setup_env.sh first." + ) from exc + + return { + "AgentLoop": AgentLoop, + "ContextBuilder": ContextBuilder, + "_PlainTextContext": _PlainTextContext, + "_PlainTextDelivered": _PlainTextDelivered, + "_PlainTextFinal": _PlainTextFinal, + "Tool": Tool, + "MessageBus": MessageBus, + "_init_bot_data": _init_bot_data, + "_make_provider": _make_provider, + "ensure_config": ensure_config, + "SessionKey": SessionKey, + "SandboxManager": SandboxManager, + "SessionManager": SessionManager, + "get_source_workspace_path": get_source_workspace_path, + } + + +def _make_tau2_tool( + schema: dict[str, Any], + provider: Any, + *, + tool_lock: "_AsyncRWLock | None" = None, + is_write_tool: bool = False, + record_tool_timing: Callable[[str, float], None] | None = None, +): + Tool = _vikingbot_imports()["Tool"] + + class Tau2Tool(Tool): + """Bridge tau2 tool schema into VikingBot Tool interface.""" + + def __init__(self, tool_schema: dict[str, Any], tool_provider: Any): + self._schema = tool_schema + self._provider = tool_provider + function_def = tool_schema.get("function", {}) if isinstance(tool_schema, dict) else {} + self._name = function_def.get("name", "") + self._description = function_def.get("description", "") + self._parameters = function_def.get("parameters", {}) + + @property + def name(self) -> str: + return self._name + + @property + def description(self) -> str: + return self._description + + @property + def parameters(self) -> dict[str, Any]: + return self._parameters + + async def execute(self, tool_context: Any, **kwargs: Any) -> str: + del tool_context + started_at = time.perf_counter() + + try: + if tool_lock is None: + return await asyncio.to_thread(self._provider.call_tool, self._name, kwargs) + + if is_write_tool: + async with tool_lock.writer(): + return await asyncio.to_thread(self._provider.call_tool, self._name, kwargs) + + # Read path: acquire a shared (reader) lock so concurrent read tools + # don't block each other. + async with tool_lock.reader(): + return await asyncio.to_thread(self._provider.call_tool, self._name, kwargs) + finally: + if record_tool_timing is not None: + record_tool_timing(self._name, _elapsed_ms(started_at)) + + return Tau2Tool(schema, provider) + + +def _make_search_experience_tool(): + Tool = _vikingbot_imports()["Tool"] + + class SearchExperienceTool(Tool): + @property + def name(self) -> str: + return "search_experience" + + @property + def description(self) -> str: + return ( + "Search OpenViking case memories under the current user, read each matched " + "case's Linked Experiences section, and return candidate case summaries plus " + "linked experience URIs. Use read_experience to open selected experience URIs." + ) + + @property + def parameters(self) -> dict[str, Any]: + return { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Natural-language query describing the current task intent, target object, operation, policy/tool keywords.", + }, + "limit": { + "type": "integer", + "description": "Maximum candidate cases to inspect and return.", + "default": 10, + }, + }, + "required": ["query"], + } + + async def execute( + self, tool_context: Any, query: str, limit: int = 10, **kwargs: Any + ) -> str: + del kwargs + client = None + try: + from vikingbot.openviking_mount.ov_server import VikingClient + + client = await VikingClient.create() + target_uri = _current_cases_uri(client) + result = await client.search(query, target_uri=target_uri, limit=max(1, int(limit))) + memories = result.get("memories", []) if isinstance(result, dict) else [] + candidates = [ + await _experience_search_summary(client, item, rank) + for rank, item in enumerate(memories, start=1) + ] + return json.dumps( + { + "query": query, + "target_uri": target_uri, + "count": len(candidates), + "candidates": candidates, + }, + ensure_ascii=False, + indent=2, + ) + except Exception as exc: + logger.warning("search_experience failed: %s", exc) + return f"Error searching experience candidates: {exc}" + finally: + if client is not None: + await client.close() + + return SearchExperienceTool() + + +def _make_read_experience_tool(): + Tool = _vikingbot_imports()["Tool"] + + class ReadExperienceTool(Tool): + @property + def name(self) -> str: + return "read_experience" + + @property + def description(self) -> str: + return "Read one OpenViking experience memory by full URI. Returns Markdown." + + @property + def parameters(self) -> dict[str, Any]: + return { + "type": "object", + "properties": { + "experience_uri": { + "type": "string", + "description": "Full Viking URI of the experience memory to read.", + }, + }, + "required": ["experience_uri"], + } + + async def execute(self, tool_context: Any, experience_uri: str, **kwargs: Any) -> str: + del tool_context, kwargs + client = None + try: + from vikingbot.openviking_mount.ov_server import VikingClient + + client = await VikingClient.create() + experience_uri = str(experience_uri or "").strip() + if "/memories/experiences/" not in experience_uri: + return f"Error: URI is not an experience memory: {experience_uri}" + content = await client.read_content(experience_uri, level="read") + if not content: + return ( + "# Loaded Experience\n\n" + f"Experience URI: `{experience_uri}`\n\n" + "Error: experience content not found." + ) + return "\n".join( + [ + "# Loaded Experience", + "", + f"Experience URI: `{experience_uri}`", + "", + content.rstrip(), + ] + ).rstrip() + except Exception as exc: + logger.warning("read_experience failed: %s", exc) + return f"Error reading experience memory: {exc}" + finally: + if client is not None: + await client.close() + + return ReadExperienceTool() + + +def _current_cases_uri(client: Any) -> str: + return f"{client._memory_target_uri(None).rstrip('/')}/cases" + + +def _case_uri(item: Any) -> str: + if isinstance(item, dict): + return str(item.get("uri") or "") + return str(getattr(item, "uri", "") or "") + + +def _case_score(item: Any) -> float: + value = item.get("score", 0.0) if isinstance(item, dict) else getattr(item, "score", 0.0) + try: + return float(value) + except (TypeError, ValueError): + return 0.0 + + +def _case_abstract(item: Any) -> str: + return str( + item.get("abstract", "") if isinstance(item, dict) else getattr(item, "abstract", "") or "" + ) + + +def _filename_name(uri: str) -> str: + return str(uri or "").rstrip("/").rsplit("/", 1)[-1].removesuffix(".md") + + +def _markdown_section(content: str, heading: str) -> str: + match = re.search( + rf"(?ims)^##\s+{re.escape(heading)}\s*\n(.*?)(?=^##\s+|\Z)", + content or "", + ) + return match.group(1).strip() if match else "" + + +def _parse_json_object(value: str) -> dict[str, Any]: + try: + parsed = json.loads(str(value or "").strip()) + except Exception: + return {} + return parsed if isinstance(parsed, dict) else {} + + +def _shorten(value: Any, limit: int = 240) -> str: + text = re.sub(r"\s+", " ", str(value or "")).strip() + return text if len(text) <= limit else text[: max(0, limit - 1)].rstrip() + "…" + + +def _linked_experience_count(content: str) -> int: + section = _markdown_section(content, "Linked Experiences") + if not section: + return 0 + links = re.findall(r"\[[^\]]+\]\(([^)\s]+)\)", section) + if links: + return len(links) + return sum(1 for line in section.splitlines() if line.strip().startswith("- ")) + + +async def _experience_search_summary(client: Any, item: Any, rank: int) -> dict[str, Any]: + case_uri = _case_uri(item) + summary: dict[str, Any] = { + "rank": rank, + "score": round(_case_score(item), 6), + "case_name": _filename_name(case_uri), + "case_uri": case_uri, + "case_abstract": _shorten(_case_abstract(item), 360), + "experiences": [], + } + if not case_uri: + return summary + try: + content = await client.read_content(case_uri, level="read") + except Exception: + return summary + input_text = _markdown_section(content, "Input") + input_obj = _parse_json_object(input_text) + exp_uris = _linked_experience_uris(content, source_uri=case_uri) + # ponytail: fetch Situation snippet per experience so agent can gate read_experience on applicability. + experiences: list[dict[str, Any]] = [] + for idx, exp_uri in enumerate(exp_uris, start=1): + exp_entry: dict[str, Any] = { + "index": idx, + "name": _filename_name(exp_uri), + "uri": exp_uri, + "situation": "", + } + try: + exp_content = await client.read_content(exp_uri, level="read") + except Exception: + exp_content = "" + situation = _markdown_section(exp_content, "Situation") if exp_content else "" + # ponytail: cap at ~600 chars per exp to bound search-result tokens; exclusions ("不适用于"/"not apply") are preserved. + exp_entry["situation"] = _shorten(situation, 600) + experiences.append(exp_entry) + summary.update( + { + "task_signature": _shorten(_markdown_section(content, "Task Signature")), + "input_summary": _shorten(input_obj.get("summary") if input_obj else input_text), + "experiences": experiences, + } + ) + return summary + + +def _linked_experience_uris(content: str, *, source_uri: str) -> list[str]: + section = _markdown_section(content, "Linked Experiences") + if not section: + return [] + targets = re.findall(r"\[[^\]]+\]\(([^)\s]+)\)", section) + if not targets: + targets = [ + line.lstrip("- ").strip() + for line in section.splitlines() + if line.strip().startswith("- ") + ] + uris: list[str] = [] + for target in targets: + uri = _resolve_case_link_uri(target, source_uri=source_uri) + if "/memories/experiences/" in uri and uri not in uris: + uris.append(uri) + return uris + + +def _resolve_case_link_uri(target: str, *, source_uri: str) -> str: + target = str(target or "").strip() + if not target: + return "" + if "://" in target: + return target + if "/" not in target: + target = f"../experiences/{target.removesuffix('.md')}.md" + if not source_uri.startswith("viking://"): + return target + source_dir = source_uri.removeprefix("viking://").rsplit("/", 1)[0] + return "viking://" + posixpath.normpath(f"{source_dir}/{target}") + + +class _AsyncRWLock: + """A simple asyncio reader/writer lock. + + - Multiple readers may hold the lock concurrently. + - Writers get exclusive access; new readers are blocked while a writer is waiting + to avoid writer starvation. + - Not reentrant. + """ + + def __init__(self) -> None: + self._readers = 0 + self._writers_waiting = 0 + self._writing = False + self._lock = asyncio.Lock() + self._readers_ok = asyncio.Condition(self._lock) + self._writer_ok = asyncio.Condition(self._lock) + + def reader(self) -> "_ReaderCtx": + return _ReaderCtx(self) + + def writer(self) -> "_WriterCtx": + return _WriterCtx(self) + + async def _acquire_reader(self) -> None: + async with self._lock: + while self._writing or self._writers_waiting > 0: + await self._readers_ok.wait() + self._readers += 1 + + async def _release_reader(self) -> None: + async with self._lock: + self._readers -= 1 + if self._readers == 0: + self._writer_ok.notify() + + async def _acquire_writer(self) -> None: + async with self._lock: + self._writers_waiting += 1 + try: + while self._readers > 0 or self._writing: + await self._writer_ok.wait() + self._writing = True + finally: + self._writers_waiting -= 1 + + async def _release_writer(self) -> None: + async with self._lock: + self._writing = False + if self._writers_waiting > 0: + self._writer_ok.notify() + else: + self._readers_ok.notify_all() + + +class _ReaderCtx: + __slots__ = ("_rw",) + + def __init__(self, rw: _AsyncRWLock) -> None: + self._rw = rw + + async def __aenter__(self) -> "_ReaderCtx": + await self._rw._acquire_reader() + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + await self._rw._release_reader() + + +class _WriterCtx: + __slots__ = ("_rw",) + + def __init__(self, rw: _AsyncRWLock) -> None: + self._rw = rw + + async def __aenter__(self) -> "_WriterCtx": + await self._rw._acquire_writer() + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + await self._rw._release_writer() + + +@dataclass(slots=True) +class VikingBotTau2RolloutExecutor: + """Execute tau2 cases with VikingBot agent loop and tau2 tools.""" + + config_path: str | None = None + concurrency: int = 20 + keep_default_tools: bool = True + max_iterations: int = 30 + log_timings: bool = True + rollout_language: str = "default" + + def __post_init__(self) -> None: + if self.rollout_language not in {"default", "zh"}: + raise ValueError("rollout_language must be 'default' or 'zh'") + + async def execute( + self, + cases: list[Case], + policy_set: ExperienceSet, + context: ExecutionContext, + ) -> list[Rollout]: + del policy_set + if self.concurrency <= 0: + raise ValueError("concurrency must be > 0") + semaphore = asyncio.Semaphore(self.concurrency) + + async def run_one(case: Case) -> Rollout: + async with semaphore: + return await self._execute_one(case, context) + + return list(await asyncio.gather(*(run_one(case) for case in cases))) + + async def _execute_one(self, case: Case, context: ExecutionContext) -> Rollout: + return await self._execute_one_async(case, context) + + async def _execute_one_async(self, case: Case, context: ExecutionContext) -> Rollout: + domain = str(case.input["domain"]) + task_id = str(case.input["task_id"]) + task_no = int(case.input["task_no"]) + data_split = str(case.input["data_split"]) + data_root = case.input.get("data_root") + trial = _case_trial(case) + + timings = _RolloutTiming(case=case.name, enabled=self.log_timings) + total_started_at = time.perf_counter() + + stage_started_at = time.perf_counter() + Tau2BenchToolProvider = _tool_provider_cls() + provider = Tau2BenchToolProvider(domain, task_id, data_root=data_root) + await asyncio.to_thread(provider.reset) + timings.record("provider_reset", stage_started_at) + + stage_started_at = time.perf_counter() + agent = await asyncio.to_thread( + _build_agent, + self.config_path, + max_iterations=self.max_iterations, + ) + timings.record("build_agent", stage_started_at) + + stage_started_at = time.perf_counter() + _configure_tools( + agent, + provider, + keep_default_tools=self.keep_default_tools, + record_tool_timing=timings.record_tool, + task_id=task_id, + task_no=task_no, + data_split=data_split, + ) + timings.record("configure_tools", stage_started_at) + + stage_started_at = time.perf_counter() + system_prompt = _build_system_prompt( + provider.policy, + keep_default_tools=self.keep_default_tools, + rollout_language=self.rollout_language, + ) + user_prompt = provider.user_query + SessionKey = _vikingbot_imports()["SessionKey"] + trial_suffix = "" if trial is None else f"_r{int(trial)}" + stage = _safe_session_fragment(str(context.metadata.get("stage") or "rollout")) + session_key = SessionKey( + type="cli", + channel_id="tau2", + chat_id=f"tau2_{stage}_{data_split}_{task_no}{trial_suffix}", + ) + timings.record("prepare_prompt", stage_started_at) + + ( + final_content, + final_reasoning_content, + tools_used, + token_usage, + iteration, + memory_content, + experience_reminder, + experience_loader_skill, + ) = await _run_agent( + agent=agent, + system_prompt=system_prompt, + user_prompt=user_prompt, + session_key=session_key, + sender_id="tau2_user", + keep_default_tools=self.keep_default_tools, + timings=timings, + case_lookup=_tau2_case_lookup(case), + ) + + reward = None + evaluation_result = None + stage_started_at = time.perf_counter() + if provider.env is not None: + try: + # Customer-facing content should be sent before `done`; do not append + # the post-done final response to tau2's simulator/evaluator. + reward, evaluation_result = await asyncio.to_thread(provider.env._get_reward) + reward = _to_jsonable(reward) + evaluation_result = _to_jsonable(evaluation_result) + except Exception as exc: + logger.exception( + "tau2 reward calculation failed case=%s domain=%s task_id=%s", + case.name, + domain, + task_id, + ) + evaluation_result = {"error": str(exc), "type": type(exc).__name__} + timings.record("reward", stage_started_at) + + stage_started_at = time.perf_counter() + rollout = Rollout( + case=case, + messages=_build_rollout_messages( + system_prompt=system_prompt, + user_prompt=user_prompt, + tools_used=tools_used, + final_content=final_content, + evaluation_result=evaluation_result, + reward=reward, + experience_reminder=experience_reminder, + artifact_created_at=_tau2_policy_current_time_iso(system_prompt), + ), + policy_snapshot_id=context.policy_snapshot_id, + evaluation=_tau2_evaluation(reward=reward, evaluation_result=evaluation_result), + metadata={ + "domain": domain, + "data_split": data_split, + "task_no": task_no, + "task_id": task_id, + "eval_trial": case.input.get("eval_trial"), + "eval_trial_count": case.input.get("eval_trial_count"), + "train_trial": case.input.get("train_trial"), + "train_trial_count": case.input.get("train_trial_count"), + "original_case_name": case.input.get("original_case_name"), + "reward": reward, + "evaluation_result": evaluation_result, + "tools_used": tools_used, + "token_usage": token_usage, + "iterations": iteration, + "memory": memory_content, + "system_prompt": system_prompt, + "business_current_time": _tau2_policy_current_time_iso(system_prompt), + "user_prompt": user_prompt, + "final_content": final_content, + "final_reasoning_content": final_reasoning_content, + "keep_default_tools": self.keep_default_tools, + "ov_tools_enable": False, + "experience_recall_enable": self.keep_default_tools, + "experience_loader_skill": experience_loader_skill, + "execution_metadata": dict(context.metadata), + }, + ) + timings.record("build_rollout", stage_started_at) + timings.log_summary( + total_ms=_elapsed_ms(total_started_at), + task_id=task_id, + task_no=task_no, + data_split=data_split, + iterations=iteration, + reward=reward, + message_count=len(rollout.messages), + ) + rollout.metadata["timing_ms"] = timings.snapshot( + total_ms=_elapsed_ms(total_started_at), + iterations=iteration, + ) + return rollout + + +def _tau2_case_lookup(case: Case) -> dict[str, Any]: + case_input = dict(case.input or {}) + domain = case_input.get("domain") + split = case_input.get("split") + task_id = case_input.get("task_id") + # Trial cases append a trial suffix to Case.task_signature; case memories are + # keyed by the stable tau2题目 identity, so use the base signature. + task_signature = ( + f"tau2:{domain}:{split}:{task_id}" + if domain is not None and split is not None and task_id is not None + else case.task_signature + ) + data_split = case_input.get("data_split") + task_no = case_input.get("task_no") + case_name = case_input.get("original_case_name") or case.name + case_names = [case_name] + if data_split is not None and task_no is not None: + case_names.append(f"tau2_{data_split}_{task_no}") + return { + "benchmark": "tau2", + "strict": True, + "case_names": case_names, + "domain": domain, + "split": split, + "data_split": data_split, + "task_no": task_no, + "task_id": task_id, + "case_name": case_name, + "task_signature": task_signature, + "original_case_name": case_input.get("original_case_name"), + "expected_fields": { + "input.domain": domain, + "input.split": split, + "input.data_split": data_split, + "input.task_no": task_no, + "input.task_id": task_id, + }, + } + + +def _append_final_answer_for_tau2_evaluation(provider_env: Any, final_content: str | None) -> None: + if not final_content or not str(final_content).strip(): + return + target = getattr(provider_env, "_impl", provider_env) + append_message = getattr(target, "append_agent_message", None) + if callable(append_message): + append_message(str(final_content)) + + +# Tokens tau2's user simulator emits to signal that the conversation should end. +_TAU2_USER_STOP_TOKENS = ("###STOP###",) +_TAU2_USER_TRANSFER_TOKENS = ("###TRANSFER###",) + + +def _tau2_user_reply_terminates(reply: Any) -> bool: + text = str(reply or "") + return any(tok in text for tok in _TAU2_USER_STOP_TOKENS + _TAU2_USER_TRANSFER_TOKENS) + + +def _make_tau2_plain_text_router(*, publish_events: bool, bus: Any, session_key: Any): + """Build an `on_plain_text` callback that forwards assistant text via communicate_with_user. + + In tau2 bench, plain assistant text is semantically equivalent to calling + `communicate_with_user`: both should be delivered to the user simulator so the + simulated user can reply and the conversation can continue. This router is owned by + the tau2 executor so vikingbot's generic AgentLoop stays benchmark-agnostic. + """ + imports = _vikingbot_imports() + PlainTextContext = imports["_PlainTextContext"] + PlainTextDelivered = imports["_PlainTextDelivered"] + PlainTextFinal = imports["_PlainTextFinal"] + OutboundMsgType = imports["MessageBus"] # only used for type/attr access + del OutboundMsgType + + async def _route( + ctx: PlainTextContext, # type: ignore[valid-type] + ): + text = ctx.text + # If the assistant text itself contains STOP (unlikely in tau2), treat as final. + if any(tok in text for tok in _TAU2_USER_STOP_TOKENS): + return PlainTextFinal(content=text) + if not ctx.tools.has("communicate_with_user"): + return PlainTextFinal(content=text) + + messages = list(ctx.messages) + # Record the assistant text using the same dict shape vikingbot uses elsewhere. + assistant_entry: dict[str, Any] = {"role": "assistant", "content": text} + if ctx.reasoning_content: + assistant_entry["reasoning_content"] = ctx.reasoning_content + messages.append(assistant_entry) + from vikingbot.utils.helpers import cal_str_tokens as _cal + + started_at = time.perf_counter() + user_reply = await ctx.tools.execute( + "communicate_with_user", + {"content": text}, + session_key=ctx.session_key, + sandbox_manager=ctx.sandbox_manager, + sender_id=ctx.sender_id, + memory_peer_ids=ctx.memory_peer_ids, + memory_owner_user_ids=ctx.memory_owner_user_ids, + openviking_connection=ctx.openviking_connection, + ) + duration_ms = (time.perf_counter() - started_at) * 1000 + args_str = json.dumps({"content": text}, ensure_ascii=False) + logger.info("[TAU2_PLAIN_TEXT]: routed assistant text through communicate_with_user") + logger.info(f"[TOOL_CALL]: communicate_with_user({args_str[:200]})") + logger.info(f"[RESULT]: {str(user_reply)[:600]}") + if publish_events: + from vikingbot.bus.events import OutboundEventType, OutboundMessage + + await bus.publish_outbound( + OutboundMessage( + session_key=session_key, + content=f"communicate_with_user({args_str})", + event_type=OutboundEventType.TOOL_CALL, + ) + ) + await bus.publish_outbound( + OutboundMessage( + session_key=session_key, + content=str(user_reply), + event_type=OutboundEventType.TOOL_RESULT, + ) + ) + tools_used = [ + { + "tool_name": "communicate_with_user", + "args": args_str, + "result": user_reply, + "duration": duration_ms, + "execute_success": _viking_is_tool_result_success(user_reply), + "input_token": 0, + "output_token": _cal(user_reply, text_type="mixed"), + "auto": True, + } + ] + messages.append({"role": "user", "content": str(user_reply)}) + terminates = _tau2_user_reply_terminates(user_reply) + return PlainTextDelivered( + messages=messages, + tools_used=tools_used, + user_terminates=terminates, + ) + + return _route + + +def _build_agent(config_path: str | None, *, max_iterations: int): + imports = _vikingbot_imports() + config = imports["ensure_config"](Path(config_path).expanduser() if config_path else None) + imports["_init_bot_data"](config) + bus = imports["MessageBus"]() + session_manager = imports["SessionManager"](config.bot_data_path) + sandbox_parent_path = config.workspace_path + source_workspace_path = imports["get_source_workspace_path"]() + sandbox_manager = imports["SandboxManager"](config, sandbox_parent_path, source_workspace_path) + provider = imports["_make_provider"](config) + return imports["AgentLoop"]( + bus=bus, + provider=provider, + workspace=config.workspace_path, + model=config.agents.model, + max_iterations=max_iterations, + memory_window=config.agents.memory_window, + brave_api_key=config.tools.web.search.api_key or None, + exa_api_key=None, + gen_image_model=config.agents.gen_image_model, + exec_config=config.tools.exec, + cron_service=None, + session_manager=session_manager, + sandbox_manager=sandbox_manager, + config=config, + eval=True, + mcp_servers=None, + ) + + +def _configure_tools( + agent: Any, + provider: Any, + *, + keep_default_tools: bool, + record_tool_timing: Callable[[str, float], None] | None = None, + task_id: str | None = None, + task_no: int | None = None, + data_split: str | None = None, +) -> None: + # Tau2 rollout may keep generic VikingBot tools, but OpenViking access is + # restricted to automatic experience recall during prompt construction. + # No openviking_* tool should be callable by the agent. + del keep_default_tools + for tool_name in list(agent.tools.tool_names): + if str(tool_name).startswith("openviking_"): + agent.tools.unregister(tool_name) + agent.tools.register(_make_search_experience_tool()) + agent.tools.register(_make_read_experience_tool()) + tool_lock = _AsyncRWLock() + write_tool_names = _classify_write_tools(provider) + for schema in provider.list_openai_tools(): + fn_name = str((schema.get("function") or {}).get("name") or "") + agent.tools.register( + _make_tau2_tool( + schema, + provider, + tool_lock=tool_lock, + is_write_tool=fn_name in write_tool_names, + record_tool_timing=record_tool_timing, + ) + ) + + +def _classify_write_tools(provider: Any) -> set[str]: + """Classify which tau2 tools mutate environment state. + + Pure read/lookup tools can run in parallel within a single rollout; state-mutating + tools (book/update/cancel/etc.) plus communicate_with_user and ``done`` must run + exclusively because they advance the user simulator and tau2 DB state. + """ + write_names: set[str] = {"communicate_with_user", "done"} + + # 1) Introspect the underlying tau2 ToolKit: tau2 marks tools with __tool_type__ and + # __mutates_state__. Prefer this when available (covers both gym and native envs). + env = getattr(provider, "env", None) + inner = getattr(env, "_impl", None) if env is not None else None + inner_env = getattr(inner, "env", None) if inner is not None else None + for toolkit_attr in ("tools", "user_tools"): + toolkit = getattr(inner_env, toolkit_attr, None) if inner_env is not None else None + if toolkit is None: + continue + get_tools_fn = getattr(toolkit, "get_tools", None) + tool_type_fn = getattr(toolkit, "tool_type", None) + mutates_fn = getattr(toolkit, "tool_mutates_state", None) + try: + tools_dict = get_tools_fn() if callable(get_tools_fn) else None + except Exception: + tools_dict = None + if isinstance(tools_dict, dict): + for name, tool_fn in tools_dict.items(): + mutates = getattr(tool_fn, "__mutates_state__", None) + tool_type = getattr(tool_fn, "__tool_type__", None) + if mutates is None and mutates_fn is not None: + try: + mutates = mutates_fn(name) + except Exception: + mutates = None + if tool_type is None and tool_type_fn is not None: + try: + tool_type = tool_type_fn(name) + except Exception: + tool_type = None + is_write = mutates is True or str(tool_type) in { + "write", + "ToolType.WRITE", + "ToolType.WRITE.value", + } + if is_write: + write_names.add(str(name)) + + # 2) Heuristic fallback for tools not introspected above: any tool not starting + # with a read-y prefix is assumed to be a writer. This is conservative (pessimistic + # about parallelism) rather than risking races on stateful tools. + try: + schemas = list(provider.list_openai_tools() or []) + except Exception: + schemas = [] + _READ_PREFIXES = ( + "get_", + "search_", + "list_", + "find_", + "retrieve_", + "lookup_", + "check_", + "view_", + "describe_", + "think", + "summary", + ) + for schema in schemas: + fn = schema.get("function") or {} + name = str(fn.get("name") or "") + if not name or name in write_names: + continue + if not any(name.startswith(p) for p in _READ_PREFIXES): + write_names.add(name) + return write_names + + +def _build_system_prompt(policy: str, *, keep_default_tools: bool, rollout_language: str) -> str: + del keep_default_tools + instructions = [] + if policy: + instructions.append(policy) + instructions.append("Use the provided tools to interact with the environment.") + instructions.append( + "Before taking task actions, you MUST use the required `experience_loader` skill. " + "It explains how to search OpenViking case memories with the `search_experience` tool, return linked experience URIs, and read selected experiences using the `read_experience` tool." + ) + instructions.append( + "Loaded experiences are guidance from prior training runs. " + "Use them only when their situation and applicability boundaries match the current " + "task; current policy, current tool results, and current user facts override prior " + "experience." + ) + if rollout_language == "zh": + instructions.append( + "Communicate with the user and write the final response in Chinese. " + "Do not translate tool names, identifiers, JSON field names, reservation IDs, " + "flight numbers, or other structured values used by tools." + ) + instructions.append( + "If you need to communicate with the user, you MUST call tool `communicate_with_user`." + ) + instructions.append( + "When communicating numbers, prices, reservation IDs, flight numbers, airport codes, " + "dates, names, or other values from tool results, include the exact original value " + "verbatim even if the surrounding response is in another language." + ) + instructions.append( + "When the task is finished or terminated, send any final customer-facing message " + "through `communicate_with_user` before calling `done`. After `done`, do not call " + "any more tools and do not emit extra ending content." + ) + return "\n".join(instructions) + + +EXPERIENCE_LOADER_TEMPLATE_DIR = Path(__file__).resolve().parent / "experience_loader_template" +EXPERIENCE_LOADER_SKILL_PATH = "skills/experience_loader/SKILL.md" + + +async def _prepare_experience_loader_skill( + *, + agent: Any, + session_key: Any, +) -> Any: + """Install the generic experience_loader skill into the rollout sandbox. + + The loader does not contain per-task memory. It instructs the LLM to use the + tau2-only `search_experience` and `read_experience` tools to search case-linked experiences and load selected experience memories. + """ + + imports = _vikingbot_imports() + sandbox_manager = getattr(agent, "sandbox_manager", None) + workspace_path = ( + sandbox_manager.get_workspace_path(session_key) + if sandbox_manager + else agent.context.workspace + ) + skill_content = _read_experience_loader_template_file("SKILL.md") + if sandbox_manager: + try: + sandbox = await sandbox_manager.get_sandbox(session_key) + await sandbox.write_file(EXPERIENCE_LOADER_SKILL_PATH, skill_content) + except Exception as exc: + logger.warning("failed to write experience_loader skill to sandbox: %s", exc) + _write_experience_loader_files( + workspace_path=workspace_path, + skill_content=skill_content, + ) + else: + _write_experience_loader_files( + workspace_path=workspace_path, + skill_content=skill_content, + ) + + context_builder = imports["ContextBuilder"]( + workspace_path, + sandbox_manager=sandbox_manager, + eval=True, + ) + context_builder.latest_experience_loader_skill_content = skill_content + return context_builder + + +def _read_experience_loader_template_file(relative_path: str) -> str: + return (EXPERIENCE_LOADER_TEMPLATE_DIR / relative_path).read_text(encoding="utf-8") + + +def _write_experience_loader_files( + *, + workspace_path: Path, + skill_content: str, +) -> None: + skill_dir = workspace_path / "skills" / "experience_loader" + skill_dir.mkdir(parents=True, exist_ok=True) + skill_dir.joinpath("SKILL.md").write_text(skill_content, encoding="utf-8") + + +async def _execute_required_experience_loader_read( + *, + agent: Any, + messages: list[dict[str, Any]], + session_key: Any, + sender_id: str, +) -> dict[str, Any]: + """Force-load the generic experience_loader skill before task actions.""" + + path = EXPERIENCE_LOADER_SKILL_PATH + tool_id = "required-experience-loader-skill-read" + messages.append( + { + "role": "assistant", + "content": "Reading required experience_loader skill before task actions.", + "tool_calls": [ + { + "id": tool_id, + "type": "function", + "function": { + "name": "read_file", + "arguments": json.dumps({"path": path}, ensure_ascii=False), + }, + } + ], + } + ) + started_at = time.perf_counter() + result = await agent.tools.execute( + "read_file", + {"path": path}, + session_key=session_key, + sandbox_manager=agent.sandbox_manager, + sender_id=sender_id, + ) + duration_ms = _elapsed_ms(started_at) + messages.append( + { + "role": "tool", + "tool_call_id": tool_id, + "name": "read_file", + "content": result, + } + ) + execute_success = not (isinstance(result, str) and result.lstrip().startswith("Error")) + if not execute_success: + logger.warning("required experience_loader skill read failed: %s", str(result)[:300]) + return { + "tool_name": "read_file", + "args": json.dumps({"path": path}, ensure_ascii=False), + "result": result, + "duration": duration_ms, + "execute_success": execute_success, + "input_token": 0, + "output_token": 0, + "auto": True, + "required_skill": "experience_loader", + } + + +async def _run_agent( + *, + agent: Any, + system_prompt: str, + user_prompt: str, + session_key: Any, + sender_id: str, + keep_default_tools: bool, + timings: "_RolloutTiming | None" = None, + case_lookup: dict[str, Any] | None = None, +): + stage_started_at = time.perf_counter() + message_context = agent.context + del case_lookup + message_context = await _prepare_experience_loader_skill( + agent=agent, + session_key=session_key, + ) + experience_loader_skill = getattr( + message_context, + "latest_experience_loader_skill_content", + None, + ) + messages = await message_context.build_messages( + history=[], + current_message=user_prompt, + session_key=session_key, + ov_tools_enable=False, + experience_recall_enable=False, + media=None, + profile_user_list=[], + ) + if timings is not None: + timings.record("build_messages", stage_started_at) + if system_prompt: + messages.insert(1, {"role": "system", "content": system_prompt}) + _override_vikingbot_current_time_messages( + messages, + business_current_time=_tau2_policy_current_time_display(system_prompt), + ) + user_memory = None + experience_reminder_text = None # 完整的 [Experience Reminder] 消息文本(用于 messages.json) + for msg in messages: + content = msg.get("content", "") if isinstance(msg, dict) else "" + if not isinstance(content, str): + continue + # Experience Reminder (经验记忆) - role=user, starts with [Experience Reminder] + if "[Experience Reminder]" in content and "## Relevant Agent Experience" in content: + experience_reminder_text = content + continue + # User memory (用户记忆) - starts with "## Current Session" + if content.startswith("## Current Session"): + user_memory = _extract_memory_content(content) + + # 合并用户记忆 + 经验记忆正文,去重 + exp_content = ( + _extract_experience_content(experience_reminder_text) if experience_reminder_text else None + ) + memory_content = _merge_memories(user_memory, exp_content) + stage_started_at = time.perf_counter() + required_skill_tool = None + if experience_loader_skill and experience_loader_skill.strip(): + required_skill_tool = await _execute_required_experience_loader_read( + agent=agent, + messages=messages, + session_key=session_key, + sender_id=sender_id, + ) + plain_text_router = _make_tau2_plain_text_router( + publish_events=False, + bus=getattr(agent, "bus", None), + session_key=session_key, + ) + result = await agent._run_agent_loop( + messages=messages, + session_key=session_key, + publish_events=False, + sender_id=sender_id, + ov_tools_enable=False, + stop_tool_names=["done"], + on_plain_text=plain_text_router, + ) + if timings is not None: + timings.record("agent_loop", stage_started_at) + final_content, final_reasoning_content, tools_used, token_usage, iteration = result + if required_skill_tool is not None: + tools_used = [required_skill_tool, *tools_used] + case_memory_context = _case_memory_context_from_tools(tools_used) + memory_content = _merge_memories(memory_content, case_memory_context) + if _last_tool_name(tools_used) == "done": + final_content = None + final_reasoning_content = None + return ( + final_content, + final_reasoning_content, + tools_used, + token_usage, + iteration, + memory_content, + experience_reminder_text, + experience_loader_skill, + ) + + +def _override_vikingbot_current_time_messages( + messages: list[dict[str, Any]], + *, + business_current_time: str | None, +) -> None: + """Replace VikingBot's wall-clock prompt time with tau2's business time. + + VikingBot's generic context builder includes ``## Current Time: `` in the user-memory wrapper. For tau2, the domain policy owns the + business clock; leaving the host clock in the prompt can make the agent + interpret unqualified dates against the run date. + """ + if not business_current_time: + return + replacement = f"## Current Time: {business_current_time}" + for msg in messages: + content = msg.get("content") if isinstance(msg, dict) else None + if not isinstance(content, str) or "## Current Time:" not in content: + continue + msg["content"] = re.sub( + r"(?m)^## Current Time: .*$", + replacement, + content, + count=1, + ) + + +@dataclass(slots=True) +class _RolloutTiming: + case: str + enabled: bool + stages: dict[str, float] = field(default_factory=dict) + tool_durations: list[tuple[str, float]] = field(default_factory=list) + + def record(self, stage: str, started_at: float) -> None: + if self.enabled: + self.stages[stage] = _elapsed_ms(started_at) + + def record_tool(self, tool_name: str, duration_ms: float) -> None: + if self.enabled: + self.tool_durations.append((tool_name, duration_ms)) + + def snapshot(self, *, total_ms: float, iterations: int | None) -> dict[str, Any]: + """Return a JSON-serializable timing breakdown for rollout.metadata.""" + tool_total_ms = sum(duration for _, duration in self.tool_durations) + tool_counts: dict[str, int] = {} + tool_total_by_name: dict[str, float] = {} + tool_max_by_name: dict[str, float] = {} + for name, duration in self.tool_durations: + tool_counts[name] = tool_counts.get(name, 0) + 1 + tool_total_by_name[name] = tool_total_by_name.get(name, 0.0) + duration + cur = tool_max_by_name.get(name, 0.0) + if duration > cur: + tool_max_by_name[name] = duration + tools_by_name = { + name: { + "count": tool_counts[name], + "total_ms": round(tool_total_by_name[name], 2), + "avg_ms": round(tool_total_by_name[name] / tool_counts[name], 2), + "max_ms": round(tool_max_by_name[name], 2), + } + for name in tool_counts + } + slowest = max(self.tool_durations, key=lambda item: item[1], default=None) + return { + "total_ms": round(total_ms, 2), + "iterations": iterations, + "stages_ms": {k: round(v, 2) for k, v in self.stages.items()}, + "tool_count": len(self.tool_durations), + "tool_total_ms": round(tool_total_ms, 2), + "slowest_tool": ( + {"name": slowest[0], "duration_ms": round(slowest[1], 2)} + if slowest is not None + else None + ), + "tools_by_name": tools_by_name, + } + + def log_summary(self, *, total_ms: float, **metadata: Any) -> None: + if not self.enabled: + return + tool_total_ms = sum(duration for _, duration in self.tool_durations) + slowest_tool = max(self.tool_durations, key=lambda item: item[1], default=None) + logger.info( + "tau2 rollout timing case=%s total_ms=%.1f stages=%s tool_count=%d " + "tool_total_ms=%.1f slowest_tool=%s metadata=%s", + self.case, + total_ms, + _format_stage_timings(self.stages), + len(self.tool_durations), + tool_total_ms, + _format_tool_timing(slowest_tool), + metadata, + ) + + +def _elapsed_ms(started_at: float) -> float: + return (time.perf_counter() - started_at) * 1000.0 + + +def _format_stage_timings(stages: dict[str, float]) -> str: + return ",".join(f"{stage}:{duration_ms:.1f}" for stage, duration_ms in stages.items()) + + +def _format_tool_timing(item: tuple[str, float] | None) -> str | None: + if item is None: + return None + tool_name, duration_ms = item + return f"{tool_name}:{duration_ms:.1f}" + + +def _safe_session_fragment(value: str) -> str: + return "".join(ch if ch.isalnum() or ch in "_.-" else "_" for ch in value)[:80] or "rollout" + + +MEMORY_PROMPT_PREFIX = "## Current Session\nChannel: cli\n\n---\n\n" +MEMORY_PROMPT_SUFFIX = ( + "---\n\nReply in the same language as the user's query, ignoring the language of " + "the reference materials. User's query:" +) + + +def _extract_memory_content(content: str) -> str | None: + start = content.find(MEMORY_PROMPT_PREFIX) + end = content.rfind(MEMORY_PROMPT_SUFFIX) + if start == -1 or end == -1: + return None + start += len(MEMORY_PROMPT_PREFIX) + if start > end: + return None + return content[start:end] + + +def _extract_experience_content(content: str) -> str | None: + """从 Experience Reminder 消息中提取经验记忆正文。""" + prefix = "[Experience Reminder]\n## Relevant Agent Experience\n" + start = content.find(prefix) + if start == -1: + return None + start += len(prefix) + return content[start:].strip() or None + + +def _case_memory_context_from_tools(tools_used: list[dict] | None) -> str: + blocks: list[str] = [] + for tool in tools_used or []: + if not isinstance(tool, dict) or tool.get("tool_name") != "read_experience": + continue + result = str(tool.get("result") or "").strip() + if not result: + continue + args = tool.get("args") + blocks.append( + "\n".join( + [ + "## Loaded Experience", + "", + "Tool: `read_experience`", + "", + "Args:", + "```json", + str(args or "{}"), + "```", + "", + result, + ] + ) + ) + if not blocks: + return "" + return "# Experience Loader Context\n\n" + "\n\n---\n\n".join(blocks) + + +def _merge_memories(user_memory: str | None, exp_memory: str | None) -> str | None: + """合并用户记忆和经验记忆,去重。 + + 两者都为 None 时返回 None;只有一个时直接返回它;都有时拼接并标记类型。 + """ + parts: list[str] = [] + if user_memory and user_memory.strip(): + parts.append(f"## User Memories\n{user_memory.strip()}") + if exp_memory and exp_memory.strip(): + parts.append(f"## Experience Memories\n{exp_memory.strip()}") + if not parts: + return None + return "\n\n---\n\n".join(parts) + + +def _build_rollout_messages( + *, + system_prompt: str, + user_prompt: str, + tools_used: Any, + final_content: str | None, + evaluation_result: Any, + reward: Any, + experience_reminder: str | None = None, + artifact_created_at: str | None = None, +) -> list[Message]: + messages = [ + _metadata_message( + "tau2-system", + f"system:\n{system_prompt}", + created_at=artifact_created_at, + ), + ] + # Experience Reminder 放在 system 之后、user 之前,与 agent 实际看到的顺序一致 + if experience_reminder: + messages.append( + _message("tau2-experience", "user", experience_reminder, created_at=artifact_created_at) + ) + messages.append(_message("tau2-user", "user", user_prompt, created_at=artifact_created_at)) + if isinstance(tools_used, list): + for idx, tool_info in enumerate(tools_used): + if not isinstance(tool_info, dict): + continue + tool_name = str(tool_info.get("tool_name") or "unknown") + if not tool_name or tool_name == "unknown" and not tool_info.get("result"): + continue + args = tool_info.get("args", "") + tool_input = _as_tool_input(args) + result = tool_info.get("result") + has_result = result is not None + if _is_communicate_with_user(tool_name): + assistant_text = _communicate_text_from_tool_input(tool_input) + if assistant_text.strip(): + messages.append( + _message( + f"tau2-communicate-assistant-{idx}", + "assistant", + assistant_text, + created_at=artifact_created_at, + ) + ) + if has_result: + user_text = _stringify(result) + if user_text.strip(): + messages.append( + _message( + f"tau2-communicate-user-{idx}", + "user", + user_text, + created_at=artifact_created_at, + ) + ) + continue + messages.append( + Message( + id=f"tau2-tool-{idx}", + role="user" if has_result else "assistant", + parts=[ + ToolPart( + tool_id=f"tau2-tool-{idx}", + tool_name=tool_name, + tool_input=tool_input, + tool_output=_stringify(result) if has_result else "", + tool_status="completed" if has_result else "running", + ) + ], + created_at=artifact_created_at, + ) + ) + if final_content and str(final_content).strip(): + messages.append( + _message("tau2-final", "assistant", str(final_content), created_at=artifact_created_at) + ) + reward_jsonable = _to_jsonable(reward) + evaluation_jsonable = _to_jsonable(evaluation_result) + success = reward_jsonable == 1 or reward_jsonable == 1.0 + messages.append( + _message( + "tau2-reward", + "user", + f"task_success: {success}\ntask_reward: {reward_jsonable}\n" + f"evaluation report: {_stringify(evaluation_jsonable)}", + created_at=artifact_created_at, + ) + ) + return messages + + +def _tau2_evaluation(*, reward: Any, evaluation_result: Any) -> RubricEvaluation: + return _tau2_evaluation_helper( + reward=reward, evaluation_result=evaluation_result, source="tau2_executor" + ) + + +def _last_tool_name(tools_used: Any) -> str: + if not isinstance(tools_used, list) or not tools_used: + return "" + last = tools_used[-1] + if not isinstance(last, dict): + return "" + return str(last.get("tool_name") or "") + + +# Backwards-compatible alias for existing imports. +Tau2RolloutExecutor = VikingBotTau2RolloutExecutor diff --git a/benchmark/tau2/train/run_batch_train_eval.sh b/benchmark/tau2/train/run_batch_train_eval.sh new file mode 100755 index 0000000000..85f4eda5ca --- /dev/null +++ b/benchmark/tau2/train/run_batch_train_eval.sh @@ -0,0 +1,20 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Tau2 convenience launcher for the generic OpenViking session/train batch pipeline. +# Start the tau2 runtime service first: +# bash benchmark/tau2/train/run_service.sh --host 127.0.0.1 --port 1944 +# Pass --rollout-backend native|vikingbot to override per run (default: vikingbot). + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +TAU2_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)" +REPO_ROOT="$(cd "${TAU2_DIR}/../.." && pwd)" + +exec "${REPO_ROOT}/openviking/session/train/run_batch_train_eval.sh" \ + --dataset tau2 \ + --domain airline \ + --eval-each-epoch \ + --concurrency 200 \ + --commit-concurrency 200 \ + --benchmark-service-url "${BENCHMARK_SERVICE_URL:-http://127.0.0.1:1944}" \ + "$@" diff --git a/benchmark/tau2/train/run_service.sh b/benchmark/tau2/train/run_service.sh new file mode 100755 index 0000000000..ec263782d5 --- /dev/null +++ b/benchmark/tau2/train/run_service.sh @@ -0,0 +1,224 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +TAU2_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)" +REPO_ROOT="$(cd "${TAU2_DIR}/../.." && pwd)" + +load_user_env_file() { + local env_file="${OPENVIKING_ENV_FILE:-${HOME}/.openviking_benchmark_env}" + if [[ -z "${env_file}" || ! -f "${env_file}" ]]; then + return + fi + + local -a preserved_env=() + local entry + while IFS= read -r -d '' entry; do + if [[ "${entry}" != *= ]]; then + preserved_env+=("${entry}") + fi + done < <(env -0) + + echo "[tau2-service] loading env file ${env_file}" + set +u + set -a + # shellcheck source=/dev/null + source "${env_file}" + set +a + set -euo pipefail + + for entry in "${preserved_env[@]}"; do + export "${entry}" + done +} + +load_user_env_file + +PYTHON_BIN="${PYTHON_BIN:-python}" +HOST="127.0.0.1" +PORT="1944" +DATA_ROOT="${TAU2_DATA_ROOT:-}" +CONFIG="${OPENVIKING_CONFIG_FILE:-${HOME}/.openviking/ov.conf}" +KILL_EXISTING=1 +ROLLOUT_LANGUAGE="default" +ROLLOUT_BACKEND="${TAU2_ROLLOUT_BACKEND:-vikingbot}" +NATIVE_THREAD_WORKERS="${TAU2_NATIVE_THREAD_WORKERS:-128}" +MAX_ROLLOUT_CONCURRENCY="${TAU2_MAX_ROLLOUT_CONCURRENCY:-200}" +ROLLOUT_THREAD_WORKERS="${TAU2_ROLLOUT_THREAD_WORKERS:-200}" +REPAIR_VIKINGBOT_GYM="${TAU2_REPAIR_VIKINGBOT_GYM:-1}" + +while [[ $# -gt 0 ]]; do + case "$1" in + --host) HOST="$2"; shift 2 ;; + --port) PORT="$2"; shift 2 ;; + --data-root) DATA_ROOT="$2"; shift 2 ;; + --config) CONFIG="$2"; shift 2 ;; + --rollout-language) ROLLOUT_LANGUAGE="$2"; shift 2 ;; + --rollout-backend) ROLLOUT_BACKEND="$2"; shift 2 ;; + --native-thread-workers) NATIVE_THREAD_WORKERS="$2"; shift 2 ;; + --max-rollout-concurrency) MAX_ROLLOUT_CONCURRENCY="$2"; shift 2 ;; + --rollout-thread-workers) ROLLOUT_THREAD_WORKERS="$2"; shift 2 ;; + --repair-vikingbot-gym) REPAIR_VIKINGBOT_GYM=1; shift 1 ;; + --no-repair-vikingbot-gym) REPAIR_VIKINGBOT_GYM=0; shift 1 ;; + --no-kill-existing) KILL_EXISTING=0; shift 1 ;; + -h|--help) + cat <<'EOF' +Usage: + bash benchmark/tau2/train/run_service.sh [--host 127.0.0.1] [--port 1944] + +Options: + --data-root PATH tau2-bench data/tau2 root. Default auto-detect/TAU2_DATA_ROOT + --config PATH ov.conf for VikingBot/OpenViking access. Default ~/.openviking/ov.conf + --rollout-language default|zh + Rollout response language. Use zh for Chinese user-facing replies. + --rollout-backend native|vikingbot + Rollout implementation backend. Default: vikingbot. + --native-thread-workers N + Default thread pool workers for native rollout. Default: 128. + --max-rollout-concurrency N + Maximum concurrent rollout executions hosted by the service. + Default: 200. + --rollout-thread-workers N + Worker threads used to host rollouts off the uvicorn event loop. + Default: 200. Use 0 to disable threaded hosting. + --repair-vikingbot-gym + If --rollout-backend=vikingbot and tau2.gym/gymnasium is missing, + install tau2-bench[gym] into the current Python environment. + Default: enabled. Set TAU2_REPAIR_VIKINGBOT_GYM=0 or pass + --no-repair-vikingbot-gym to disable automatic repair. + --no-kill-existing Do not stop existing process listening on --port +EOF + exit 0 ;; + *) echo "Unknown argument: $1" >&2; exit 1 ;; + esac +done + +if [[ "${ROLLOUT_LANGUAGE}" != "default" && "${ROLLOUT_LANGUAGE}" != "zh" ]]; then + echo "[tau2-service] invalid --rollout-language: ${ROLLOUT_LANGUAGE}. Expected default or zh" >&2 + exit 1 +fi + +if [[ "${ROLLOUT_BACKEND}" != "native" && "${ROLLOUT_BACKEND}" != "vikingbot" ]]; then + echo "[tau2-service] invalid --rollout-backend: ${ROLLOUT_BACKEND}. Expected native or vikingbot" >&2 + exit 1 +fi + +if ! [[ "${NATIVE_THREAD_WORKERS}" =~ ^[0-9]+$ ]] || [[ "${NATIVE_THREAD_WORKERS}" -le 0 ]]; then + echo "[tau2-service] invalid --native-thread-workers: ${NATIVE_THREAD_WORKERS}. Expected positive integer" >&2 + exit 1 +fi + +if ! [[ "${MAX_ROLLOUT_CONCURRENCY}" =~ ^[0-9]+$ ]] || [[ "${MAX_ROLLOUT_CONCURRENCY}" -le 0 ]]; then + echo "[tau2-service] invalid --max-rollout-concurrency: ${MAX_ROLLOUT_CONCURRENCY}. Expected positive integer" >&2 + exit 1 +fi + +if ! [[ "${ROLLOUT_THREAD_WORKERS}" =~ ^[0-9]+$ ]]; then + echo "[tau2-service] invalid --rollout-thread-workers: ${ROLLOUT_THREAD_WORKERS}. Expected non-negative integer" >&2 + exit 1 +fi + +if [[ "${REPAIR_VIKINGBOT_GYM}" != "0" && "${REPAIR_VIKINGBOT_GYM}" != "1" ]]; then + echo "[tau2-service] invalid TAU2_REPAIR_VIKINGBOT_GYM: ${REPAIR_VIKINGBOT_GYM}. Expected 0 or 1" >&2 + exit 1 +fi + +if [[ -z "${DATA_ROOT}" ]]; then + for _candidate in \ + "${REPO_ROOT}/tau2-bench/data/tau2" \ + "${REPO_ROOT}/../tau2-bench/data/tau2" \ + "${HOME}/workspace/tau2-bench/data/tau2"; do + if [[ -d "${_candidate}/domains" ]]; then + DATA_ROOT="${_candidate}" + break + fi + done +fi +if [[ -z "${DATA_ROOT}" || ! -d "${DATA_ROOT}/domains" ]]; then + echo "[tau2-service] tau2 data root not found. Pass --data-root /data/tau2" >&2 + exit 1 +fi + +TAU2_BENCH_ROOT="${TAU2_BENCH_ROOT:-}" +if [[ -z "${TAU2_BENCH_ROOT}" ]]; then + _maybe_root="$(cd "${DATA_ROOT}/../.." && pwd)" + if [[ -d "${_maybe_root}/src/tau2" ]]; then + TAU2_BENCH_ROOT="${_maybe_root}" + fi +fi +VIKINGBOT_ROOT="${VIKINGBOT_ROOT:-${REPO_ROOT}/bot}" +export PYTHONPATH="${REPO_ROOT}:${VIKINGBOT_ROOT}:${TAU2_BENCH_ROOT:+${TAU2_BENCH_ROOT}/src:}${PYTHONPATH:-}" +export TAU2_DATA_ROOT="${DATA_ROOT}" +export OPENVIKING_CONFIG_FILE="${CONFIG}" +export OPENAI_API_KEY="${OPENAI_API_KEY:-${ARK_API_KEY:-}}" +export OPENAI_API_BASE="${OPENAI_API_BASE:-https://ark.cn-beijing.volces.com/api/v3}" +export OPENAI_BASE_URL="${OPENAI_BASE_URL:-${OPENAI_API_BASE}}" +export AGENT_API_BASE="${AGENT_API_BASE:-${OPENAI_API_BASE}}" +export USER_API_BASE="${USER_API_BASE:-${OPENAI_API_BASE}}" + +check_vikingbot_user_simulator() { + "${PYTHON_BIN}" - <<'PY' >/dev/null 2>&1 +from tau2.gym.gym_agent import AgentGymEnv +import gymnasium # noqa: F401 +assert AgentGymEnv is not None +PY +} + +repair_vikingbot_user_simulator() { + if [[ -z "${TAU2_BENCH_ROOT}" || ! -d "${TAU2_BENCH_ROOT}" ]]; then + echo "[tau2-service] cannot repair vikingbot user simulator: tau2-bench root not found." >&2 + echo "[tau2-service] set TAU2_BENCH_ROOT or pass --data-root /data/tau2." >&2 + return 1 + fi + echo "[tau2-service] repairing vikingbot user simulator dependency: ${PYTHON_BIN} -m pip install -e ${TAU2_BENCH_ROOT}[gym]" + "${PYTHON_BIN}" -m pip install -e "${TAU2_BENCH_ROOT}[gym]" +} + +if [[ "${ROLLOUT_BACKEND}" == "vikingbot" ]]; then + if ! check_vikingbot_user_simulator && [[ "${REPAIR_VIKINGBOT_GYM}" == "1" ]]; then + if ! repair_vikingbot_user_simulator; then + echo "[tau2-service] vikingbot user simulator repair failed; validating again before exit." >&2 + fi + fi + if ! check_vikingbot_user_simulator; then + cat >&2 <}[gym]" +[tau2-service] bash benchmark/tau2/train/run_service.sh --rollout-backend vikingbot --repair-vikingbot-gym +EOF + exit 1 + fi + if [[ -z "${OPENAI_API_KEY:-}" ]]; then + echo "[tau2-service] WARNING: OPENAI_API_KEY/ARK_API_KEY is empty; tau2 user simulator LLM calls may fail." >&2 + fi +fi + +cd "${REPO_ROOT}" +export TAU2_ROLLOUT_BACKEND="${ROLLOUT_BACKEND}" +export TAU2_NATIVE_THREAD_WORKERS="${NATIVE_THREAD_WORKERS}" +export TAU2_MAX_ROLLOUT_CONCURRENCY="${MAX_ROLLOUT_CONCURRENCY}" +export TAU2_ROLLOUT_THREAD_WORKERS="${ROLLOUT_THREAD_WORKERS}" +echo "[tau2-service] host=${HOST} port=${PORT} data_root=${DATA_ROOT} config=${CONFIG} rollout_language=${ROLLOUT_LANGUAGE} rollout_backend=${ROLLOUT_BACKEND} native_thread_workers=${NATIVE_THREAD_WORKERS} max_rollout_concurrency=${MAX_ROLLOUT_CONCURRENCY} rollout_thread_workers=${ROLLOUT_THREAD_WORKERS}" +if [[ "${KILL_EXISTING}" == "1" ]]; then + EXISTING_PIDS="$(lsof -tiTCP:"${PORT}" -sTCP:LISTEN 2>/dev/null || true)" + if [[ -n "${EXISTING_PIDS}" ]]; then + echo "[tau2-service] stopping existing listener(s) on port ${PORT}: ${EXISTING_PIDS}" + kill ${EXISTING_PIDS} 2>/dev/null || true + for _ in {1..20}; do + sleep 0.2 + if ! lsof -tiTCP:"${PORT}" -sTCP:LISTEN >/dev/null 2>&1; then + break + fi + done + REMAINING_PIDS="$(lsof -tiTCP:"${PORT}" -sTCP:LISTEN 2>/dev/null || true)" + if [[ -n "${REMAINING_PIDS}" ]]; then + echo "[tau2-service] force stopping listener(s) on port ${PORT}: ${REMAINING_PIDS}" + kill -9 ${REMAINING_PIDS} 2>/dev/null || true + fi + fi +fi +exec "${PYTHON_BIN}" "${SCRIPT_DIR}/service_app.py" --host "${HOST}" --port "${PORT}" --data-root "${DATA_ROOT}" --config "${CONFIG}" --rollout-language "${ROLLOUT_LANGUAGE}" --rollout-backend "${ROLLOUT_BACKEND}" --native-thread-workers "${NATIVE_THREAD_WORKERS}" --max-rollout-concurrency "${MAX_ROLLOUT_CONCURRENCY}" --rollout-thread-workers "${ROLLOUT_THREAD_WORKERS}" diff --git a/benchmark/tau2/train/service_app.py b/benchmark/tau2/train/service_app.py new file mode 100644 index 0000000000..f8ffc29952 --- /dev/null +++ b/benchmark/tau2/train/service_app.py @@ -0,0 +1,217 @@ +#!/usr/bin/env python3 +"""HTTP service exposing tau2 cases and rollout execution.""" + +# ruff: noqa: E402 + +from __future__ import annotations + +import argparse +import asyncio +import logging +import os +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Any + +import uvicorn + +DEFAULT_NATIVE_THREAD_WORKERS = 128 +DEFAULT_MAX_ROLLOUT_CONCURRENCY = 200 +DEFAULT_ROLLOUT_THREAD_WORKERS = 200 +TAU2_SERVICE_LOG_LEVEL = "WARNING" + +REPO_ROOT = Path(__file__).resolve().parents[3] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from benchmark.tau2.train.case_loader import Tau2CaseLoader +from benchmark.tau2.train.rollout_executor import ( + DEFAULT_TAU2_ROLLOUT_BACKEND, + make_tau2_rollout_executor, + normalize_tau2_rollout_backend, +) +from openviking.session.train.components.dataset_service import create_dataset_service_app + + +def configure_tau2_service_logging() -> None: + """Keep third-party tau2/loguru service output at warning and above.""" + logging.getLogger("tau2").setLevel(logging.WARNING) + try: + from loguru import logger as loguru_logger + except Exception: + return + + loguru_logger.remove() + loguru_logger.add(sys.stderr, level=TAU2_SERVICE_LOG_LEVEL) + + +def create_app( + *, + data_root: str | None = None, + config_path: str | None = None, + rollout_language: str = "default", + rollout_backend: str | None = None, + max_rollout_concurrency: int | None = None, + rollout_thread_workers: int | None = None, +): + if rollout_language not in {"default", "zh"}: + raise ValueError("rollout_language must be 'default' or 'zh'") + default_backend = normalize_tau2_rollout_backend( + rollout_backend or os.getenv("TAU2_ROLLOUT_BACKEND") or DEFAULT_TAU2_ROLLOUT_BACKEND + ) + + def make_case_loader( + dataset: str, + domain: str, + split: str, + filters: dict[str, Any], + ) -> Tau2CaseLoader: + if dataset != "tau2": + raise ValueError(f"Unsupported dataset: {dataset}") + return Tau2CaseLoader( + domain=domain, + split=split, + data_root=data_root, + task_indices=_task_indices_from_filters(filters), + ) + + def make_rollout_executor(options: dict[str, Any]): + backend = normalize_tau2_rollout_backend( + options.get("rollout_backend") or options.get("backend") or default_backend + ) + return make_tau2_rollout_executor( + backend=backend, + options={ + **options, + "show_progress": options.get("show_progress", False), + "progress_label": options.get("progress_label") or "tau2", + }, + config_path=config_path, + concurrency=1, + rollout_language=rollout_language, + ) + + return create_dataset_service_app( + service_name="tau2", + make_case_loader=make_case_loader, + make_rollout_executor=make_rollout_executor, + max_rollout_concurrency=max_rollout_concurrency, + rollout_thread_workers=rollout_thread_workers, + ) + + +def _task_indices_from_filters(filters: dict[str, Any]) -> list[int] | None: + raw = filters.get("task_indices") + if raw is None: + return None + if not isinstance(raw, list): + raise ValueError("task_indices filter must be a list") + indices: list[int] = [] + for value in raw: + index = int(value) + if index < 0: + raise ValueError("task index must be >= 0") + indices.append(index) + return indices + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Start tau2 rollout HTTP service") + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=1944) + parser.add_argument("--data-root", default=os.getenv("TAU2_DATA_ROOT")) + parser.add_argument("--config", default=os.getenv("OPENVIKING_CONFIG_FILE")) + parser.add_argument("--rollout-language", choices=["default", "zh"], default="default") + parser.add_argument( + "--rollout-backend", + choices=["native", "vikingbot"], + default=os.getenv("TAU2_ROLLOUT_BACKEND", DEFAULT_TAU2_ROLLOUT_BACKEND), + help="Rollout implementation backend (default: native).", + ) + parser.add_argument( + "--native-thread-workers", + type=int, + default=int(os.getenv("TAU2_NATIVE_THREAD_WORKERS", str(DEFAULT_NATIVE_THREAD_WORKERS))), + help="Default thread pool workers for native tau2 rollout execution (default: 128).", + ) + parser.add_argument( + "--max-rollout-concurrency", + type=int, + default=int( + os.getenv( + "TAU2_MAX_ROLLOUT_CONCURRENCY", + str(DEFAULT_MAX_ROLLOUT_CONCURRENCY), + ) + ), + help=( + "Maximum concurrent rollout executions hosted by this service. " + f"Default: {DEFAULT_MAX_ROLLOUT_CONCURRENCY}." + ), + ) + parser.add_argument( + "--rollout-thread-workers", + type=int, + default=int( + os.getenv( + "TAU2_ROLLOUT_THREAD_WORKERS", + str(DEFAULT_ROLLOUT_THREAD_WORKERS), + ) + ), + help=( + "Worker threads used to host rollout executions off the uvicorn event loop. " + "Set to 0 to disable threaded hosting. " + f"Default: {DEFAULT_ROLLOUT_THREAD_WORKERS}." + ), + ) + return parser.parse_args() + + +class Tau2ServiceServer(uvicorn.Server): + def __init__(self, config: uvicorn.Config, *, native_thread_workers: int) -> None: + super().__init__(config) + self._native_thread_workers = native_thread_workers + self._default_executor: ThreadPoolExecutor | None = None + + async def serve(self, sockets=None) -> None: + if self._native_thread_workers <= 0: + raise ValueError("native_thread_workers must be > 0") + loop = asyncio.get_running_loop() + self._default_executor = ThreadPoolExecutor( + max_workers=self._native_thread_workers, + thread_name_prefix="tau2-native", + ) + loop.set_default_executor(self._default_executor) + try: + await super().serve(sockets=sockets) + finally: + self._default_executor.shutdown(wait=False, cancel_futures=True) + + +def main() -> None: + args = parse_args() + + configure_tau2_service_logging() + app = create_app( + data_root=args.data_root, + config_path=args.config, + rollout_language=args.rollout_language, + rollout_backend=args.rollout_backend, + max_rollout_concurrency=args.max_rollout_concurrency, + rollout_thread_workers=( + None if args.rollout_thread_workers == 0 else args.rollout_thread_workers + ), + ) + config = uvicorn.Config( + app, + host=args.host, + port=args.port, + access_log=False, + log_level="warning", + ) + server = Tau2ServiceServer(config, native_thread_workers=args.native_thread_workers) + server.run() + + +if __name__ == "__main__": + main() diff --git a/benchmark/tau2/vikingbot/README.md b/benchmark/tau2/vikingbot/README.md index 27700f538a..2651846f54 100644 --- a/benchmark/tau2/vikingbot/README.md +++ b/benchmark/tau2/vikingbot/README.md @@ -58,7 +58,7 @@ env vars. Override any of these by exporting before sourcing: The tau2 **user simulator** talks to an OpenAI-compatible endpoint — set `ARK_API_KEY` (e.g. Doubao through volcengine ARK) before sourcing, or the simulator will fail. The user-simulator -model is configured in [`tau2_env/tau2_environment.py`](tau2_env/tau2_environment.py). +model is configured in [`../common/tau2_env/tau2_environment.py`](../common/tau2_env/tau2_environment.py). > Note: the sibling `llm/` harness ([`../llm/README.md`](../llm/README.md)) pins a tau2-bench ref > with a confirmation-aware user-simulator prompt (sierra-research/tau2-bench#297). Set @@ -234,16 +234,16 @@ Two VikingBot behaviours need to change for tau2 self-improvement: turn. For tau2 we want accumulated **experience** memory pulled once per task, with a larger character budget per experience. -Both are now controlled by config — set these three flags in the `bot.ov_server` section of the -`ov.conf` pointed to by `OPENVIKING_CONFIG_FILE`: +Experience recall is enabled by default and retrieves experience memory once on the first user turn. +You can tune these two flags in the `bot.ov_server` section of the `ov.conf` pointed to by +`OPENVIKING_CONFIG_FILE`: ```jsonc { "bot": { "ov_server": { - "recall_exp_first_round_only": true, // skip per-turn recall; inject exp once on the first user turn - "exp_recall_limit": 2, // fetch 2 experiences per task (default: 5) - "exp_recall_max_chars": 10000 // character budget for the injected experience block (default: 2000) + "exp_recall_limit": 2, // fetch 2 experiences per task (default: 5) + "exp_recall_max_chars": 10000 // character budget for the injected experience block (default: 2000) } } } @@ -251,10 +251,6 @@ Both are now controlled by config — set these three flags in the `bot.ov_serve What each flag does: -- **`recall_exp_first_round_only`** — when `true`, `ContextBuilder._build_user_memory` skips the - default per-turn memory recall and instead calls `get_viking_experience_context` once, on the - first user-turn of the session. The runner only ever sends one user turn per task, so this - becomes "fetch experience once per task." - **`exp_recall_limit`** — how many experiences to retrieve. tau2 prefers **fewer but longer** (2) over many shallow hits (5). - **`exp_recall_max_chars`** — total character budget for the formatted experience block. Bumped @@ -271,7 +267,7 @@ runtime identity, so each domain reads and writes `viking://user//. - `setup_env.sh` — environment setup (PYTHONPATH, tau2 data root, simulator LLM) - `run_full_test.sh` — full pipeline for one epoch (run → eval → commit) - `run_airline_2epochs.sh` — multi-epoch example (cold start → memory-augmented epochs) -- `tau2_env/` — tau2 environment integration (`tau2_environment.py`, `tau2_tool_provider.py`) +- `../common/tau2_env/` — tau2 environment integration (`tau2_environment.py`, `tau2_tool_provider.py`) - `scripts/` - `provision_openviking_user.py` — create/refresh a benchmark user and write a user-key config - `vikingbot_tau2_runner.py` — runs a single tau2 task through the VikingBot agent loop diff --git a/benchmark/tau2/vikingbot/scripts/vikingbot_tau2_runner.py b/benchmark/tau2/vikingbot/scripts/vikingbot_tau2_runner.py index 8cb13e935b..344a80d752 100644 --- a/benchmark/tau2/vikingbot/scripts/vikingbot_tau2_runner.py +++ b/benchmark/tau2/vikingbot/scripts/vikingbot_tau2_runner.py @@ -18,11 +18,16 @@ from typing import Any SCRIPT_DIR = Path(__file__).resolve().parent -REPO_ROOT = SCRIPT_DIR.parent -if str(REPO_ROOT) not in sys.path: - sys.path.insert(0, str(REPO_ROOT)) - -from tau2_env.tau2_tool_provider import Tau2BenchToolProvider, load_task_id +TAU2_ROOT = SCRIPT_DIR.parents[1] +OPENVIKING_ROOT = TAU2_ROOT.parents[1] +for _path in (OPENVIKING_ROOT,): + if str(_path) not in sys.path: + sys.path.insert(0, str(_path)) + +from benchmark.tau2.common.tau2_env.tau2_tool_provider import ( + Tau2BenchToolProvider, + load_task_id, +) try: from vikingbot.agent.loop import AgentLoop diff --git a/benchmark/tau2/vikingbot/setup_env.sh b/benchmark/tau2/vikingbot/setup_env.sh index 42c039fdaf..b3c758be58 100755 --- a/benchmark/tau2/vikingbot/setup_env.sh +++ b/benchmark/tau2/vikingbot/setup_env.sh @@ -42,6 +42,15 @@ TAU2_BENCH_REPO="${TAU2_BENCH_REPO:-https://github.com/sierra-research/tau2-benc (return 0 2>/dev/null) && _SOURCED=1 || _SOURCED=0 _abort() { echo "[setup_env] ERROR: $*" >&2; if [[ "${_SOURCED}" -eq 1 ]]; then return 1; else exit 1; fi; } +_tau2_user_simulator_ready() { + local py="$1" + "${py}" - <<'PYEOF' >/dev/null 2>&1 +from tau2.gym.gym_agent import AgentGymEnv +import gymnasium # noqa: F401 +assert AgentGymEnv is not None +PYEOF +} + # --- parse args --- REINSTALL=0 for _arg in "$@"; do @@ -61,7 +70,26 @@ _setup_install() { fi if [[ -f "${SETUP_MARKER}" && "${REINSTALL}" -eq 0 ]]; then - return 0 # already installed; nothing to do + local PY="${VENV}/bin/python" + if [[ -x "${PY}" ]] && _tau2_user_simulator_ready "${PY}"; then + return 0 # already installed; nothing to do + fi + echo "[setup_env] Existing setup marker found, but tau2 gym user simulator is missing; repairing tau2-bench[gym] install." + if [[ ! -x "${PY}" ]]; then + echo "[setup_env] venv python not found at ${PY}; continuing with full setup." + else + if [[ ! -d "${TAU2_BENCH_ROOT}" ]]; then + echo "[setup_env] tau2-bench checkout missing at ${TAU2_BENCH_ROOT}; continuing with full setup." + else + echo "[setup_env] Installing tau2-bench with gym extra (pip install -e ${TAU2_BENCH_ROOT}[gym])" + "${PY}" -m pip install -e "${TAU2_BENCH_ROOT}[gym]" || { echo "[setup_env] tau2-bench gym repair failed"; return 1; } + if ! _tau2_user_simulator_ready "${PY}"; then + echo "[setup_env] tau2 gym user simulator still unavailable after repair; try: source ${BASH_SOURCE[0]} --reinstall" + return 1 + fi + return 0 + fi + fi fi if [[ ! -f "${VENV}/bin/activate" ]]; then @@ -145,6 +173,10 @@ PYEOF # tau2-bench: install the [gym] extra so tau2.gym (gymnasium) is available to the runner. echo "[setup_env] Installing tau2-bench with gym extra (pip install -e ${TAU2_BENCH_ROOT}[gym])" "${PY}" -m pip install -e "${TAU2_BENCH_ROOT}[gym]" || { echo "[setup_env] tau2-bench install failed"; return 1; } + if ! _tau2_user_simulator_ready "${PY}"; then + echo "[setup_env] tau2 gym user simulator is unavailable after installing tau2-bench[gym]" + return 1 + fi echo "[setup_env] Installing smolagents" "${PY}" -m pip install smolagents || { echo "[setup_env] smolagents install failed"; return 1; } @@ -156,7 +188,7 @@ PYEOF if ! _setup_install; then _abort "environment install failed (see messages above)" - unset -f _setup_install _abort + unset -f _setup_install _abort _tau2_user_simulator_ready return 1 2>/dev/null || exit 1 fi unset -f _setup_install @@ -187,6 +219,9 @@ export OPENVIKING_CONFIG_FILE="${OPENVIKING_CONFIG_FILE:-${HOME}/.openviking/ov. # Provide your own key via ARK_API_KEY (do NOT commit real keys). export OPENAI_API_KEY="${OPENAI_API_KEY:-${ARK_API_KEY:-}}" export OPENAI_API_BASE="${OPENAI_API_BASE:-https://ark.cn-beijing.volces.com/api/v3}" +export OPENAI_BASE_URL="${OPENAI_BASE_URL:-${OPENAI_API_BASE}}" +export AGENT_API_BASE="${AGENT_API_BASE:-${OPENAI_API_BASE}}" +export USER_API_BASE="${USER_API_BASE:-${OPENAI_API_BASE}}" if [[ -z "${OPENAI_API_KEY}" ]]; then echo "[setup_env] WARNING: OPENAI_API_KEY/ARK_API_KEY is empty; the tau2 user simulator will fail." fi @@ -195,4 +230,4 @@ echo "[setup_env] PYTHONPATH includes openviking (${OPENVIKING_TAU2_ROOT}) and v echo "[setup_env] TAU2_DATA_ROOT=${TAU2_DATA_ROOT}" echo "[setup_env] OPENAI_API_BASE=${OPENAI_API_BASE}" -unset -f _abort 2>/dev/null || true +unset -f _abort _tau2_user_simulator_ready 2>/dev/null || true diff --git a/benchmark/tau2/vikingbot/tau2_env/tau2_environment.py b/benchmark/tau2/vikingbot/tau2_env/tau2_environment.py deleted file mode 100644 index 27a89143a6..0000000000 --- a/benchmark/tau2/vikingbot/tau2_env/tau2_environment.py +++ /dev/null @@ -1,96 +0,0 @@ -#!/usr/bin/env python3 - -from __future__ import annotations - -import json - -from tau2.gym.gym_agent import AgentGymEnv - - -class CommunicateWithUser: - """The agent's only channel for talking to the user. - - tau2's environment has no native "speak to the user" action, so we add this - tool: whatever ``content`` the agent passes is delivered to the tau2 user - simulator, and the simulator's reply comes back as the tool result. This class - owns both the tool's schema (``openai_schema``, consumed by - ``Tau2BenchToolProvider``) and its execution (``forward``, invoked by - ``Tau2BenchEnv.tool_call``). - """ - - name = "communicate_with_user" - description = ( - "say something to the user. Note that the customer cannot see the answer " - "returned in `final_answer`. You must communicate with the customer " - "exclusively through this tool." - ) - parameters = { - "type": "object", - "properties": { - "content": { - "type": "string", - "description": "the content to say to the user", - } - }, - "required": ["content"], - } - - def __init__(self, env): - # ``env`` is the underlying tau2 AgentGymEnv. - self.env = env - - def forward(self, content: str): - """Deliver ``content`` to the tau2 user simulator. - - Returns the raw gym step tuple ``(obs, reward, terminated, truncated, info)``; - ``Tau2BenchEnv.tool_call`` cleans the observation and tracks termination. - """ - response = self.env.tool_call(self.name, {"content": content}) - return response - - @classmethod - def openai_schema(cls) -> dict: - return { - "type": "function", - "function": { - "name": cls.name, - "description": cls.description, - "parameters": cls.parameters, - }, - } - - -class Tau2BenchEnv: - def __init__(self, domain: str, task_id: str): - self.env = AgentGymEnv(domain=domain, task_id=task_id, user_llm="openai/doubao-seed-2-0-pro-260215") - self.terminated = False - - def reset(self): - user_query, info_dict = self.env.reset() - self.user_query = user_query.lstrip("user: ") - self.task = info_dict["task"] - self.simulation_run = info_dict["simulation_run"] - self.policy = info_dict["policy"] - # All OpenAI tool schemas exposed to the agent: tau2's native tools plus - # communicate_with_user (the agent's only channel to the user). - self.tool_schemas = [tool.openai_schema for tool in info_dict["tools"]] - self.tool_schemas.append(CommunicateWithUser.openai_schema()) - self.ground_truth = str(self.task.evaluation_criteria) - self.user_scenario = self.task.user_scenario - - def tool_call(self, tool_name: str, arguments: dict) -> str: - if self.terminated: - return "Task Terminated" - - if tool_name == CommunicateWithUser.name: - obs, reward, terminated, truncated, info = self.env.step(arguments["content"]) - else: - action = {"name": tool_name, "arguments": arguments} - obs, reward, terminated, truncated, info = self.env.step(json.dumps(action)) - - if "tool: " in obs: - obs = obs.lstrip("tool: ") - if "user: " in obs: - obs = obs.lstrip("user: ") - self.terminated = terminated - return obs diff --git a/bot/scripts/restart_openviking_server.sh b/bot/scripts/restart_openviking_server.sh index 38171f9ef7..5ba9e4d5b3 100755 --- a/bot/scripts/restart_openviking_server.sh +++ b/bot/scripts/restart_openviking_server.sh @@ -1,28 +1,67 @@ #!/bin/bash # Restart OpenViking Server with Bot API enabled -# Usage: ./restart_openviking_server.sh [--port PORT] [--bot-port PORT] +# Usage: ./restart_openviking_server.sh [--port PORT] [--bot-port PORT] [--config PATH] [--data-dir PATH] set -e # Default values PORT="1933" BOT_PORT="18790" +CONFIG="${OPENVIKING_CONFIG_FILE:-${HOME}/.openviking/ov.conf}" +DATA_DIR="${OPENVIKING_DATA_DIR:-${HOME}/.openviking/data}" +KILL_ALL_VIKINGBOT="${OPENVIKING_KILL_ALL_VIKINGBOT:-1}" + +usage() { + echo "Usage: $0 [--port PORT] [--bot-port PORT] [--config PATH] [--data-dir PATH] [--no-kill-all-vikingbot]" +} # Parse arguments while [[ $# -gt 0 ]]; do case $1 in --port) + if [[ $# -lt 2 ]]; then + echo "Missing value for --port" + usage + exit 1 + fi PORT="$2" shift 2 ;; --bot-port) + if [[ $# -lt 2 ]]; then + echo "Missing value for --bot-port" + usage + exit 1 + fi BOT_PORT="$2" shift 2 ;; + --config) + if [[ $# -lt 2 ]]; then + echo "Missing value for --config" + usage + exit 1 + fi + CONFIG="$2" + shift 2 + ;; + --data-dir) + if [[ $# -lt 2 ]]; then + echo "Missing value for --data-dir" + usage + exit 1 + fi + DATA_DIR="$2" + shift 2 + ;; + --no-kill-all-vikingbot) + KILL_ALL_VIKINGBOT=0 + shift 1 + ;; *) echo "Unknown option: $1" - echo "Usage: $0 [--port PORT] [--bot-port PORT]" + usage exit 1 ;; esac @@ -33,13 +72,15 @@ echo "Restarting OpenViking Server with Bot API" echo "==========================================" echo "OpenViking Server Port: $PORT" echo "Bot Port: $BOT_PORT" +echo "Config: $CONFIG" +echo "Data Directory: $DATA_DIR" echo "" # Step 0: Kill process on port and delete data directory echo "Step 0: Killing process on port $PORT..." -if lsof -i :"$PORT" > /dev/null 2>&1; then - pid=$(lsof -ti :"$PORT") - kill -9 "$pid" 2>/dev/null || true +if lsof -nP -iTCP:"$PORT" -sTCP:LISTEN > /dev/null 2>&1; then + pid=$(lsof -tiTCP:"$PORT" -sTCP:LISTEN) + kill -9 $pid 2>/dev/null || true sleep 1 echo " ✓ Killed process $pid on port $PORT" else @@ -47,33 +88,49 @@ else fi echo "" -echo "Step 0b: Deleting data directory /Users/bytedance/.openviking/data..." -if [ -d "/Users/bytedance/.openviking/data" ]; then - rm -rf /Users/bytedance/.openviking/data - echo " ✓ Deleted /Users/bytedance/.openviking/data" +echo "Step 0b: Deleting data directory $DATA_DIR..." +if [ -d "$DATA_DIR" ]; then + rm -rf "$DATA_DIR" + echo " ✓ Deleted $DATA_DIR" else echo " ✓ Data directory does not exist" fi -# Kill existing vikingbot processes +# Kill existing vikingbot processes. Multi-slot launchers pass +# --no-kill-all-vikingbot so one slot restart does not kill other slots. echo "" echo "Step 0c: Stopping existing vikingbot processes..." -if pgrep -f "vikingbot.*openapi" > /dev/null 2>&1 || pgrep -f "vikingbot.*gateway" > /dev/null 2>&1; then - pkill -f "vikingbot.*openapi" 2>/dev/null || true - pkill -f "vikingbot.*gateway" 2>/dev/null || true - sleep 2 - echo " ✓ Stopped existing vikingbot processes" +if [[ "$KILL_ALL_VIKINGBOT" == "1" ]]; then + if pgrep -f "vikingbot.*openapi" > /dev/null 2>&1 || pgrep -f "vikingbot.*gateway" > /dev/null 2>&1; then + pkill -f "vikingbot.*openapi" 2>/dev/null || true + pkill -f "vikingbot.*gateway" 2>/dev/null || true + sleep 2 + echo " ✓ Stopped existing vikingbot processes" + else + echo " ✓ No existing vikingbot processes found" + fi +else + echo " ✓ Skipped global vikingbot kill" +fi + +echo "" +echo "Step 0d: Killing process on bot port $BOT_PORT..." +if lsof -nP -iTCP:"$BOT_PORT" -sTCP:LISTEN > /dev/null 2>&1; then + bot_pid=$(lsof -tiTCP:"$BOT_PORT" -sTCP:LISTEN) + kill -9 $bot_pid 2>/dev/null || true + sleep 1 + echo " ✓ Killed process $bot_pid on bot port $BOT_PORT" else - echo " ✓ No existing vikingbot processes found" + echo " ✓ No process found on bot port $BOT_PORT" fi # Step 1: Verify port is free echo "" echo "Step 1: Verifying port $PORT is free..." -if lsof -i :"$PORT" > /dev/null 2>&1; then +if lsof -nP -iTCP:"$PORT" -sTCP:LISTEN > /dev/null 2>&1; then echo " ✗ Port $PORT is still in use, trying to force kill..." - pid=$(lsof -ti :"$PORT") - kill -9 "$pid" 2>/dev/null || true + pid=$(lsof -tiTCP:"$PORT" -sTCP:LISTEN) + kill -9 $pid 2>/dev/null || true sleep 1 fi echo " ✓ Port $PORT is free" @@ -81,7 +138,7 @@ echo " ✓ Port $PORT is free" # Step 2: Start openviking-server with --with-bot echo "" echo "Step 2: Starting openviking-server with Bot API..." -echo " Command: openviking-server --with-bot --port $PORT --bot-port $BOT_PORT" +echo " Command: OPENVIKING_CONFIG_FILE=$CONFIG openviking-server --with-bot --port $PORT --bot-port $BOT_PORT --config $CONFIG" echo "" # Start in background and log to file @@ -91,10 +148,12 @@ echo "" # --bot-port "$BOT_PORT" \ # > /tmp/openviking-server.log 2>&1 & +export OPENVIKING_CONFIG_FILE="$CONFIG" openviking-server \ --with-bot \ --port "$PORT" \ - --bot-port "$BOT_PORT" + --bot-port "$BOT_PORT" \ + --config "$CONFIG" SERVER_PID=$! @@ -105,25 +164,23 @@ echo "" echo "Step 3: Waiting for server to be ready..." sleep 3 -# First check if server is responding at all +# First check if bot proxy health reports healthy for i in {1..10}; do - if curl -s http://localhost:"$PORT"/api/v1/bot/health > /dev/null 2>&1; then + health_response=$(curl -fsS http://localhost:"$PORT"/bot/v1/health 2>/dev/null || true) + if echo "${health_response//[[:space:]]/}" | grep -q '"status":"healthy"'; then echo "" echo "==========================================" echo "✓ OpenViking Server started successfully!" echo "==========================================" echo "" echo "Server URL: http://localhost:$PORT" - echo "Health Check: http://localhost:$PORT/api/v1/bot/health" + echo "Health Check: http://localhost:$PORT/bot/v1/health" echo "Logs: tail -f /tmp/openviking-server.log" echo "" exit 0 fi # Check actual health response - health_response=$(curl -s http://localhost:"$PORT"/api/v1/bot/health 2>/dev/null) - if echo "$health_response" | grep -q "Vikingbot"; then - echo " ✓ Vikingbot is healthy" - elif echo "$health_response" | grep -q "Bot service unavailable"; then + if echo "$health_response" | grep -q "Bot service unavailable"; then echo " ⏳ Waiting for Vikingbot to start (attempt $i/10)..." fi sleep 2 diff --git a/bot/scripts/test_restart_openviking_server.sh b/bot/scripts/test_restart_openviking_server.sh index 7e3875f7d3..950be0f3d6 100755 --- a/bot/scripts/test_restart_openviking_server.sh +++ b/bot/scripts/test_restart_openviking_server.sh @@ -117,9 +117,10 @@ echo "" echo "Step 6: Waiting for server to be ready..." sleep 3 -# First check if server is responding at all +# First check if bot proxy health reports healthy for i in {1..10}; do - if curl -s http://localhost:"$PORT"/api/v1/bot/health > /dev/null 2>&1; then + health_response=$(curl -fsS http://localhost:"$PORT"/bot/v1/health 2>/dev/null || true) + if echo "${health_response//[[:space:]]/}" | grep -q '"status":"healthy"'; then echo "" echo "==========================================" echo "✓ OpenViking Server started successfully! (TEST MODE)" @@ -128,15 +129,12 @@ for i in {1..10}; do echo "Server URL: http://localhost:$PORT" echo "Config File: $TEST_CONFIG" echo "Data Dir: $TEST_DATA_DIR" - echo "Health Check: http://localhost:$PORT/api/v1/bot/health" + echo "Health Check: http://localhost:$PORT/bot/v1/health" echo "" exit 0 fi # Check actual health response - health_response=$(curl -s http://localhost:"$PORT"/api/v1/bot/health 2>/dev/null) - if echo "$health_response" | grep -q "Vikingbot"; then - echo " ✓ Vikingbot is healthy" - elif echo "$health_response" | grep -q "Bot service unavailable"; then + if echo "$health_response" | grep -q "Bot service unavailable"; then echo " ⏳ Waiting for Vikingbot to start (attempt $i/10)..." fi sleep 2 diff --git a/bot/tests/test_openviking_api_key_type.py b/bot/tests/test_openviking_api_key_type.py index 0cafa12b48..155f3d83ed 100644 --- a/bot/tests/test_openviking_api_key_type.py +++ b/bot/tests/test_openviking_api_key_type.py @@ -119,6 +119,7 @@ def _make_config(api_key_type: str, mode: str = "remote", **ov_overrides): agents = SimpleNamespace(**{**agent_defaults, **agent_overrides}) ov_server = SimpleNamespace( mode=mode, + auth_mode="", api_key_type=api_key_type, server_url="http://ov.local", api_key="root-key" if api_key_type == "root" else "user-key", @@ -2556,6 +2557,9 @@ def test_openviking_search_description_allows_follow_up_memory_queries(): @pytest.mark.asyncio async def test_context_reminds_agent_to_search_current_memory_question(tmp_path): class _EmptyMemory: + async def get_viking_experience_context(self, **_kwargs): + return "" + async def get_viking_memory_context(self, **_kwargs): return "" @@ -2567,7 +2571,6 @@ async def get_viking_memory_context(self, **_kwargs): current_message="我会哪些语言", sender_id="sender-1", ov_tools_enable=True, - is_first_round=False, ) assert "OpenViking Memory Retrieval" in user_info @@ -2585,8 +2588,14 @@ async def get_viking_memory_context(self, **_kwargs): @pytest.mark.asyncio async def test_context_memory_prefix_tells_agent_to_read_summary_and_uri_details(tmp_path): + calls = [] + class _Memory: - async def get_viking_memory_context(self, **_kwargs): + async def get_viking_experience_context(self, **_kwargs): + return "" + + async def get_viking_memory_context(self, **kwargs): + calls.append(kwargs) return ( "### user memories:\n" '\n' @@ -2603,9 +2612,21 @@ async def get_viking_memory_context(self, **_kwargs): current_message="问题", sender_id="sender-1", ov_tools_enable=True, - is_first_round=False, + memory_peer_ids=["peer-1"], + memory_owner_user_ids=["owner-1"], ) + assert calls == [ + { + "current_message": "问题", + "workspace_id": "cli__default__chat-1", + "sender_id": "sender-1", + "peer_ids": ["peer-1"], + "user_ids": ["owner-1"], + "openviking_connection": None, + } + ] + assert context.latest_relevant_memories assert "## openviking_search(query=[user_query])" in user_info assert "grouped by memory_type" in user_info assert "full means the full memory content is already shown" in user_info @@ -2781,3 +2802,313 @@ async def test_openviking_request_connection_client_is_closed_after_tool_call(mo assert result == "No results found for query: hello" assert len(_DummyHTTPClient.instances) == 1 assert _DummyHTTPClient.instances[0].closed is True + + +@pytest.mark.asyncio +async def test_experience_reminder_follows_direct_case_experience_links(monkeypatch, tmp_path): + from openviking.session.memory.dataclass import MemoryFile, StoredLink + from openviking.session.memory.utils.memory_file_utils import MemoryFileUtils + + case_uri = "viking://user/admin/memories/cases/case1.md" + exp_uri = "viking://user/admin/memories/experiences/exp1.md" + direct_exp_uri = "viking://user/admin/memories/experiences/direct.md" + case_content = MemoryFileUtils.write( + MemoryFile( + uri=case_uri, + content="# case1\n\n## Linked Experiences\n- exp1", + memory_type="cases", + links=[ + StoredLink( + from_uri=case_uri, + to_uri=exp_uri, + link_type="related_to", + weight=1.0, + match_text=None, + description="", + ).model_dump() + ], + ) + ) + + class _FakeClient: + admin_user_id = "admin" + + def __init__(self): + self.search_calls = [] + self.search_experience_calls = [] + self.read_calls = [] + + async def search(self, *, query, target_uri, limit): + self.search_calls.append((query, target_uri, limit)) + if target_uri.endswith("/cases/"): + return {"memories": [{"uri": case_uri, "score": 0.9}]} + return {"memories": []} + + async def search_experiences(self, query, limit=5): + self.search_experience_calls.append((query, limit)) + return [{"uri": direct_exp_uri, "score": 0.6}] if limit > 0 else [] + + async def read_content(self, uri, level="read"): + self.read_calls.append((uri, level)) + if uri == case_uri: + return "# case1\n\n## Linked Experiences\n- [exp1](../experiences/exp1.md)" + if uri == exp_uri: + return "linked exp content" + if uri == direct_exp_uri: + return "direct exp content" + return "" + + async def close(self): + return None + + fake_client = _FakeClient() + + async def _fake_create(**_kwargs): + return fake_client + + monkeypatch.setattr( + "vikingbot.agent.memory.load_config", + lambda: _make_config( + "root", + case_recall_limit=1, + exp_recall_limit=2, + exp_recall_max_chars=4000, + ), + ) + monkeypatch.setattr("vikingbot.agent.memory.VikingClient.create", _fake_create) + + content, uris = await MemoryStore(tmp_path).get_viking_experience_reminder( + query="hello", + workspace_id="workspace", + ) + + assert fake_client.search_calls == [("hello", "viking://user/admin/memories/cases/", 1)] + assert fake_client.search_experience_calls == [] + assert (case_uri, "read") in fake_client.read_calls + assert "linked exp content" in content + assert "direct exp content" not in content + assert exp_uri in uris + assert direct_exp_uri not in uris + + +@pytest.mark.asyncio +async def test_tau2_experience_reminder_loads_case_by_exact_task_not_query(monkeypatch, tmp_path): + from openviking.session.memory.dataclass import MemoryFile, StoredLink + from openviking.session.memory.utils.memory_file_utils import MemoryFileUtils + + matched_case_uri = "viking://user/admin/memories/cases/tau2_airline_train_1.md" + wrong_case_uri = "viking://user/admin/memories/cases/tau2_airline_train_9.md" + matched_exp_uri = "viking://user/admin/memories/experiences/matched.md" + wrong_exp_uri = "viking://user/admin/memories/experiences/wrong.md" + + def _case_content(case_uri, exp_uri, *, task_no, task_id): + return MemoryFileUtils.write( + MemoryFile( + uri=case_uri, + content="", + memory_type="cases", + extra_fields={ + "case_name": case_uri.rsplit("/", 1)[-1].removesuffix(".md"), + "task_signature": f"tau2:airline:train:{task_id}", + "input": json.dumps( + { + "domain": "airline", + "split": "train", + "data_split": "airline_train", + "task_no": task_no, + "task_id": str(task_id), + } + ), + }, + links=[ + StoredLink( + from_uri=case_uri, + to_uri=exp_uri, + link_type="related_to", + weight=1.0, + match_text=None, + description="", + ).model_dump() + ], + ), + content_template=( + "# {{ case_name }}\n\n" + "## Task Signature\n{{ task_signature }}\n\n" + "## Input\n{{ input }}\n\n" + "## Linked Experiences\n" + "{% for link in links or [] %}" + "- [{{ uri_basename(link.to_uri) }}]({{ link_target(link.to_uri) }})\n" + "{% endfor %}" + ), + ) + + raw_cases = { + matched_case_uri: _case_content(matched_case_uri, matched_exp_uri, task_no=1, task_id=7), + wrong_case_uri: _case_content(wrong_case_uri, wrong_exp_uri, task_no=9, task_id=99), + } + visible_cases = { + uri: raw.split(" + + Tau2 Runtime Service :1944 + 依赖 tau2 / VikingBot / 环境工具 + + + Batch Train/Eval Runner + 通用训练评测编排,不直接依赖 tau2 runtime + + + OpenViking Server :1933 + session.commit / memory extraction / experience update + + + + tau2 data + tasks / split / domain fixtures + + + POST /v1/cases/query + 返回通用 Case JSON + + + POST /v1/rollouts/execute + 执行 case batch + + + Tau2RolloutExecutor + VikingBot + tau2 env tools + 生成 Message / ToolPart + + + RubricEvaluation + reward / passed / feedback + + + + RemoteCaseLoader + 拉取 train/test cases + + + OfflinePolicyOptimizationPipeline + baseline_eval / train / final_eval + + + RemoteRolloutExecutor + HTTP 调用 tau2 rollout service + + + SessionCommitPolicyTrainer + CaseSpec + Rollout + OutcomeEvaluation + + + Report + accuracy / avg_reward / delta + + + + session.commit + 接收训练 rollout + + + SessionCompressorV3 + commit 后台抽取与训练入口 + + + TrajectoryRolloutAnalyzer + extract trajectories + + + ExperienceGradientEstimator + trajectory → gradient + + + PatchMergePolicyOptimizer + 全局 merge / split / delete plan + + + memories / experiences + PolicyStore / latest experience set + + + + + + + + + + + + + + + + + + + query cases + + + execute rollout + + + Rollout + evaluation + + + session.commit + + + read current OV memories via VikingBot / openviking tools + + + final_eval sees latest experiences + + + + + runner 编排调用 + + 训练写入链路 + + rollout 执行时读取当前 OpenViking memories + diff --git a/docs/design/assets/train-execution-details.png b/docs/design/assets/train-execution-details.png new file mode 100644 index 0000000000..8b9baeb367 Binary files /dev/null and b/docs/design/assets/train-execution-details.png differ diff --git a/docs/design/assets/train-execution-details.svg b/docs/design/assets/train-execution-details.svg new file mode 100644 index 0000000000..502e01c5e2 --- /dev/null +++ b/docs/design/assets/train-execution-details.svg @@ -0,0 +1,187 @@ + + OpenViking session.train 训练执行细节图 + 展示训练框架中的输入、并行 rollout、并行信号抽取、串行 policy update,以及 trajectories、experiences、remote service、session archive 等存储关系。 + + + + + + + + + + OpenViking session.train 训练执行细节 + Domain model 作为流程中的数据产物出现;图中标明哪些阶段并行、哪些阶段被 ExperienceSet.lock() 串行化,以及每个阶段读写哪些存储。 + + + + 1. Input / Rollout + parallel + + + CaseLoader + ListCaseLoader / RemoteCaseLoader + + + Case[] + + + PolicySnapshotter + snapshot current ExperienceSet + + + RolloutExecutor.execute + case batch can run concurrently + Remote executor / agent loop / simulator + + + + 2. Learning Signal Extraction + parallel + + + Rollout[] + + + RolloutAnalyzer.analyze + parallel per rollout · LLM ExtractLoop + uses rollout.evaluation / evaluator feedback + + + LLM + + + RolloutAnalysis[] + + + GradientEstimator.estimate + parallel per analysis · LLM ExtractLoop + proposes experience operations + + + LLM + + + + 3. Policy Update Transaction + serialized + + + SemanticGradient[] + + + ExperienceSet.lock() + wait indefinitely; serialize writers + + + ExperienceSet.reload() + read latest experiences under lock + + + PatchMergePolicyOptimizer.plan + serialized · LLM ExtractLoop + merge / split / delete via PatchMergeContextProvider + + + LLM + + + PolicyUpdatePlan + + + PolicyUpdater.apply + MemoryFilePolicyUpdater + upsert/delete with before_content guard + + + PolicyApplyResult + + + + Stores + + + Case Source + local list + remote service + + + Trajectory + memories/ + trajectories + + + Experience + memories/ + experiences + policy set + + + Session + archive + memory_diff + trace_id + + + + Realtime / Streaming path + + submit_rollout + requires Rollout.case + + + analyze + estimate + immediate per rollout + + + StreamingBatcher + max gradients / max wait + + submit_rollout 会等待当前 rollout 所在 batch 完成 flush + apply 后返回;flush 后进入同一个 serialized Policy Update Transaction。 + + + + + + + + + + + + + + + + + + + + case query/read + + write trajectory memory + + reload latest + + read required + candidates + + write/delete + + via session.commit + + + + + + flush batch + + + + eval bypasses analyzer/trainer and uses Rollout.evaluation only + + + + 并行边界:case rollout、rollout analysis、gradient estimation。LLM 调用:trajectory ExtractLoop、experience gradient ExtractLoop、patch merge ExtractLoop。串行边界:ExperienceSet.lock() 内的 reload → optimizer.plan → updater.apply。 + diff --git a/docs/design/assets/train-framework-overview.svg b/docs/design/assets/train-framework-overview.svg new file mode 100644 index 0000000000..1218c2a54e --- /dev/null +++ b/docs/design/assets/train-framework-overview.svg @@ -0,0 +1,188 @@ + + OpenViking session.train 框架模块关系图 + 展示 CaseLoader、RolloutExecutor、PolicyTrainer、RolloutAnalyzer、GradientEstimator、PolicyOptimizer、PolicyUpdater、ExperienceSet 和 StreamingPolicyTrainer 等模块之间的关系。 + + + + + + + + + + + + + + + + + + OpenViking session.train 框架模块关系 + 离线 batch、实时 streaming、远程 session.commit 三种训练入口共享同一套 rollout → trajectory → semantic gradient → policy update 语义。 + + + + Domain Model + + + Case + task + input + rubric + + + Rubric + what good means + + + Rollout + case + messages + eval + + + Trajectory + trainable trace + + + SemanticGradient + before/after MemoryFile + + + ExperienceSet + experiences directory + lock() + reload() + + + + OfflinePolicyOptimizationPipeline + + + CaseLoader + ListCaseLoader / RemoteCaseLoader + + + PolicySnapshotter + ContentHashPolicySnapshotter + + + RolloutExecutor + execute cases under snapshot + + + PolicyTrainer + Batch / Streaming / SessionCommit + train_rollouts(...) + + + RolloutAnalyzer + extract trajectory + + + GradientEstimator + trajectory → gradient + + + PolicyOptimizer + patch merge plan + + + PolicyUpdater + write/delete experiences + + + PipelineResult + accuracy / score / plan / apply + + + + Integration Paths + + + Batch train/eval + CaseLoader → executor + baseline/train/final eval + + + Realtime rollout + train_from_rollouts + StreamingPolicyTrainer + + + session.commit + CaseSpec + rollout + OutcomeEvaluation + + + Remote benchmark + HTTP cases / rollouts + tau2 service + + + Report + accuracy + avg_reward + + + + snapshot + + + Case[] + + + policy_snapshot_id + + + Rollout[] + + + + + + + + + + lock + reload before plan/apply + + current policy set + + + + + + + + + + + + + + + + + eval uses rollout.evaluation only + + + + 蓝色:rollout 执行 / eval;绿色:训练信号生成;橙色:policy plan/apply 写入;虚线:数据依赖或可选集成路径。 + diff --git a/docs/design/session-memory-extraction-flow.md b/docs/design/session-memory-extraction-flow.md index 5323704c0e..c951543de5 100644 --- a/docs/design/session-memory-extraction-flow.md +++ b/docs/design/session-memory-extraction-flow.md @@ -6,12 +6,13 @@ code-modification reference, so it avoids proposed or removed flows. ## Policy `memory_policy` carries target switches plus an optional global memory type -whitelist: +whitelist, and can disable per-archive Working Memory summaries: ```json { "self": { "enabled": true }, "peer": { "enabled": false }, + "working_memory": { "enabled": false }, "memory_types": ["profile", "preferences"] } ``` @@ -19,15 +20,23 @@ whitelist: When `memory_types` is omitted or `null`, all enabled schemas from `MemoryTypeRegistry` are allowed, including custom prompt/schema types. When it is set, extraction is limited to those names for both self and peer writes. +When `working_memory.enabled` is `false`, commit still archives messages and +runs configured memory extraction, but skips the archive summary. ## Memory Type Groups | Group | Types | Target | | --- | --- | --- | -| Long-term memory extraction | Enabled registry schemas without `agent_only` | Self and peer | +| Long-term memory extraction | Enabled registry schemas with `stage: user` | Self and peer | | Execution memory extraction | Execution-derived schemas, currently `trajectories`, `experiences` | Self only | | Session skills | `SESSION_SKILL_MEMORY_TYPE` output | Self only | +Memory schemas default to `stage: user` and `peer_enabled: true`. Set +`stage: agent` for schemas that are extracted only by the execution-memory +providers. Set `peer_enabled: false` for user-stage schemas that should ignore +`peer_id` and `ranges` peer targets and remain under the current user space +(for example `cases`). + Trajectory/experience extraction is controlled by `memory_types`: omitted or `null` allows both, while an explicit list must include those names. Session skill extraction also requires self memory to be enabled and only runs when the @@ -69,6 +78,7 @@ independently: | Unsafe `peer_id` | Skip | | Safe but unallowed `peer_id` | Skip | | `ranges` present | Read the message range; no-peer messages route to self, allowed peer messages route to peer | +| Schema has `peer_enabled: false` | Ignore `peer_id` and `ranges` peer targets; write self if self memory is enabled | | Only disabled targets found | Skip | The router does not rewrite message roles. A `role=user` message remains user diff --git a/docs/design/traj-exp-experience-learning-redesign.md b/docs/design/traj-exp-experience-learning-redesign.md new file mode 100644 index 0000000000..fa8b179689 --- /dev/null +++ b/docs/design/traj-exp-experience-learning-redesign.md @@ -0,0 +1,1291 @@ +# Trajectory / Experience 经验学习框架重构 +> +> 目标:把真实或离线环境中的 agent rollout 转换为 trajectory,再从 trajectory 估计 experience 更新信号,最终通过可审查、可合并、可并发安全的 policy update 机制更新 `experiences` 目录。 + +## 1. 总体定位 + +当前框架把 `experiences` 目录视为一个可优化的 **Experience Policy Set**: + +```text +viking://user//memories/experiences/ +``` + +目录中的每个 experience 文件是一个 `Experience`,整个目录共同构成 agent 的经验策略。训练框架不直接绑定某个 agent loop;它只约束以下抽象链路: + +```text +CaseLoader + -> RolloutExecutor + -> PolicyTrainer + -> RolloutAnalyzer + -> GradientEstimator + -> PolicyOptimizer + -> PolicyUpdater +``` + +其中 `PolicyTrainer` 是训练入口。默认本地实现会在进程内执行 `analyze -> estimate -> plan -> apply`;远程实现可以把 rollout 通过 `session.commit` 提交给 OpenViking 服务端,由服务端完成分析和训练。 + +### 1.1 训练执行细节图 + +OpenViking session.train 训练执行细节 + +这张图强调三个实现边界: + +- **并行边界**:case rollout、rollout analysis、gradient estimation 可以并行。 +- **串行边界**:`ExperienceSet.lock()` 内的 `reload -> PolicyOptimizer.plan -> PolicyUpdater.apply` 必须串行。 +- **存储边界**:trajectory 写入发生在 `RolloutAnalyzer`;experience 读取/合并/写入发生在 optimizer/updater;session archive 和 `memory_diff.json` 只出现在 `session.commit` 路径。 +- **LLM 边界**:红色特殊框表示该模块会调用 LLM / `ExtractLoop`,包括 trajectory 抽取、experience gradient 估计和 patch merge。 + + +## 2. 代码结构 + +当前模块结构: + +```text +openviking/session/train/ + context.py # PipelineContext / ExecutionContext + domain.py # domain dataclass + engine.py # PolicyTrainingEngine:共享 analyze/estimate/plan/apply 内核 + gradients.py # PatchSemanticGradient + interfaces.py # Protocol 接口 + pipeline.py # OfflinePolicyOptimizationPipeline + + components/ # 可替换组件实现 + case_loader.py + gradient_estimator.py + memory_store.py + policy_optimizer.py + policy_trainer.py + policy_updater.py + remote.py + rollout_executor.py + session_commit.py + snapshotter.py + trajectory_analyzer.py +``` + +设计边界: + +- 根目录保留框架内核、domain、接口和编排。 +- `components/` 放所有具体实现。 +- `openviking.session.train` 顶层继续导出常用类,便于外部使用。 + +## 3. 核心 Domain Model + +### 3.1 Experience / ExperienceSet + +`Experience` 对应 experiences 目录下的一个 experience 文件。 + +```python +@dataclass(slots=True) +class Experience: + name: str + uri: str + version: int + status: PolicyStatus + content: str + metadata: dict[str, Any] = field(default_factory=dict) + links: list[dict[str, Any]] = field(default_factory=list) + backlinks: list[dict[str, Any]] = field(default_factory=list) +``` + +`ExperienceSet` 是某个 experiences 根目录的快照: + +```python +@dataclass(slots=True) +class ExperienceSet: + root_uri: str + policies: list[Experience] + metadata: dict[str, Any] = field(default_factory=dict) + viking_fs: Any | None = field(default=None, repr=False, compare=False) + request_context: Any | None = field(default=None, repr=False, compare=False) +``` + +当前实现中,`ExperienceSet` 还负责提供并发安全能力: + +```python +async with policy_set.lock(): + latest_policy_set = await policy_set.reload() +``` + +约定: + +- `root_uri` 是 experiences 目录 URI。 +- `links/backlinks` 对应 memory file 中的 `MEMORY_FIELDS.links/backlinks`,用于在 train 域快照内保留 v2 link 协议数据。 +- `policies` 是当前目录下所有 experience 文件解析后的快照。 +- `viking_fs` / `request_context` 是运行时依赖,用于 `lock()` 和 `reload()`,不参与 equality/repr。 +- `PolicyTrainingEngine.plan_and_apply(...)` 会先加 policy tree lock,再 reload 最新 policy set,然后 plan/apply。 + +### 3.2 Trajectory + +`Trajectory` 是从 rollout 中抽取并持久化的可训练轨迹样本,对应 trajectories 目录下的 memory 文件。 + +```python +@dataclass(slots=True) +class Trajectory: + name: str + uri: str + content: str + outcome: TrajectoryOutcome | str + retrieval_anchor: str + metadata: dict[str, Any] = field(default_factory=dict) +``` + +约定: + +- `Rollout` 是原始执行记录。 +- `Trajectory` 是从 rollout messages 中抽取出的训练样本。 +- trajectory 文件由 `TrajectoryRolloutAnalyzer` 通过 `ExtractLoop + MemoryUpdater` 写入 `memories/trajectories`。 + +### 3.3 Case / Rubric + +`Case` 是可执行、可复现、可评估的训练/评测样例。 + +```python +@dataclass(slots=True) +class Case: + name: str + task_signature: str + input: dict[str, Any] + rubric: Rubric + metadata: dict[str, Any] = field(default_factory=dict) +``` + +`Rubric` 定义“什么叫做好”和“怎么检查”。当前不再保留独立 `Outcome` 概念。 + +```python +@dataclass(slots=True) +class Rubric: + name: str + description: str + criteria: list[RubricCriterion] + metadata: dict[str, Any] = field(default_factory=dict) + +@dataclass(slots=True) +class RubricCriterion: + name: str + description: str + required: bool + weight: float + metadata: dict[str, Any] = field(default_factory=dict) +``` + +### 3.4 Rollout + +`Rollout` 是某个 policy snapshot 在某个 case 上执行后的记录。 + +```python +@dataclass(slots=True) +class Rollout: + case: Case + messages: list[Message] + policy_snapshot_id: str + evaluation: RubricEvaluation | None = None + metadata: dict[str, Any] = field(default_factory=dict) +``` + +当前关键变化:`Rollout.evaluation` 是一等可选字段。 + +- 如果环境本身能给 reward / evaluation,`RolloutExecutor` 应直接填入 `rollout.evaluation`。 +- 训练时 `TrajectoryRolloutAnalyzer` 优先沿用 `rollout.evaluation`;没有时才通过注入的 `RolloutEvaluator` 评估;再没有时用“是否抽取到 trajectory”作为 fallback evaluation。 +- `pipeline.eval(...)` 不再调用 `RolloutAnalyzer`,只依赖 `RolloutExecutor` 返回的 `rollout.evaluation`;如果 eval rollout 缺 evaluation,会直接报错。 + +### 3.5 RubricEvaluation + +```python +@dataclass(slots=True) +class RubricEvaluation: + passed: bool + score: float + criterion_results: list[CriterionResult] + feedback: list[str] + metadata: dict[str, Any] = field(default_factory=dict) + +@dataclass(slots=True) +class CriterionResult: + criterion_name: str + passed: bool + score: float + feedback: list[str] + evidence: list[str] + metadata: dict[str, Any] = field(default_factory=dict) +``` + +在 tau2 集成中: + +- `passed = reward >= 1.0` +- `score = reward` +- report 展示以 `accuracy = passed_count / case_count` 为主,`average_reward` 为辅助指标。 + +## 4. SemanticGradient + +`SemanticGradient` 是针对一个目标 experience 的语义更新信号。当前接口以 `MemoryFile` before/after 表达,而不是文本 patch 对象。 + +```python +class SemanticGradient(Protocol): + @property + def before_file(self) -> MemoryFile | None: ... + + @property + def after_file(self) -> MemoryFile: ... + + @property + def target_experience_name(self) -> str: ... + @property + def target_experience_uri(self) -> str | None: ... + @property + def base_version(self) -> int | None: ... + @property + def rationale(self) -> str: ... + @property + def links(self) -> list[StoredLink]: ... + @property + def confidence(self) -> float: ... + @property + def metadata(self) -> dict[str, Any]: ... +``` + +当前具体实现: + +```python +@dataclass(slots=True) +class PatchSemanticGradient: + before_file: MemoryFile | None + after_file: MemoryFile + base_version: int | None + rationale: str + links: list[StoredLink] + confidence: float + metadata: dict[str, Any] = field(default_factory=dict) +``` + +约定: + +- `before_file is None` 表示建议新建。 +- `after_file` 是建议的目标 memory file 状态。 +- `links` 承载 exp→traj 的 provenance,沿用 v2 `MEMORY_FIELDS.links/backlinks` 协议;来源轨迹关系使用 `StoredLink(from_uri=exp_uri, to_uri=traj_uri, link_type="derived_from", weight=1.0)`,不再引入单独的轨迹 URI 列表字段。 +- patch 文本不是 gradient 自身字段,而是由 `PatchMergeContextProvider` 在 merge 阶段把 before/after memory file 渲染为字段级 unified diff。 + +## 5. PolicyUpdatePlan / PolicyUpdater + +`PolicyOptimizer.plan(...)` 输出 `PolicyUpdatePlan`,`PolicyUpdater.apply(...)` 负责真正写文件。 + +```python +PolicyPlanItemKind = Literal["upsert_experience", "delete_experience"] + +@dataclass(slots=True) +class PolicyPlanItem: + kind: PolicyPlanItemKind + target_experience_name: str + target_experience_uri: str | None + before_content: str | None + after_content: str | None + base_version: int | None = None + confidence: float | None = None + links: list[StoredLink] = field(default_factory=list) + metadata: dict[str, Any] = field(default_factory=dict) + +@dataclass(slots=True) +class PolicyUpdatePlan: + items: list[PolicyPlanItem] = field(default_factory=list) + metadata: dict[str, Any] = field(default_factory=dict) + +@dataclass(slots=True) +class PolicyApplyResult: + updated_policy_set: ExperienceSet + written_uris: list[str] = field(default_factory=list) + deleted_uris: list[str] = field(default_factory=list) + errors: list[str] = field(default_factory=list) + metadata: dict[str, Any] = field(default_factory=dict) +``` + +当前 `MemoryFilePolicyUpdater` 支持: + +- `upsert_experience` +- `delete_experience` +- 基于 `before_content` 的轻量 base-content guard,避免覆盖已发散内容。 + +## 6. 接口定义 + +### 6.1 CaseLoader + +```python +class CaseLoader(Protocol): + async def batches(self, context: Any) -> AsyncIterator[list[Case]]: ... +``` + +实现: + +- `ListCaseLoader` +- `RemoteCaseLoader`:通过 HTTP 服务拉取 cases。 + +### 6.2 RolloutExecutor + +```python +class RolloutExecutor(Protocol): + async def execute( + self, + cases: list[Case], + policy_set: ExperienceSet, + context: ExecutionContext, + ) -> list[Rollout]: ... +``` + +实现: + +- `SingleTurnLLMRolloutExecutor` +- `RemoteRolloutExecutor` +- `Tau2RolloutExecutor`(benchmark/tau2 内部实现,通过 tau2 service 暴露给训练流程) + +### 6.3 RolloutEvaluator + +```python +class RolloutEvaluator(Protocol): + async def evaluate(self, rollout: Rollout, context: Any) -> RubricEvaluation: ... +``` + +用途:环境不能直接提供 `rollout.evaluation` 时,`RolloutAnalyzer` 可注入 evaluator 进行评估。 + +### 6.4 RolloutAnalyzer + +```python +class RolloutAnalyzer(Protocol): + async def analyze(self, rollout: Rollout, context: Any) -> RolloutAnalysis: ... +``` + +当前实现:`TrajectoryRolloutAnalyzer`。 + +职责: + +1. 确定 rollout evaluation: + - 优先使用 `rollout.evaluation` + - 否则使用注入的 `RolloutEvaluator` + - 否则基于是否抽取到 trajectory 生成默认 evaluation +2. 将 evaluation feedback 追加到 trajectory extraction messages。 +3. 通过 `AgentTrajectoryContextProvider + ExtractLoop` 只抽取 `trajectories` memory type。 +4. 通过 `MemoryUpdater.apply_operations(...)` 写入 trajectory memory。 +5. 读取写入的 trajectory 文件并返回 `RolloutAnalysis`。 + +### 6.5 GradientEstimator + +```python +class GradientEstimator(Protocol): + async def estimate( + self, + analysis: RolloutAnalysis, + experience_set: ExperienceSet, + context: Any, + ) -> list[SemanticGradient]: ... +``` + +当前实现:`ExperienceGradientEstimator`。 + +它复用: + +- `AgentExperienceContextProvider` +- `ExtractLoop` +- `MemoryIsolationHandler(allowed_memory_types={"experiences"})` + +但不调用 `MemoryUpdater.apply_operations(...)`。它把 ExtractLoop 产生的 upsert operations 转成 `PatchSemanticGradient`。 + +### 6.6 PolicyOptimizer + +```python +class PolicyOptimizer(Protocol): + async def plan( + self, + gradients: list[SemanticGradient], + policy_set: ExperienceSet, + context: Any, + ) -> PolicyUpdatePlan: ... +``` + +当前实现:`PatchMergePolicyOptimizer`。 + +它不按 target 分组限制输出,而是把一批 gradients 一次性交给 `PatchMergeContextProvider + ExtractLoop` 进行全局 merge。LLM 可以: + +- 合并多个 patch 到一个 experience。 +- 把一个臃肿 patch 拆成多个 experience。 +- 合并相似新文件。 +- 主动输出删除操作。 + +### 6.7 PolicyUpdater + +```python +class PolicyUpdater(Protocol): + async def apply( + self, + plan: PolicyUpdatePlan, + policy_set: ExperienceSet, + context: Any, + ) -> PolicyApplyResult: ... +``` + +实现: + +- `DryRunPolicyUpdater` +- `MemoryFilePolicyUpdater` + +### 6.8 PolicyTrainer + +```python +class PolicyTrainer(Protocol): + async def train_rollouts( + self, + rollouts: list[Rollout], + policy_set: ExperienceSet, + context: Any, + analyses: list[RolloutAnalysis] | None = None, + ) -> RolloutTrainingResult: ... +``` + +实现: + +- `BatchPolicyTrainer`:显式 batch,本地执行 analyze/estimate/plan/apply。 +- `StreamingPolicyTrainer`:实时 rollout 输入,先 analyze/estimate,再按梯度数量和时间窗口攒批,批量 plan/apply。 +- `SessionCommitPolicyTrainer`:把 rollout 写入远端 OpenViking session,通过 `session.commit` 让服务端完成训练。 + +### 6.9 PolicyOptimizationPipeline + +```python +class PolicyOptimizationPipeline(Protocol): + async def train(...) -> PipelineResult: ... + async def eval(...) -> PipelineEvaluationResult: ... + async def train_from_rollouts(...) -> RolloutTrainingResult: ... +``` + +当前实现:`OfflinePolicyOptimizationPipeline`。 + +## 7. PipelineContext / ExecutionContext + +```python +@dataclass(slots=True) +class PipelineContext: + case_load_context: Any = None + snapshot_context: Any = None + analysis_context: Any = None + gradient_context: Any = None + optimization_context: Any = None + apply_context: Any = None + execution_metadata: dict[str, Any] = field(default_factory=dict) + max_epochs: int = 1 + +@dataclass(slots=True) +class ExecutionContext: + policy_snapshot_id: str + metadata: dict[str, Any] = field(default_factory=dict) +``` + +`max_epochs` 是训练迭代次数。之前文档中的 `max_iterations` 已改为 epoch 概念。 + +## 8. 训练流程 + +### 8.1 OfflinePolicyOptimizationPipeline.train + +```text +for epoch in range(ctx.max_epochs): + for cases in case_loader.batches(...): + snapshot_id = snapshotter.snapshot(policy_set) + rollouts = rollout_executor.execute(cases, policy_set, ExecutionContext(snapshot_id)) + training_result = policy_trainer.train_rollouts(rollouts, policy_set, ctx) + policy_set = training_result.apply_result.updated_policy_set +``` + +默认 `policy_trainer` 是 `BatchPolicyTrainer`,因此本地训练链路为: + +```text +Rollout[] + -> RolloutAnalyzer.analyze(...) + -> GradientEstimator.estimate(...) + -> PolicyTrainingEngine.plan_and_apply(...) + -> async with ExperienceSet.lock() + -> ExperienceSet.reload() + -> PolicyOptimizer.plan(...) + -> PolicyUpdater.apply(...) +``` + +### 8.2 OfflinePolicyOptimizationPipeline.eval + +```text +CaseLoader -> RolloutExecutor -> Rollout.evaluation -> PipelineEvaluationResult +``` + +eval 阶段不会调用 `RolloutAnalyzer`,不会抽 trajectory,也不会写 policy。它要求 `RolloutExecutor` 返回带 `evaluation` 的 rollout。 + +### 8.3 train_from_rollouts + +实时场景或外部系统已经产生 rollout 时,可以绕过 `CaseLoader / PolicySnapshotter / RolloutExecutor`: + +```text +Rollout[] -> policy_trainer.train_rollouts(...) +``` + +约束:每个 rollout 必须包含 `case`。 + +## 9. Batch 与 Streaming + +### 9.1 BatchPolicyTrainer + +适合离线训练,输入一批 rollout 后直接完成一次: + +```text +analyze -> estimate -> plan -> apply +``` + +### 9.2 StreamingPolicyTrainer + +适合实时 commit / 并发 rollout 场景。 + +流程: + +```text +submit_rollout(rollout) + -> analyze rollout + -> estimate gradients + -> submit gradients to StreamingBatcher + -> 等待该 rollout 所在 batch 被 flush 并 apply +``` + +flush 触发条件: + +- `max_gradients_per_update` 达到阈值 +- 最老 gradient 等待超过 `max_wait_seconds` +- `close()` 时 flush 剩余内容 + +默认配置: + +```python +@dataclass(slots=True) +class StreamingPolicyTrainerConfig: + max_gradients_per_update: int = 8 + max_wait_seconds: float = 10.0 + timer_check_interval_seconds: float = 1.0 + trace_console: bool = False +``` + +进程内全局共享: + +```python +get_streaming_policy_trainer(...) +make_streaming_policy_trainer_key(policy_root_uri, request_context) +``` + +并发安全由 `PolicyTrainingEngine.plan_and_apply(...)` 中的 `ExperienceSet.lock()` 保证。 + +## 10. Patch Merge 机制 + +### 10.1 PatchSemanticGradient 到 PatchMergePatch + +`PatchMergePolicyOptimizer` 会把每个 `SemanticGradient` 转为: + +```python +@dataclass(slots=True) +class PatchMergePatch: + before_file: MemoryFile | None + after_file: MemoryFile + metadata: dict[str, Any] +``` + +### 10.2 PatchMergeContextProvider + +位置:`openviking/session/memory/patch_merge_context_provider.py` + +职责: + +- 给 LLM 提供待合并 patch 相关的原始 memory 文件。 +- 将 `MemoryFile` before/after 渲染为字段级 unified diff。 +- 用 embedding 检索额外候选文件,帮助发现相似/重复 memory。 +- 暴露指定 memory type 的 schema,让 ExtractLoop 输出合法 memory operations。 + +输入文件选择: + +```text +required_file_uris = patch target uri / superseded policy uri +extra_candidate_files = embedding search 当前 memory_type 下的相似文件 +max_extra_candidate_files = max(5, len(required_file_uris)) +search_limit = max_extra_candidate_files * 2 +``` + +字段 diff 规则: + +- 只展示发生变化的字段。 +- 字符串按行 diff。 +- dict/list 先 JSON 格式化再 diff。 +- `content` 已在 `Field Diff: content` 中展示,因此不会额外在 metadata 中重复塞完整 content。 + +### 10.3 PatchMergePolicyOptimizer + +```text +SemanticGradient[] + -> PatchMergeContextProvider.prefetch() + -> ExtractLoop(max_iterations=1) + -> ResolvedOperations + -> PolicyPlanItem[] +``` + +输出支持: + +- upsert experience +- delete experience + +merge 输入/输出日志通过 `tracer.info(..., console=False)` 记录,避免默认污染 console。 + +## 11. session.commit 实时训练接入 + +`SessionCompressorV3` 已把用户记忆抽取和实时训练接起来。 + +### 11.1 用户记忆抽取 + +`SessionCompressorV3._extract_user_memories(...)`: + +1. 通过原用户记忆 `ExtractLoop` 抽取用户记忆。 +2. case 不再额外单独调用 LLM,而是作为一种普通 memory type:`cases`。 +3. 抽取结果交给 `StreamingMemoryUpdater` 做 patch merge 写入用户记忆。 +4. 如有 `archive_uri`,写入 `memory_diff.json`,其中包含顶层 `trace_id`。 + +`memory_diff.json` 顶层结构包含: + +```json +{ + "archive_uri": "...", + "trace_id": "...", + "extracted_at": "...", + "operations": {...}, + "summary": {...} +} +``` + +### 11.2 从 cases 触发 streaming train + +`SessionCompressorV3.train_from_extracted_cases(...)`: + +```text +extracted Case[] + original commit messages + -> Rollout(case, messages, policy_snapshot_id=session-commit:...) + -> StreamingPolicyTrainer.submit_rollout(...) +``` + +即真实 session.commit 产生的对话可以被转为 rollout 输入训练框架。 + +## 12. SessionCommitPolicyTrainer:远程服务端训练 + +`SessionCommitPolicyTrainer` 是一个 `PolicyTrainer` 实现,用于“训练框架在外部,OpenViking 服务端负责训练”的场景。 + +它会把 rollout 写成一个临时 session: + +```text +[CaseSpec message] +[Rollout messages] +[OutcomeEvaluation message] +``` + +其中: + +- `CaseSpec` 放在开头,只含 case/rubric/task context,不含 evaluation。 +- `OutcomeEvaluation` 放在最后,只含 evaluation,作为训练信号。 +- rollout 的工具结果会通过 `ToolPart` 的 `tool_output` 上传,而不是普通 text。 + +然后执行: + +```text +client.create_session(...) +client.batch_add_messages(...) +client.commit_session(...) +client.get_task(...) until completed/failed/timeout +``` + +CaseSpec 会做精简,避免传入巨大或重复字段: + +- 不传 `policy` +- 不传 `data_root` +- 不传 `rollout_metadata` +- 不传 `policy_snapshot_id` +- 保留 `domain/split/data_split/task_id/task_no/user_query/ground_truth/rubric` + +## 13. Remote HTTP 组件 + +`components/remote.py` 提供通用 HTTP 组件: + +- `RemoteCaseLoader` +- `RemoteRolloutExecutor` + +它们面向一个环境/benchmark service: + +```text +POST /v1/cases/query +POST /v1/rollouts/execute +GET /v1/rollouts/executions/{execution_id} +``` + +其中 `/v1/rollouts/execute` 只负责提交单个 case 的 rollout execution,返回 +`execution_id`;`RemoteRolloutExecutor` 会并发提交多个 case,并通过 +`/v1/rollouts/executions/{execution_id}` 轮询状态。这样长耗时 rollout 不会占用 +一个超长 HTTP request,也便于未来 benchmark service 做多机部署和负载均衡。 + +这样训练框架不需要直接依赖 tau2 或其他 benchmark 的代码,只依赖通用 +Case/Rollout JSON 协议。 + +## 14. tau2 集成 + +### 14.1 架构 + +当前 tau2 训练分为两个进程: + +```text +tau2 service + - 依赖 tau2 / vikingbot + - 暴露 case query 和 rollout execute HTTP API + +train/eval runner + - 使用 RemoteCaseLoader / RemoteRolloutExecutor + - 使用 SessionCommitPolicyTrainer 提交 OpenViking session.commit + - 本身不直接依赖 tau2 runtime +``` + +### 14.2 tau2 service + +位置: + +```text +benchmark/tau2/train/service_app.py +benchmark/tau2/train/run_service.sh +``` + +启动: + +```bash +benchmark/tau2/train/run_service.sh \ + --host 127.0.0.1 \ + --port 1944 +``` + +### 14.3 remote train/eval + +位置: + +```text +benchmark/tau2/train/run_batch_train_eval.sh +openviking/session/train/run_batch_train_eval.py +openviking/session/train/batch_runner.py +``` + +预先只跑 test 分数(不训练): + +```bash +benchmark/tau2/train/run_batch_train_eval.sh \ + --epochs 0 \ + --eval-index 24 \ + --trials 8 +``` + +训练前先跑一次 test baseline,再训练并跑最终 test: + +```bash +benchmark/tau2/train/run_batch_train_eval.sh \ + --baseline-eval \ + --epochs 4 \ + --trials 8 +``` + +输出以 accuracy 为主,阶段日志由 session/train lifecycle hooks 统一输出: + +```text +[baseline_rollout] epoch=-1 trials=8 cases_per_trial=25 total_rollouts=200 accuracy=... ± ... avg_reward=... ± ... +================= epoch 0 ================= +[train_rollout] epoch=0 cases=25 accuracy=... passed=... avg_reward=... +[train] epoch=0 commits=25 errors=0 +[final_test_rollout] epoch=4 trials=8 cases_per_trial=25 total_rollouts=200 accuracy=... ± ... avg_reward=... ± ... +``` + +### 14.4 tau2 rollout messages + +`Tau2RolloutExecutor` 会把工具结果转成真正的 `ToolPart`: + +```json +{ + "type": "tool", + "tool_id": "tau2-tool-0", + "tool_name": "get_reservation_details", + "tool_input": {...}, + "tool_output": "...", + "tool_status": "completed" +} +``` + +这样上传到 `session.commit` 后,服务端可以复用已有 tool output 外部化和 memory extraction 逻辑。 + + +## 15. tau2 接入新评测框架示意图 + +tau2 的接入方式体现了推荐的 benchmark 集成模式:benchmark runtime 独立成 HTTP service,训练框架只通过通用 `RemoteCaseLoader` / `RemoteRolloutExecutor` 接入。 + +tau2 接入 OpenViking 新训练评测框架 + + +图中需要特别注意:tau2 runtime service 虽然不负责训练写入,但它执行 rollout 时会通过 VikingBot / OpenViking tools 读取当前 OpenViking memories。因此 final_eval 能看到 train epoch 后写入的最新 experiences。 + +### 15.1 接入分层 + +```text +tau2 service + - 依赖 tau2 / vikingbot + - 负责 case 查询、rollout 执行、环境 reward 评估 + - 输出通用 Case / Rollout / RubricEvaluation JSON + +train/eval runner + - 不直接依赖 tau2 runtime + - 使用 RemoteCaseLoader 查询 case + - 使用 RemoteRolloutExecutor 执行 rollout + - 使用 SessionCommitPolicyTrainer 把训练 rollout 提交给 OpenViking 服务端 + +OpenViking server + - 通过 session.commit 接收 rollout messages + - 服务端内部执行 trajectory extraction / gradient estimation / patch merge / policy update +``` + +### 15.2 train/eval 时序 + +```text +baseline_eval: + RemoteCaseLoader(test) + -> RemoteRolloutExecutor + -> Tau2RolloutExecutor + -> rollout.evaluation + -> accuracy / avg_reward report + +train epoch: + RemoteCaseLoader(train) + -> RemoteRolloutExecutor + -> Tau2RolloutExecutor + -> SessionCommitPolicyTrainer + -> session.commit + -> SessionCompressorV3 + -> StreamingPolicyTrainer + -> experiences update + +final_eval: + RemoteCaseLoader(test) + -> RemoteRolloutExecutor + -> Tau2RolloutExecutor reads latest OpenViking experiences + -> rollout.evaluation + -> accuracy delta report +``` + +### 15.3 为什么 eval 不走 RolloutAnalyzer + +在 tau2 场景中,环境执行完 rollout 后可以直接给出 reward,因此 `Tau2RolloutExecutor` 会返回: + +```python +Rollout( + case=case, + messages=messages, + policy_snapshot_id=snapshot_id, + evaluation=RubricEvaluation(...), +) +``` + +所以 `OfflinePolicyOptimizationPipeline.eval(...)` 只统计 `rollout.evaluation`: + +```text +accuracy = passed_count / case_count +average_reward = mean(evaluation.score) +``` + +eval 不抽 trajectory、不估计 gradient、不写 experience。 + +### 15.4 训练如何通过 session.commit 进入服务端 + +`SessionCommitPolicyTrainer` 会把 rollout 转成临时 session messages: + +```text +[OpenViking Training CaseSpec] +[Rollout messages: user / assistant / ToolPart] +[OpenViking OutcomeEvaluation] +``` + +其中: + +- `CaseSpec` 放在开头,只描述任务和 rubric,不包含 evaluation。 +- `OutcomeEvaluation` 放在最后,作为训练信号。 +- tau2 工具结果使用 `ToolPart.tool_output` 上传,服务端可以复用已有 tool output 外部化和 memory extraction 逻辑。 + +### 15.5 指标展示 + +tau2 runner 的报告以正确率为主: + +```text +[baseline_eval] epoch=-1 cases=10 accuracy=20.00% passed=2/10 avg_reward=0.200000 +[train_epoch] epoch=0 cases=50 accuracy=18.00% passed=9/50 avg_reward=0.180000 commits=50 errors=0 +[final_eval] epoch=1 cases=10 accuracy=30.00% passed=3/10 avg_reward=0.300000 + +baseline accuracy: 20.00% (2/10) +final accuracy: 30.00% (3/10) +accuracy delta: +10.00pp +``` + +`average_reward` 保留为辅助指标;主指标是 `accuracy`。 + +### 15.6 以 tau2 为例:新场景接入需要实现的接口 + +一个新的 benchmark / domain / environment 接入训练评测框架时,推荐复用 tau2 +的分层方式:把场景 runtime 独立成一个 HTTP service,训练进程继续使用通用 +`RemoteCaseLoader` / `RemoteRolloutExecutor`。训练框架不关心场景内部怎么启动 +agent、怎么调用工具、怎么计算 reward,只要求 service 实现下面这些协议。 + +#### 15.6.1 Case 查询接口 + +```text +POST /v1/cases/query +``` + +请求: + +```json +{ + "dataset": "tau2", + "domain": "airline", + "split": "train", + "cursor": null, + "limit": 100, + "filters": {} +} +``` + +响应: + +```json +{ + "cases": [ + { + "name": "tau2_airline_train_0", + "task_signature": "tau2:airline:train:0", + "input": { + "domain": "airline", + "split": "train", + "task_id": "0", + "task_no": 0, + "user_query": "...", + "ground_truth": "..." + }, + "rubric": { + "name": "tau2_airline_train_0_rubric", + "description": "...", + "criteria": [ + { + "name": "tau2_reward", + "description": "The tau2 environment reward is 1.0.", + "required": true, + "weight": 1.0, + "metadata": {} + } + ], + "metadata": {} + }, + "metadata": { + "dataset": "tau2", + "domain": "airline", + "source": "tau2", + "split": "train" + } + } + ], + "next_cursor": "100" +} +``` + +接入要求: + +- `dataset/domain/split` 用于定位数据集切片。 +- `cursor/limit` 用于分页;没有下一页时 `next_cursor = null`。 +- `Case.input` 只放 rollout 必需的任务输入和场景元信息,不要塞训练框架已经能从 + 上下文拿到的内容,例如完整 system prompt、完整 rollout metadata、evaluation + 结果或 policy snapshot。 +- `Case.rubric` 必须能描述评测目标;如果环境能直接给 reward,也仍然要提供 + rubric,便于训练侧把 reward 转成统一的 `RubricEvaluation`。 + +tau2 中对应实现是: + +```text +benchmark/tau2/train/service_app.py::query_cases +benchmark/tau2/train/case_loader.py::Tau2CaseLoader +``` + +#### 15.6.2 Rollout 提交接口 + +```text +POST /v1/rollouts/execute +``` + +请求: + +```json +{ + "case": { "...": "Case JSON" }, + "policy_set": { + "root_uri": "viking://user/default/memories/experiences", + "policies": [], + "metadata": {} + }, + "execution_context": { + "policy_snapshot_id": "tau2-policy-snapshot:...", + "metadata": { + "epoch": 0, + "training": true + } + }, + "options": { + "config_path": "/path/to/ov.conf", + "max_iterations": 30, + "keep_default_tools": true, + "rollout_language": "default" + } +} +``` + +响应: + +```json +{ + "execution_id": "rollout_exec_...", + "status": "running", + "case_name": "tau2_airline_train_0", + "created_at": 1781097747.0, + "updated_at": 1781097747.0, + "error": null +} +``` + +接入要求: + +- 该接口只提交一个 case 的 rollout execution,不需要同步等待 rollout 完成。 +- 客户端会对多个 case 发起多个请求,service 端可以自行排队、限流、调度到不同 + worker 或机器。 +- `policy_set.root_uri` 告诉 runtime 当前 experiences 根目录;tau2 rollout 期间 + VikingBot 会通过 OpenViking recall 读取这里的最新经验。 +- `execution_context.policy_snapshot_id` 必须原样写入返回的 `Rollout.policy_snapshot_id`, + 用于追踪这次 rollout 使用的是哪次 policy snapshot。 + +tau2 中对应实现是: + +```text +benchmark/tau2/train/service_app.py::execute_rollout +benchmark/tau2/train/service_app.py::_run_rollout_execution +benchmark/tau2/train/rollout_executor.py::Tau2RolloutExecutor +``` + +#### 15.6.3 Rollout 状态轮询接口 + +```text +GET /v1/rollouts/executions/{execution_id} +``` + +运行中响应: + +```json +{ + "execution_id": "rollout_exec_...", + "status": "running", + "case_name": "tau2_airline_train_0", + "created_at": 1781097747.0, + "updated_at": 1781097750.0, + "error": null +} +``` + +完成响应: + +```json +{ + "execution_id": "rollout_exec_...", + "status": "completed", + "case_name": "tau2_airline_train_0", + "created_at": 1781097747.0, + "updated_at": 1781097760.0, + "error": null, + "rollout": { + "case": { "...": "Case JSON" }, + "messages": [ + { + "role": "user", + "parts": [ + { + "type": "text", + "text": "..." + } + ] + }, + { + "role": "assistant", + "parts": [ + { + "type": "tool", + "tool_id": "tau2-tool-0", + "tool_name": "get_reservation_details", + "tool_input": {"reservation_id": "EHGLP3"}, + "tool_output": "...", + "tool_status": "completed" + } + ] + } + ], + "policy_snapshot_id": "tau2-policy-snapshot:...", + "evaluation": { + "passed": false, + "score": 0.0, + "criterion_results": [ + { + "criterion_name": "tau2_reward", + "passed": false, + "score": 0.0, + "feedback": ["tau2 environment reward is below 1.0."], + "evidence": [], + "metadata": {"reward": 0.0} + } + ], + "feedback": ["tau2 environment reward is below 1.0."], + "metadata": { + "source": "tau2_executor", + "reward": 0.0 + } + }, + "metadata": { + "memory": "...", + "tools_used": [], + "iterations": 6 + } + } +} +``` + +失败响应: + +```json +{ + "execution_id": "rollout_exec_...", + "status": "failed", + "case_name": "tau2_airline_train_0", + "created_at": 1781097747.0, + "updated_at": 1781097752.0, + "error": "..." +} +``` + +接入要求: + +- `status` 至少支持 `running/completed/failed`。 +- `completed` 时必须返回完整 `rollout`。 +- `failed` 时必须返回可读 `error`,训练侧会把它归入该 case 的 rollout 失败。 +- `Rollout.messages` 应使用 OpenViking `Message` / `Part` 结构;工具调用和工具结果 + 用 `ToolPart`,不要把 `tool-call:\nname: ...` 塞进普通 text content。 +- `Rollout.evaluation` 在 eval 阶段是必需字段;如果没有 evaluation, + `OfflinePolicyOptimizationPipeline.eval(...)` 会失败。 + +#### 15.6.4 RolloutExecutor 内部职责 + +新场景自己的 rollout executor 需要完成这些事情: + +1. 根据 `Case.input` 初始化环境和用户模拟器。 +2. 根据 `policy_set.root_uri` / OpenViking 配置让 agent 读取当前 experiences。 +3. 执行 agent loop,记录 user/assistant/tool messages。 +4. 把环境 reward 或 judge 结果转成 `RubricEvaluation`。 +5. 返回统一 `Rollout`: + +```python +Rollout( + case=case, + messages=messages, + policy_snapshot_id=context.policy_snapshot_id, + evaluation=RubricEvaluation(...), + metadata={ + "tools_used": [...], + "iterations": ..., + "memory": "...", + }, +) +``` + +tau2 的 `Tau2RolloutExecutor` 就是这个适配层:它一侧依赖 tau2/VikingBot runtime, +另一侧只输出训练框架理解的 `Rollout`。 + +#### 15.6.5 最小接入清单 + +接入一个新场景,最少需要实现: + +| 接口/组件 | 必需 | 作用 | +|---|---:|---| +| `POST /v1/cases/query` | 是 | 分页返回 `Case[]` | +| `POST /v1/rollouts/execute` | 是 | 提交单个 rollout execution | +| `GET /v1/rollouts/executions/{execution_id}` | 是 | 轮询 rollout 状态并取回 `Rollout` | +| `RubricEvaluation` 转换 | eval 必需 | 把场景 reward/judge 结果转成统一 evaluation | +| `Message` / `ToolPart` 转换 | 训练必需 | 保留 agent 行为和工具证据,供 session.commit 抽取 trajectory/experience | +| `GET /health` | 建议 | 方便 runner 或部署系统做 preflight | + +如果新场景不想提供 HTTP service,也可以在同进程内直接实现 +`CaseLoader` / `RolloutExecutor` Protocol;但跨进程、多机或重 runtime 依赖的场景, +推荐采用 tau2 这种 service 方式。 + +## 16. 当前主要组件清单 + +| 组件 | 文件 | 说明 | +|---|---|---| +| `OfflinePolicyOptimizationPipeline` | `pipeline.py` | 离线 train/eval 编排 | +| `PolicyTrainingEngine` | `engine.py` | 共享 analyze/estimate/plan/apply 内核 | +| `ListCaseLoader` | `components/case_loader.py` | 内存 case loader | +| `RemoteCaseLoader` | `components/remote.py` | HTTP case loader | +| `RemoteRolloutExecutor` | `components/remote.py` | HTTP rollout executor | +| `SingleTurnLLMRolloutExecutor` | `components/rollout_executor.py` | 简单单轮 LLM rollout | +| `TrajectoryRolloutAnalyzer` | `components/trajectory_analyzer.py` | 抽取 trajectory memory | +| `ExperienceGradientEstimator` | `components/gradient_estimator.py` | trajectory -> PatchSemanticGradient | +| `PatchMergePolicyOptimizer` | `components/policy_optimizer.py` | 多 gradient 全局 merge | +| `DryRunPolicyUpdater` | `components/policy_updater.py` | dry-run apply | +| `MemoryFilePolicyUpdater` | `components/policy_updater.py` | VikingFS 写回 experiences | +| `BatchPolicyTrainer` | `components/policy_trainer.py` | batch rollout 训练 | +| `StreamingPolicyTrainer` | `components/policy_trainer.py` | 实时攒批训练 | +| `SessionCommitPolicyTrainer` | `components/session_commit.py` | 通过 session.commit 远程训练 | +| `ContentHashPolicySnapshotter` | `components/snapshotter.py` | 内容 hash snapshot id | +| `ExperienceSetLoader` | `components/memory_store.py` | 从 experiences 目录加载 policy set | + +## 17. 端到端本地训练伪代码 + +```python +policy_set = await ExperienceSetLoader(viking_fs).load( + "viking://user/default/memories/experiences", + ctx=request_context, +) + +pipeline = OfflinePolicyOptimizationPipeline( + snapshotter=ContentHashPolicySnapshotter(), + rollout_executor=SomeRolloutExecutor(), + rollout_analyzer=TrajectoryRolloutAnalyzer(viking_fs=viking_fs, vikingdb=vikingdb), + gradient_estimator=ExperienceGradientEstimator(viking_fs=viking_fs), + policy_optimizer=PatchMergePolicyOptimizer(viking_fs=viking_fs), + policy_updater=MemoryFilePolicyUpdater(viking_fs=viking_fs), +) + +result = await pipeline.train( + case_loader=ListCaseLoader(cases, batch_size=8), + policy_set=policy_set, + context=PipelineContext( + max_epochs=1, + analysis_context=TrajectoryAnalyzerContext(request_context=request_context), + gradient_context=ExperienceGradientContext( + request_context=request_context, + messages=[], + ), + optimization_context=PatchMergePolicyOptimizerContext( + request_context=request_context, + ), + apply_context=request_context, + ), +) +``` + +## 18. 设计原则 + +- `Case` 是训练/评测样本,不再使用 `Outcome` 概念。 +- `Rubric` 定义验收标准;`RubricEvaluation` 是一次 rollout 的评估结果。 +- `Rollout` 保留原始执行消息和可选 evaluation;`Trajectory` 是从 rollout 中抽取的可训练样本。 +- `SemanticGradient` 是 memory-file before/after 级别的语义更新信号。 +- `PolicyOptimizer` 只规划,不写文件;`PolicyUpdater` 才是写入边界。 +- batch 和 streaming 共用同一个 `PolicyTrainingEngine`。 +- 并发写入通过 `ExperienceSet.lock() + reload()` 串行化 optimizer/apply 阶段。 +- 远程 benchmark 集成应走 `RemoteCaseLoader / RemoteRolloutExecutor`,不要让训练框架直接依赖 benchmark runtime。 diff --git a/docs/en/about/02-changelog.md b/docs/en/about/02-changelog.md index af72e3ca3c..d471bdb16f 100644 --- a/docs/en/about/02-changelog.md +++ b/docs/en/about/02-changelog.md @@ -98,7 +98,7 @@ This changelog is automatically generated from [GitHub Releases](https://github. - **Native `ov` CLI refresh**: `ov config` is now the interactive configuration manager for adding, editing, deleting, and switching saved configs, while `ov config show`, `ov config validate`, and `ov config switch` remain explicit subcommands. New `ov language` / `ov lang` selects the display language, `ov status [--verbose]` provides aggregated diagnostics, and `ov health` plus runtime errors render with clearer guidance. - **Web Studio Playground and identity management**: Studio adds a Playground with a context tree, Terminal, and Agent panel, plus a Connection & Identity page that can save connection state, select account/user identity, create accounts and users, and copy or regenerate API keys. -- **Config-driven VikingBot experience recall**: New `bot.ov_server.recall_exp_first_round_only`, `exp_recall_limit`, and `exp_recall_max_chars` inject agent experience only on the first turn, and both local and remote modes isolate experience namespaces by incoming `agent_id`. +- **Config-driven VikingBot experience recall**: New `bot.ov_server.exp_recall_limit` and `exp_recall_max_chars` tune agent experience recall, and both local and remote modes isolate experience namespaces by incoming `agent_id`. - **Simpler resource watches**: `add_resource` no longer requires an explicit `to` when `watch_interval > 0`; when the import returns a stable `root_uri`, the watch task binds to it automatically, with CLI, MCP, and docs examples updated to match. - **Structured plugin tool results and CJK token estimation**: Claude Code and OpenClaw plugins now write structured tool parts instead of flattening calls and results into text only, and CJK-aware token estimation is shared across Python and plugin code to reduce budget underestimation for Chinese, Japanese, and Korean sessions. diff --git a/docs/en/api/01-overview.md b/docs/en/api/01-overview.md index fd02d6ee89..415b8112af 100644 --- a/docs/en/api/01-overview.md +++ b/docs/en/api/01-overview.md @@ -135,7 +135,7 @@ Configuration field description: | `api_key` | API Key | `null` (no auth) | | `account` | Default account header for tenant-scoped requests | `null` | | `user` | Default user header for tenant-scoped requests | `null` | -| `timeout` | HTTP request timeout in seconds | `60.0` | +| `timeout` | HTTP request timeout in seconds | `600.0` | | `output` | Default output format: `"table"` or `"json"` | `"table"` | See the [Configuration Guide](../guides/01-configuration.md#ovcliconf) for details. @@ -150,7 +150,7 @@ import openviking as ov client = ov.SyncHTTPClient( url="http://localhost:1933", # Explicitly provided api_key="your-key", # Explicitly provided (api_key usually identifies user identity) - timeout=30.0, # Don't use default 60.0 + timeout=30.0, # Don't use default 600.0 extra_headers={} # Pass empty dict instead of None, useful for gateway auth in some scenarios ) client.initialize() @@ -159,7 +159,7 @@ client.initialize() ⚠️ **Note**: The client will attempt to load the configuration file if any of the following conditions are met: - `url` is `None` - `api_key` is `None` -- `timeout` equals `60.0` (default value) +- `timeout` equals `600.0` (default value) - `extra_headers` is `None` #### HTTP Call Examples diff --git a/docs/en/api/05-sessions.md b/docs/en/api/05-sessions.md index d74463c347..3ab7ee36d0 100644 --- a/docs/en/api/05-sessions.md +++ b/docs/en/api/05-sessions.md @@ -42,7 +42,7 @@ Create a new session. Sessions are containers for conversations, storing message | Parameter | Type | Required | Default | Description | |-----------|------|----------|---------|-------------| | session_id | str | No | None | Session ID. Creates new session with auto-generated ID if None | -| memory_policy | object | No | None | Default memory extraction policy for the session. Optional `self` and `peer` switches control write targets, and optional top-level `memory_types` limits extraction to specific enabled memory schemas. When `memory_types` is omitted or `null`, all enabled memory schemas are allowed. Invalid shapes or unknown memory types are rejected with `InvalidArgumentError`. | +| memory_policy | object | No | None | Default memory extraction policy for the session. Optional `self` and `peer` switches control write targets, optional `working_memory.enabled=false` skips archive summaries, and optional top-level `memory_types` limits extraction to specific enabled memory schemas. When `memory_types` is omitted or `null`, all enabled memory schemas are allowed. Invalid shapes or unknown memory types are rejected with `InvalidArgumentError`. | #### 3. Usage Examples diff --git a/docs/en/guides/01-configuration.md b/docs/en/guides/01-configuration.md index 67057ae8f5..83fbe78512 100644 --- a/docs/en/guides/01-configuration.md +++ b/docs/en/guides/01-configuration.md @@ -588,7 +588,7 @@ Vision Language Model for semantic extraction (L0/L1 generation). | `max_concurrent` | int | Maximum concurrent semantic LLM calls (default: `64`) | | `max_retries` | int | Maximum retry attempts for transient VLM provider errors (default: `3`; `0` disables retry) | | `backup` | object | Optional backup VLM configuration (same shape as `vlm`) for automatic failover when the primary fails with retryable errors such as rate limits, `5xx` responses, or connection/timeout failures. Only one level of failover is supported — the backup itself cannot define a nested `backup` | -| `timeout` | float | Per-request HTTP timeout in seconds passed to the underlying OpenAI/LiteLLM client. Increase for slow endpoints (e.g., DashScope, local inference). Must be `> 0` (default: `60.0`) | +| `timeout` | float | Per-request HTTP timeout in seconds passed to the underlying OpenAI/LiteLLM client. Increase for slow endpoints (e.g., DashScope, local inference). Must be `> 0` (default: `600.0`) | | `extra_headers` | object | Custom HTTP headers for compatible HTTP providers. `kimi` also accepts header overrides, but already injects the required subscription headers by default | | `extra_request_body` | object | Extra JSON body fields for OpenAI-compatible completion requests, useful for provider-specific options such as Ollama `{"think": false}` | | `stream` | bool | Enable streaming mode (for OpenAI-compatible providers, default: `false`) | @@ -1318,14 +1318,18 @@ For memory-related settings, add a `memory` section in `ov.conf`: ```json { "memory": { - "version": "v2" + "custom_templates_dir": "/path/to/custom-memory" } } ``` | Field | Description | Default | |-------|-------------|---------| -| `version` | Memory implementation version. Only `"v2"` is supported (legacy `"v1"` removed in #2264 — passing `"v1"` now raises a `ValueError` at config load). | `"v2"` | +| `version` | Deprecated and ignored. OpenViking always uses the v3 memory extraction pipeline; existing configs that set this field still load without error. | `"v3"` | +| `custom_templates_dir` | Custom memory templates directory. If set, templates from this directory are loaded in addition to built-in templates. | `""` | +| `extraction_enabled` | Whether session commit runs long-term memory extraction. | `true` | +| `session_skill_extraction_enabled` | Whether session commit also extracts reusable skills into the current user's skill directory. | `false` | +| `link_enabled` | Whether memory extraction writes and resolves memory links. | `false` | ### ovcli.conf diff --git a/docs/en/guides/10-prompt-guide.md b/docs/en/guides/10-prompt-guide.md index ef98325be5..cacef59fdc 100644 --- a/docs/en/guides/10-prompt-guide.md +++ b/docs/en/guides/10-prompt-guide.md @@ -113,6 +113,8 @@ embedding_template: | directory: "viking://user/{{ user_space }}/memories/..." enabled: true operation_mode: "upsert" +stage: "user" +peer_enabled: true ``` Field meanings: @@ -135,6 +137,10 @@ Field meanings: - Whether this memory type is enabled - `operation_mode` - The update mode of the memory type, such as `upsert` +- `stage` + - Extraction stage. The default is `user`, which participates in session user-memory extraction; `agent` is reserved for execution-derived schemas such as trajectories and experiences. +- `peer_enabled` + - Whether this memory type is stored separately under peer directories when `peer_id` or message ranges identify a peer. The default is `true`; set `false` for memories that must stay under the current user space. When writing a memory schema, focus on: @@ -249,7 +255,7 @@ These YAML files define the structure of different memory types. They are not si - Key fields: `tool_name`, `static_desc`, `call_count`, `success_time`, `when_to_use`, `optimal_params` - `trajectories` - - Effective stage: agent trajectory memory persistence stage (agent-only, add-only) + - Effective stage: agent trajectory memory persistence stage (`stage: agent`, add-only) - Affects: reusable operation contracts distilled from agent task trajectories — multi-step decisions, tool calls, and execution traces - Purpose: defines compact trajectory memory for "what reusable operation/contract emerged from a task trajectory" - Key fields: `trajectory_name`, `outcome`, `retrieval_anchor`, `content` diff --git a/docs/zh/about/02-changelog.md b/docs/zh/about/02-changelog.md index faf30dd6a6..d825335f52 100644 --- a/docs/zh/about/02-changelog.md +++ b/docs/zh/about/02-changelog.md @@ -98,7 +98,7 @@ OpenViking 的所有重要变更都将记录在此文件中。 - **原生 `ov` CLI 体验重构**:`ov config` 现在是配置管理入口,可交互式添加、编辑、删除、切换配置;`ov config show`、`ov config validate`、`ov config switch` 保留为显式子命令。新增 `ov language` / `ov lang` 选择显示语言,`ov status [--verbose]` 提供聚合诊断视图,`ov health` 与错误提示改为更可读的渲染。 - **Web Studio Playground 与身份管理**:Studio 侧边栏新增 Playground,可查看上下文树、运行 Terminal 操作并与 Agent 面板交互;Connection & Identity 页面支持保存连接、选择 account/user 身份、创建 account/user、复制或重新生成 API key。 -- **VikingBot 经验召回配置化**:新增 `bot.ov_server.recall_exp_first_round_only`、`exp_recall_limit`、`exp_recall_max_chars`,用于在单任务/评测场景中只在第一轮注入 agent experience;本地与远端模式都按传入 `agent_id` 做经验命名空间隔离。 +- **VikingBot 经验召回配置化**:新增 `bot.ov_server.exp_recall_limit`、`exp_recall_max_chars`,用于调整 agent experience 召回;本地与远端模式都按传入 `agent_id` 做经验命名空间隔离。 - **资源 Watch 更易用**:`add_resource` 设置 `watch_interval > 0` 时不再强制要求显式 `to`;如果导入结果返回稳定 `root_uri`,watch task 会自动绑定到该 URI,CLI/MCP/文档示例同步更新。 - **插件结构化工具结果与 CJK token 估算**:Claude Code / OpenClaw 插件改为向 OpenViking 写入结构化 tool parts,工具调用与结果不再只能内联到文本;CJK-aware token 估算覆盖 Python 与插件侧,降低中文、日文、韩文会话的预算低估风险。 diff --git a/docs/zh/api/01-overview.md b/docs/zh/api/01-overview.md index 02c9a57e40..f5645c957a 100644 --- a/docs/zh/api/01-overview.md +++ b/docs/zh/api/01-overview.md @@ -131,7 +131,7 @@ export OPENVIKING_CLI_CONFIG_FILE=/path/to/ovcli.conf | `api_key` | API Key | `null`(无认证) | | `account` | 租户级请求的默认账户请求头 | `null` | | `user` | 租户级请求的默认用户请求头 | `null` | -| `timeout` | HTTP 请求超时时间(秒) | `60.0` | +| `timeout` | HTTP 请求超时时间(秒) | `600.0` | | `output` | 默认输出格式:`"table"` 或 `"json"` | `"table"` | 详细内容请参见 [配置指南](../guides/01-configuration.md#ovcliconf)。 @@ -146,7 +146,7 @@ import openviking as ov client = ov.SyncHTTPClient( url="http://localhost:1933", # 显式传入 api_key="your-key", # 显式传入(默认情况下 api_key 已经能标识用户身份) - timeout=30.0, # 不要用默认值 60.0 + timeout=30.0, # 不要用默认值 600.0 extra_headers={} # 传空 dict 而不是 None,可用于某些场景的网关认证等 ) client.initialize() @@ -155,7 +155,7 @@ client.initialize() ⚠️ **注意**:只要以下任一条件满足,客户端就会尝试加载配置文件: - `url` 为 `None` - `api_key` 为 `None` -- `timeout` 等于 `60.0`(默认值) +- `timeout` 等于 `600.0`(默认值) - `extra_headers` 为 `None` #### HTTP 调用示例 diff --git a/docs/zh/api/05-sessions.md b/docs/zh/api/05-sessions.md index f7bc8e46b3..14232098d9 100644 --- a/docs/zh/api/05-sessions.md +++ b/docs/zh/api/05-sessions.md @@ -42,7 +42,7 @@ Session API 按认证用户作用域访问会话,并返回 canonical user sess | 参数 | 类型 | 必填 | 默认值 | 说明 | |------|------|------|--------|------| | session_id | str | 否 | None | 会话 ID。如果为 None,则创建一个自动生成 ID 的新会话 | -| memory_policy | object | 否 | None | 会话默认的记忆抽取策略。可选的 `self` 和 `peer` 开关控制写入目标;可选的顶层 `memory_types` 将抽取限制为指定的 enabled memory schema。未传或为 `null` 时允许所有 enabled memory schema。非法结构或未知 memory type 会以 `InvalidArgumentError` 拒绝。 | +| memory_policy | object | 否 | None | 会话默认的记忆抽取策略。可选的 `self` 和 `peer` 开关控制写入目标;可选的 `working_memory.enabled=false` 跳过 archive summary;可选的顶层 `memory_types` 将抽取限制为指定的 enabled memory schema。未传或为 `null` 时允许所有 enabled memory schema。非法结构或未知 memory type 会以 `InvalidArgumentError` 拒绝。 | #### 3. 使用示例 diff --git a/docs/zh/guides/01-configuration.md b/docs/zh/guides/01-configuration.md index 2ead03bcd1..f1db1f2d54 100644 --- a/docs/zh/guides/01-configuration.md +++ b/docs/zh/guides/01-configuration.md @@ -559,7 +559,7 @@ openviking-server doctor | `max_concurrent` | int | 语义处理阶段 LLM 最大并发调用数(默认:`64`) | | `max_retries` | int | VLM provider 瞬时错误的最大重试次数(默认:`3`;`0` 表示禁用重试) | | `backup` | object | 可选的备用 VLM 配置(结构与 `vlm` 相同),当主 VLM 遇到限流、`5xx`、超时或连接失败等可重试错误时自动切换。仅支持 1 层备用 — 备用 VLM 本身不能再嵌套 `backup` | -| `timeout` | float | 单次 VLM API 请求的 HTTP 超时时间(秒),传递给底层 OpenAI/LiteLLM 客户端。慢端点(如 DashScope、本地推理)可调大。必须 `> 0`(默认:`60.0`) | +| `timeout` | float | 单次 VLM API 请求的 HTTP 超时时间(秒),传递给底层 OpenAI/LiteLLM 客户端。慢端点(如 DashScope、本地推理)可调大。必须 `> 0`(默认:`600.0`) | | `extra_headers` | object | 兼容 HTTP provider 的自定义请求头。`kimi` 默认已注入所需订阅请求头,也支持在这里覆盖或扩展 | | `extra_request_body` | object | 传给 OpenAI 兼容 completion 请求的额外 JSON body 字段,可用于 Ollama `{"think": false}` 等 provider 专有参数 | | `stream` | bool | 启用流式模式(OpenAI 兼容 provider 可用,默认:`false`) | @@ -1291,14 +1291,18 @@ openviking-server --config /path/to/ov.conf ```json { "memory": { - "version": "v2" + "custom_templates_dir": "/path/to/custom-memory" } } ``` | 字段 | 说明 | 默认值 | |------|------|--------| -| `version` | 记忆实现版本。仅支持 `"v2"`(#2264 已移除旧版 `"v1"`;传入 `"v1"` 会在配置加载时抛出 `ValueError`)。 | `"v2"` | +| `version` | 已废弃且会被忽略。OpenViking 始终使用 v3 记忆抽取链路;已有配置中保留该字段仍可正常加载,不会报错。 | `"v3"` | +| `custom_templates_dir` | 自定义 memory templates 目录。设置后会在内置模板之外加载该目录中的模板。 | `""` | +| `extraction_enabled` | session commit 时是否执行长期记忆抽取。 | `true` | +| `session_skill_extraction_enabled` | session commit 时是否同时抽取可复用 skill 到当前用户的 skill 目录。 | `false` | +| `link_enabled` | 记忆抽取是否写入和解析 memory links。 | `false` | ### ovcli.conf diff --git a/docs/zh/guides/10-prompt-guide.md b/docs/zh/guides/10-prompt-guide.md index 2d7dd4301a..b0f25454cf 100644 --- a/docs/zh/guides/10-prompt-guide.md +++ b/docs/zh/guides/10-prompt-guide.md @@ -113,6 +113,8 @@ embedding_template: | directory: "viking://user/{{ user_space }}/memories/..." enabled: true operation_mode: "upsert" +stage: "user" +peer_enabled: true ``` 字段含义: @@ -135,6 +137,10 @@ operation_mode: "upsert" - 是否启用该类记忆 - `operation_mode` - 该类记忆的更新模式,例如 `upsert` +- `stage` + - 抽取阶段。默认是 `user`,参与会话用户记忆抽取;`agent` 用于 trajectories、experiences 这类执行派生 schema。 +- `peer_enabled` + - 当 `peer_id` 或消息 ranges 指向某个 peer 时,是否将该类记忆按 peer 分目录存储。默认是 `true`;如果该类记忆必须保留在当前 user 目录下,设置为 `false`。 编写 memory schema 时,建议重点关注: @@ -249,7 +255,7 @@ operation_mode: "upsert" - 关键字段:`tool_name`、`static_desc`、`call_count`、`success_time`、`when_to_use`、`optimal_params` - `trajectories` - - 生效环节:agent 轨迹型记忆落盘阶段(agent_only,仅追加) + - 生效环节:agent 轨迹型记忆落盘阶段(`stage: agent`,仅追加) - 影响能力:agent 任务轨迹中可复用的操作契约沉淀——多步决策、工具调用、执行链路 - 作用:定义"任务轨迹中提炼出哪些可复用的操作/契约"这一类轨迹型记忆 - 关键字段:`trajectory_name`、`outcome`、`retrieval_anchor`、`content` diff --git a/openviking/async_client.py b/openviking/async_client.py index 87a04f86a9..2bbece85a3 100644 --- a/openviking/async_client.py +++ b/openviking/async_client.py @@ -606,6 +606,14 @@ async def read(self, uri: str, offset: int = 0, limit: int = -1) -> str: await self._ensure_initialized() return await self._client.read(uri, offset=offset, limit=limit) + async def read_raw(self, uri: str, offset: int = 0, limit: int = -1) -> str: + """Read raw file content, including hidden MEMORY_FIELDS metadata.""" + await self._ensure_initialized() + read_raw = getattr(self._client, "read_raw", None) + if read_raw is not None: + return await read_raw(uri, offset=offset, limit=limit) + return await self._client.read(uri, offset=offset, limit=limit) + async def write( self, uri: str, diff --git a/openviking/client/local.py b/openviking/client/local.py index 546aefbc91..37dbc070fe 100644 --- a/openviking/client/local.py +++ b/openviking/client/local.py @@ -606,6 +606,10 @@ async def read(self, uri: str, offset: int = 0, limit: int = -1) -> str: """ return await self._service.fs.read(uri, ctx=self._ctx, offset=offset, limit=limit) + async def read_raw(self, uri: str, offset: int = 0, limit: int = -1) -> str: + """Read raw file content, including hidden MEMORY_FIELDS metadata.""" + return await self._service.fs.read(uri, ctx=self._ctx, offset=offset, limit=limit) + async def abstract(self, uri: str) -> str: """Read L0 abstract.""" return await self._service.fs.abstract(uri, ctx=self._ctx) diff --git a/openviking/message/message.py b/openviking/message/message.py index 97cb73e099..80d968b9e4 100644 --- a/openviking/message/message.py +++ b/openviking/message/message.py @@ -11,7 +11,7 @@ from typing import List, Literal, Optional from openviking.core.peer_id import normalize_peer_id -from openviking.message.part import ContextPart, ImagePart, Part, TextPart, ToolPart +from openviking.message.part import ContextPart, ImagePart, Part, TextPart, ToolPart, part_from_dict from openviking.utils.token_estimation import estimate_text_tokens @@ -168,62 +168,7 @@ def from_dict(cls, data: dict) -> "Message": raw_parts = [] for p in raw_parts: - if p["type"] == "text": - parts.append(TextPart(text=p.get("text", ""))) - elif p["type"] == "context": - parts.append( - ContextPart( - uri=p["uri"], - context_type=p.get("context_type", "memory"), - abstract=p.get("abstract", ""), - ) - ) - elif p["type"] == "image_url": - image_url = p.get("image_url") - url = "" - detail = None - if isinstance(image_url, dict): - url = str(image_url.get("url", "") or "") - detail = image_url.get("detail") - elif isinstance(image_url, str): - url = image_url - if not url.strip(): - raise ValueError("image_url part requires a non-empty URL") - parts.append(ImagePart(url=url, detail=detail)) - elif p["type"] == "tool": - parts.append( - ToolPart( - tool_id=p["tool_id"], - tool_name=p["tool_name"], - tool_uri=p["tool_uri"], - skill_uri=p.get("skill_uri", ""), - tool_input=p.get("tool_input"), - tool_output=p.get("tool_output", ""), - tool_status=p.get("tool_status", "pending"), - duration_ms=p.get("duration_ms"), - prompt_tokens=p.get("prompt_tokens"), - completion_tokens=p.get("completion_tokens"), - tool_output_ref=p.get("tool_output_ref", ""), - tool_output_truncated=bool(p.get("tool_output_truncated", False)), - tool_output_original_chars=p.get("tool_output_original_chars"), - tool_output_preview_chars=p.get("tool_output_preview_chars"), - tool_output_sha256=p.get("tool_output_sha256", ""), - tool_output_storage_uri=p.get("tool_output_storage_uri", ""), - tool_output_mime_type=p.get("tool_output_mime_type", "text/plain"), - tool_output_source_ref=p.get("tool_output_source_ref", ""), - tool_output_source_offset=p.get("tool_output_source_offset"), - tool_output_source_limit=p.get("tool_output_source_limit"), - tool_output_externalization_error=p.get( - "tool_output_externalization_error", "" - ), - tool_output_group_id=p.get("tool_output_group_id", ""), - tool_output_externalized_reason=p.get( - "tool_output_externalized_reason", "" - ), - tool_output_group_original_chars=p.get("tool_output_group_original_chars"), - tool_output_group_budget_chars=p.get("tool_output_group_budget_chars"), - ) - ) + parts.append(part_from_dict(p)) try: peer_id = normalize_peer_id(data.get("peer_id")) except ValueError: diff --git a/openviking/message/part.py b/openviking/message/part.py index d738e48367..d735083b21 100644 --- a/openviking/message/part.py +++ b/openviking/message/part.py @@ -141,4 +141,6 @@ def part_from_dict(data: Dict[str, Any]) -> Part: tool_output_group_budget_chars=data.get("tool_output_group_budget_chars"), ) else: + if "text" in data: + return TextPart(text=str(data.get("text", "") or "")) return TextPart(text=str(data)) diff --git a/openviking/models/vlm/backends/openai_vlm.py b/openviking/models/vlm/backends/openai_vlm.py index 012adcbf3d..b5f23696c8 100644 --- a/openviking/models/vlm/backends/openai_vlm.py +++ b/openviking/models/vlm/backends/openai_vlm.py @@ -52,7 +52,7 @@ def _build_openai_client_kwargs( api_base: str, api_version: str | None, extra_headers: Dict[str, str] | None, - timeout: float = 60.0, + timeout: float = 600.0, ) -> Dict[str, Any]: """Build kwargs dict shared by sync and async OpenAI/Azure client constructors.""" if provider == "azure": diff --git a/openviking/models/vlm/base.py b/openviking/models/vlm/base.py index 7aeb1b54ff..04762de163 100644 --- a/openviking/models/vlm/base.py +++ b/openviking/models/vlm/base.py @@ -67,7 +67,7 @@ def __init__(self, config: Dict[str, Any]): self.api_base = config.get("api_base") self.temperature = config.get("temperature", 0.0) self.max_retries = config.get("max_retries", 3) - self.timeout = config.get("timeout", 60.0) + self.timeout = config.get("timeout", 600.0) self.max_tokens = config.get("max_tokens") self.extra_headers = config.get("extra_headers") self.extra_request_body = dict(config.get("extra_request_body") or {}) diff --git a/openviking/prompts/templates/memory/cases.yaml b/openviking/prompts/templates/memory/cases.yaml new file mode 100644 index 0000000000..d291b5b066 --- /dev/null +++ b/openviking/prompts/templates/memory/cases.yaml @@ -0,0 +1,99 @@ +memory_type: cases +description: | + Language policy: generated natural-language text must use {{ language }}. If the source is multilingual, use the current task user's primary language as {{ language }}; if unclear, use the latest user message language. Preserve stable identifiers, code-like tokens, tool/API names, enum values, exact IDs, paths, dates, amounts, and other structured literals unchanged. When updating an existing case, prefer its existing name and language style to avoid translated duplicates. + + A trainable/evaluable case extracted from a real session commit. + + Extract ONLY when the conversation contains a concrete user task plus enough assistant execution, + tool-use evidence, or final result to define what a good future rollout should do. + Skip pure chit-chat, simple factual Q&A, or fragments with no observable task outcome. + + The case is not an experience instruction. It defines the scenario and rubric used by the + training pipeline to learn better experience memories from this commit rollout. + +directory: "viking://user/{{ user_space }}/memories/cases" +filename_template: "{{ case_name }}.md" +enabled: true +peer_enabled: false +content_template: | + # {{ case_name }} + + ## Task Signature + {{ task_signature }} + + ## Input + {{ input }} + + ## Rubric + {{ rubric }} + + ## Evidence + {{ evidence }} + + ## Linked Experiences + {%- for link in links or [] %} + {%- set target_uri = link.to_uri or "" %} + {%- if "/memories/experiences/" in target_uri %} + - [{{ uri_basename(target_uri) }}]({{ link_target(target_uri) }}) + {%- endif %} + {%- endfor %} +embedding_template: |- + {{ case_name }} + + {{ task_signature }} + + {{ input }} + + {{ rubric }} +overview_template: |- + # Cases Overview + {% for item in items %} + - [{{ item.file_content.case_name }}](./{{ item.file_name }}) — {{ item.file_content.task_signature }} + {% endfor %} + +fields: + - name: case_name + type: string + description: | + Short stable case name for this concrete training/evaluation sample. + Must be written in {{ language }}. + {% if language == 'en' %}Use lowercase snake_case, max 5 words.{% else %}Use a concise noun phrase, max 15 characters.{% endif %} + Name the task boundary, not the whole conversation. + merge_op: immutable + + - name: task_signature + type: string + description: | + Semantic task signature for retrieval and grouping. + Write generated natural-language text in {{ language }}. + Describe the reusable task intent, object/effect boundary, and success condition in one sentence. + Generalize private identifiers, exact IDs, names, amounts, dates, and paths. + merge_op: immutable + + - name: input + type: string + description: | + Compact JSON object string describing the executable case input. + JSON keys must remain stable English identifiers; generated string values should use {{ language }} while preserving exact source literals needed for task matching. + Include the user request summary and any stable preconditions needed to reproduce/evaluate the task. + Example: {"summary":"用户要求处理重复预订并保留有效订单","preconditions":["存在两个相似订单候选"]} + merge_op: immutable + + - name: rubric + type: string + description: | + Compact JSON object string defining what good means and how to check it. + JSON keys must remain stable English identifiers; generated string values should use {{ language }} while preserving exact tool names, required literals, and observable evaluation labels. + Required shape: + {"name":"...","description":"...","criteria":[{"name":"...","description":"observable success condition","required":true,"weight":1.0}]} + Criteria must be observable from rollout messages/tool results; include both success rate and execution efficiency when relevant. + Use 1-5 criteria. Weight values should be positive. + merge_op: immutable + + - name: evidence + type: string + description: | + Brief generalized evidence from the commit conversation explaining why this case is useful. + Write generated natural-language text in {{ language }}. + Do not include private raw identifiers, exact IDs, personal contact values, or full tool payloads. + merge_op: immutable diff --git a/openviking/prompts/templates/memory/events.yaml b/openviking/prompts/templates/memory/events.yaml index 7837cdc974..f1879edd2a 100644 --- a/openviking/prompts/templates/memory/events.yaml +++ b/openviking/prompts/templates/memory/events.yaml @@ -97,8 +97,9 @@ content_template: | {% if resource_event_content %} {{ resource_event_content }} {% else %} - Summary: {{ summary }} - {{extract_context.get_first_message_time_with_weekday_from_ranges(ranges|default(''))|default('N/A')}} ChatLog: + # Summary + {{ summary }} + # {{extract_context.get_first_message_time_with_weekday_from_ranges(ranges|default(''))|default('N/A')}} ChatLog: {{ extract_context.get_event_content(ranges, summary, 0) }} {% endif %} embedding_template: |- diff --git a/openviking/prompts/templates/memory/preferences.yaml b/openviking/prompts/templates/memory/preferences.yaml index 65bbfadb46..ca65a9e3c0 100644 --- a/openviking/prompts/templates/memory/preferences.yaml +++ b/openviking/prompts/templates/memory/preferences.yaml @@ -2,12 +2,9 @@ memory_type: preferences description: | User preference memory - captures "what the user likes/dislikes or is accustomed to". Extract specific preferences the user has expressed across conversations. - Keep each preference memory Complete but minimal. Each preference should be about a specific topic (not generic). Topics can be: code style, communication style, tools, workflow, food, commute, etc. Store different topics as separate memory files, do NOT mix unrelated preferences. - Each topic should stay small and should not become a second profile. - Keep each topic to roughly 3-8 bullets; if it grows beyond that or past 800 characters, split it into semantic subtopics. directory: "viking://user/{{ user_space }}/memories/preferences" filename_template: "{{ user }}/{{ topic }}.md" enabled: true @@ -41,16 +38,11 @@ fields: Preference topic used to uniquely identify this preference memory. Should be a semantic topic description such as "Python code style", "Communication style", "Food preference", "Commute preference", etc. Different preference topics should be stored as separate memories, do not mix unrelated preferences. - If a topic becomes too broad, split it into smaller semantic subtopics. merge_op: immutable - name: content type: string description: | Preference content in Markdown format describing "what the user prefers/is accustomed to". - Rewrite each topic into a Complete but minimal version rather than letting it grow indefinitely. - Keep each topic concise, usually 3-8 bullets; if it exceeds 8 bullets or 800 characters, split it into semantic subtopics. - Use compact statements, and when useful include light evidence such as "evidenced by ... (as of YYYY-MM-DD)". - Slight identity context is allowed only when it helps explain the preference, but this memory should not become a second profile. Example: "User has shown clear preferences for Python code style in multiple conversations: dislikes using type hints, considers them redundant; requires concise function comments, limited to 1-2 lines; prefers direct implementation, avoids excessive fallbacks and over-engineering." merge_op: patch diff --git a/openviking/prompts/templates/memory/profile.yaml b/openviking/prompts/templates/memory/profile.yaml index 4a84e1fd93..9d3c67ef7b 100644 --- a/openviking/prompts/templates/memory/profile.yaml +++ b/openviking/prompts/templates/memory/profile.yaml @@ -1,22 +1,15 @@ memory_type: profile description: | # Task Objective - User profile memory - captures an identity summary of who the user is as a person. - Extract relatively stable personal attributes that define the user's identity at a high level. - Keep the result Complete but minimal. - Profile should stay as an identity summary, not a place for endlessly growing details. + User profile memory - captures "who the user is" as a person. + Extract relatively stable personal attributes that define the user's identity, work style, and preferences. + Include: profession, experience level, technical background, communication style, work habits, etc. Do NOT include transient conversation content or temporary mood states. # Rules - - Rewrite the full profile into a Complete but minimal version each time - - Do not append to the old profile - - Keep the profile to 5-8 bullets total - Each item: self-contained, declarative sentence, < 30 words - Extract only facts stated/confirmed by user; no guesses - Focus on persistent information, not temporary situations - - Keep only high-level identity facts in profile - - If preference, habit, taste, workflow style, or other reusable behavioral detail would make profile grow, migrate it to preferences instead - - Do not keep concrete preference examples in profile - Forbidden: events, only-assistant content, sensitive/private info, trivial updates - Merge similar items; keep latest if conflicting @@ -28,13 +21,9 @@ fields: - name: content type: string description: | - User profile content in Markdown format describing "who the user is" as an identity summary. - Rewrite the full profile as a Complete but minimal version, rather than extending the old text. - Keep only the most important high-level identity facts, with 5-8 bullets total. + User profile content in Markdown format describing "who the user is". + Includes relatively stable personal attributes: profession, experience, tech stack, communication style, etc. Only record objective statuses, do not record events or similar information. - If reusable preference-oriented details overflow profile scope, migrate them to preferences instead. - Even though this field uses patch semantics, you may use patch to rewrite the whole profile into a shorter minimal version when needed. - Prefer large-scale replacement over local edits when compressing an oversized profile. [IMPORTANT] For changeable statuses, must include the last updated time in the format: (as of 2023-06-09) Example: # Caroline diff --git a/openviking/server/routers/system.py b/openviking/server/routers/system.py index da8952c51c..7d9dde77b5 100644 --- a/openviking/server/routers/system.py +++ b/openviking/server/routers/system.py @@ -12,7 +12,7 @@ from openviking.core.path_variables import resolve_path_variables from openviking.core.uri_validation import validate_viking_uri from openviking.pyagfs.exceptions import AGFSInvalidOperationError, AGFSNotSupportedError -from openviking.server.auth import get_request_context, require_role, resolve_identity +from openviking.server.auth import get_request_context, require_role from openviking.server.dependencies import get_service from openviking.server.identity import AuthMode, RequestContext, Role from openviking.server.models import Response @@ -69,40 +69,17 @@ async def _embedding_probe(embedder) -> str: @router.get("/health", tags=["system"]) async def health_check(request: Request): - """Health check endpoint (no authentication required).""" + """Health check endpoint (no authentication or identity resolution required).""" from openviking import __version__ result = {"status": "ok", "healthy": True, "version": __version__} - # Try to get user identity try: - # Extract headers manually - x_api_key = request.headers.get("X-API-Key") - authorization = request.headers.get("Authorization") - x_openviking_account = request.headers.get("X-OpenViking-Account") - x_openviking_user = request.headers.get("X-OpenViking-User") - - # Get effective auth mode from config - effective_auth_mode = AuthMode.API_KEY.value config = getattr(request.app.state, "config", None) + effective_auth_mode = AuthMode.API_KEY.value if config is not None and hasattr(config, "get_effective_auth_mode"): effective_auth_mode = config.get_effective_auth_mode() result["auth_mode"] = effective_auth_mode - - if x_api_key or authorization: - try: - identity = await resolve_identity( - request, - x_api_key=x_api_key, - authorization=authorization, - x_openviking_account=x_openviking_account, - x_openviking_user=x_openviking_user, - ) - result["account_id"] = str(identity.account_id) - result["user_id"] = str(identity.user_id) - result["role"] = str(identity.role) - except Exception as e: - logger.warning(f"Failed to resolve identity: {e}") except Exception as e: logger.error(f"Failed to get health check: {e}") diff --git a/openviking/session/__init__.py b/openviking/session/__init__.py index 877638b8f0..9ee6c16bef 100644 --- a/openviking/session/__init__.py +++ b/openviking/session/__init__.py @@ -11,31 +11,33 @@ logger = get_logger(__name__) if TYPE_CHECKING: - from openviking.session.compressor_v2 import SessionCompressorV2 + from openviking.session.compressor_v3 import SessionCompressorV3 def create_session_compressor( vikingdb: VikingDBManager, memory_version: Optional[str] = None, skill_processor=None, -) -> "SessionCompressorV2": +) -> "SessionCompressorV3": """ - Create the v2 session compressor. + Create the session compressor. Args: vikingdb: VikingDBManager instance - memory_version: Deprecated optional override. Only "v2" is supported. + memory_version: Deprecated and ignored; v3 is always used. Existing + configs that still set memory.version continue to load, but no + longer select the implementation. Returns: - SessionCompressorV2 instance + v3 session compressor instance """ - if memory_version not in (None, "v2"): - raise ValueError("memory.version only supports 'v2'; legacy memory v1 has been removed") + if memory_version is not None: + logger.warning("memory.version is deprecated and ignored; using v3 memory compressor") - logger.info("Using v2 memory compressor (templating system)") - from openviking.session.compressor_v2 import SessionCompressorV2 + logger.info("Using v3 memory compressor (v2 + commit streaming train)") + from openviking.session.compressor_v3 import SessionCompressorV3 - return SessionCompressorV2(vikingdb=vikingdb, skill_processor=skill_processor) + return SessionCompressorV3(vikingdb=vikingdb, skill_processor=skill_processor) __all__ = [ diff --git a/openviking/session/compressor_v2.py b/openviking/session/compressor_v2.py index 29fd1edd0d..ff0b260a64 100644 --- a/openviking/session/compressor_v2.py +++ b/openviking/session/compressor_v2.py @@ -547,6 +547,7 @@ async def extract_execution_memories( strict_extract_errors=strict_extract_errors, phase_label="trajectory", allowed_memory_types=allowed_execution_types, + thinking=True, ) if traj_result is None: return empty_result @@ -622,6 +623,7 @@ async def _append_sources_before_unlock( phase_label=f"experience({traj_uri})", post_apply=_append_sources_before_unlock, allowed_memory_types=allowed_execution_types, + thinking=True, ) if exp_result is None: fallback_uris = await self._single_existing_experience_uris( @@ -719,6 +721,7 @@ async def _run_extract_phase( phase_label: str, post_apply: Optional[ExtractPostApply] = None, allowed_memory_types: Optional[set[str]] = None, + thinking: bool = False, ): """Run one ExtractLoop phase with its own lock scope, then apply operations. @@ -758,6 +761,7 @@ async def _run_extract_phase( ctx=ctx, context_provider=provider, isolation_handler=isolation_handler, + thinking=thinking, ) lock_manager = None @@ -1208,6 +1212,7 @@ async def _build_memory_diff( return { "archive_uri": archive_uri, + "trace_id": tracer.get_trace_id() or None, "extracted_at": datetime.utcnow().isoformat() + "Z", "operations": { "adds": adds, diff --git a/openviking/session/compressor_v3.py b/openviking/session/compressor_v3.py new file mode 100644 index 0000000000..8e89c0b12f --- /dev/null +++ b/openviking/session/compressor_v3.py @@ -0,0 +1,1744 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Session Compressor V3. + +V3 keeps the V2 extraction interface while changing user-memory commits to a +patch-merge flow without directory-level memory locks. Training cases are not +extracted by a separate LLM call; they are a normal user-memory ``memory_type`` +(``cases``) emitted by the same ExtractLoop that extracts profile/preferences/ +events/etc. When such case memories are produced, the same commit rollout is +submitted to the process-global StreamingPolicyTrainer. +""" + +from __future__ import annotations + +import json +import re +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Any, List, Optional +from uuid import uuid4 + +from openviking.core.context import Context +from openviking.message import Message +from openviking.server.identity import RequestContext +from openviking.session.memory import ExtractLoop, MemoryUpdater, StreamingMemoryUpdaterConfig +from openviking.session.memory.dataclass import ( + MemoryFile, + MemoryOperationSource, + ResolvedOperation, + ResolvedOperations, + StoredLink, +) +from openviking.session.memory.memory_isolation_handler import MemoryIsolationHandler +from openviking.session.memory.memory_type_registry import create_default_registry +from openviking.session.memory.memory_updater import ExtractContext, write_stored_links +from openviking.session.memory.session_extract_context_provider import ( + SessionExtractContextProvider, +) +from openviking.session.memory.streaming_memory_updater import ( + MemoryUpdateRequest, + get_streaming_memory_updater, + make_streaming_memory_updater_key, + merge_link_lists, +) +from openviking.session.memory.utils.json_parser import JsonUtils +from openviking.session.memory.utils.memory_file_utils import MemoryFileUtils +from openviking.session.memory.utils.uri import generate_uri +from openviking.session.train import ( + Case, + ExperienceGradientContext, + ExperienceGradientEstimator, + ExperienceSetLoader, + MemoryFilePolicyUpdater, + PatchMergePolicyOptimizer, + PatchMergePolicyOptimizerContext, + PipelineContext, + PolicyApplyResult, + PolicyPlanItem, + PolicyUpdatePlan, + Rollout, + RolloutAnalysis, + RolloutTrainingResult, + Rubric, + RubricCriterion, + SkillPolicyUpdater, + SkillSetLoader, + StreamingPolicyTrainerConfig, + TrajectoryAnalyzerContext, + TrajectoryRolloutAnalyzer, + get_streaming_policy_trainer, + make_streaming_policy_trainer_key, +) +from openviking.storage.viking_fs import get_viking_fs +from openviking.telemetry import tracer +from openviking_cli.utils import get_logger +from openviking_cli.utils.config import get_openviking_config + +logger = get_logger(__name__) + +_CASES_MEMORY_TYPE = "cases" +_TRAINING_CASE_SPEC_PROTOCOL = "openviking.batch_train.case_spec.v1" +_TRAINING_CASE_SPEC_HEADER = "# OpenViking Batch Training CaseSpec v1" +_TRAINING_FAST_PATH_MEMORY_TYPES = frozenset({"cases", "trajectories", "experiences"}) +_JSON_FENCE_RE = re.compile(r"```(?:json)?\s*(.*?)\s*```", re.DOTALL | re.IGNORECASE) + + +class SessionCompressorV3: + """Session compressor with lock-free patch-merge user memory extraction.""" + + rollout_analyzer: TrajectoryRolloutAnalyzer | Any + streaming_trainer_config: StreamingPolicyTrainerConfig = field( + default_factory=StreamingPolicyTrainerConfig + ) + streaming_memory_updater_config: StreamingMemoryUpdaterConfig = field( + default_factory=StreamingMemoryUpdaterConfig + ) + + def __init__( + self, + vikingdb, + skill_processor: Optional[Any] = None, + *, + rollout_analyzer: TrajectoryRolloutAnalyzer | Any | None = None, + streaming_trainer_config: StreamingPolicyTrainerConfig | None = None, + streaming_memory_updater_config: StreamingMemoryUpdaterConfig | None = None, + ): + self.vikingdb = vikingdb + self.skill_processor = skill_processor + self.rollout_analyzer = rollout_analyzer or TrajectoryRolloutAnalyzer( + viking_fs=get_viking_fs(), + vikingdb=vikingdb, + ) + self.streaming_trainer_config = streaming_trainer_config or StreamingPolicyTrainerConfig() + self.streaming_memory_updater_config = ( + streaming_memory_updater_config or StreamingMemoryUpdaterConfig() + ) + + def _get_or_create_react( + self, + ctx: Optional[RequestContext] = None, + messages: Optional[List] = None, + latest_archive_overview: str = "", + isolation_handler: Optional[MemoryIsolationHandler] = None, + transaction_handle=None, + ) -> ExtractLoop: + config = get_openviking_config() + vlm = config.vlm.get_vlm_instance() + viking_fs = get_viking_fs() + context_provider = SessionExtractContextProvider( + messages=messages, + latest_archive_overview=latest_archive_overview, + isolation_handler=isolation_handler, + ctx=ctx, + viking_fs=viking_fs, + transaction_handle=transaction_handle, + ) + return ExtractLoop( + vlm=vlm, + viking_fs=viking_fs, + ctx=ctx, + context_provider=context_provider, + isolation_handler=isolation_handler, + ) + + def _get_or_create_updater(self, registry, transaction_handle=None) -> MemoryUpdater: + return MemoryUpdater( + registry=registry, + vikingdb=self.vikingdb, + transaction_handle=transaction_handle, + ) + + async def _build_memory_diff( + self, + result: Any, + operations: ResolvedOperations, + viking_fs: Any, + ctx: RequestContext, + archive_uri: str = "", + ) -> dict[str, Any]: + adds: list[dict[str, Any]] = [] + updates: list[dict[str, Any]] = [] + deletes: list[dict[str, Any]] = [] + + upsert_by_uri = {} + for op in operations.upsert_operations: + for uri in op.uris: + upsert_by_uri[uri] = op + delete_by_uri = {dc.uri: dc for dc in operations.delete_file_contents} + + for uri in result.written_uris: + op = upsert_by_uri.get(uri) + memory_type = ( + op.memory_type if op else MemoryUpdater.memory_type_from_uri(uri) or "unknown" + ) + old_file = op.old_memory_file_content if op else None + if old_file: + updates.append( + { + "uri": uri, + "memory_type": memory_type, + "before": old_file.content, + "after": "", + } + ) + else: + adds.append({"uri": uri, "memory_type": memory_type, "after": ""}) + + for uri in result.edited_uris: + op = upsert_by_uri.get(uri) + memory_type = ( + op.memory_type if op else MemoryUpdater.memory_type_from_uri(uri) or "unknown" + ) + old_file = op.old_memory_file_content if op and op.old_memory_file_content else None + updates.append( + { + "uri": uri, + "memory_type": memory_type, + "before": old_file.content if old_file else "", + "after": "", + } + ) + + for uri in result.deleted_uris: + deleted = delete_by_uri.get(uri) + deletes.append( + { + "uri": uri, + "memory_type": (deleted.memory_type if deleted else None) or "unknown", + "deleted_content": deleted.content if deleted else "", + } + ) + + # Read new content for adds and updates. + # Some upsert operations can be reported as successful even when the + # final file body is identical to the pre-existing content (for + # example, a no-op merge/patch or a write that only re-serializes the + # same memory). memory_diff.json should only include effective content + # changes, so filter no-op updates after the final content is known. + for item in adds: + try: + content = await viking_fs.read_file(uri=item["uri"], ctx=ctx) + item["after"] = MemoryFileUtils.read(content).content + except Exception: + pass + + effective_updates: list[dict[str, Any]] = [] + for item in updates: + op = upsert_by_uri.get(item["uri"]) + old_file = op.old_memory_file_content if op else None + new_file: Optional[MemoryFile] = None + try: + content = await viking_fs.read_file(uri=item["uri"], ctx=ctx) + new_file = MemoryFileUtils.read(content, uri=item["uri"]) + item["after"] = new_file.content + except Exception: + pass + if old_file is not None and _same_memory_file(old_file, new_file): + logger.info( + "Skipping unchanged memory file in memory_diff.json: %s", + item.get("uri"), + ) + continue + effective_updates.append(item) + updates = effective_updates + + return _make_memory_diff( + archive_uri=archive_uri, + adds=adds, + updates=updates, + deletes=deletes, + ) + + @tracer(ignore_result=True) + async def extract_long_term_memories( + self, + messages: List[Message], + user: Optional[Any] = None, + session_id: Optional[str] = None, + ctx: Optional[RequestContext] = None, + strict_extract_errors: bool = False, + latest_archive_overview: str = "", + archive_uri: Optional[str] = None, + allowed_memory_types: Optional[set[str]] = None, + allow_self_memory: bool = True, + allowed_peer_ids: Optional[set[str]] = None, + ): + message_list = list(messages) + fast_path_case = _training_case_from_first_message(message_list, allowed_memory_types) + if fast_path_case is not None: + return await self._commit_training_case_fast_path( + case=fast_path_case, + messages=message_list, + ctx=ctx, + session_id=session_id, + archive_uri=archive_uri or "", + strict_extract_errors=strict_extract_errors, + ) + + result = await self._extract_user_memories( + messages=message_list, + user=user, + session_id=session_id, + ctx=ctx, + strict_extract_errors=strict_extract_errors, + latest_archive_overview=latest_archive_overview, + archive_uri=archive_uri, + allowed_memory_types=allowed_memory_types, + allow_self_memory=allow_self_memory, + allowed_peer_ids=allowed_peer_ids, + ) + train_result = await self.train_from_extracted_cases( + cases=result.cases, + messages=message_list, + ctx=ctx, + case_uri_by_name=getattr(result, "case_uri_by_name", {}), + session_id=session_id, + archive_uri=archive_uri or "", + strict_extract_errors=strict_extract_errors, + collect_memory_diff=True, + ) + await self._write_final_memory_diff( + archive_uri=archive_uri or "", + ctx=ctx, + memory_diffs=[ + getattr(result, "memory_diff", None), + train_result.get("memory_diff"), + ], + ) + return _v3_extraction_response( + contexts=result.contexts, + train_result=train_result, + archive_uri=archive_uri or "", + ) + + async def _commit_training_case_fast_path( + self, + *, + case: Case, + messages: list[Message], + ctx: Optional[RequestContext], + session_id: Optional[str], + archive_uri: str, + strict_extract_errors: bool, + ) -> dict[str, Any]: + if ctx is None: + logger.warning("No RequestContext provided, skipping training case fast path") + return {"contexts": [], "session_skills": []} + case_write = await self._write_training_case_memory( + case=case, + ctx=ctx, + archive_uri=archive_uri, + ) + case_result = _applied_memory_result(case_write) + contexts = _contexts_from_update_result(case_result) + train_result = await self.train_from_extracted_cases( + cases=[case], + messages=_training_messages_after_case_spec(messages), + ctx=ctx, + case_uri_by_name={case.name: _first_context_uri(contexts)}, + session_id=session_id, + archive_uri=archive_uri, + strict_extract_errors=strict_extract_errors, + collect_memory_diff=True, + ) + await self._write_final_memory_diff( + archive_uri=archive_uri, + ctx=ctx, + memory_diffs=[ + _applied_memory_diff(case_write), + train_result.get("memory_diff"), + ], + ) + return _v3_extraction_response( + contexts=contexts, + train_result=train_result, + archive_uri=archive_uri, + ) + + @tracer("train.compressor_v3.fast_path.write_case", ignore_result=True, ignore_args=True) + async def _write_training_case_memory( + self, + *, + case: Case, + ctx: RequestContext, + archive_uri: str, + ) -> Any: + viking_fs = get_viking_fs() + registry = create_default_registry() + schema = registry.get(_CASES_MEMORY_TYPE) + if schema is None or not schema.enabled: + raise RuntimeError("cases memory schema is not available") + + extract_context = ExtractContext([]) + fields = _case_to_memory_fields(case) + uri = generate_uri( + memory_type=schema, + fields=fields, + user_space=getattr(getattr(ctx, "user", None), "user_id", None) or "default", + extract_context=extract_context, + ) + old_file = None + try: + raw = await viking_fs.read_file(uri, ctx=ctx) + if raw: + old_file = MemoryFileUtils.read(raw, uri=uri) + except Exception: + old_file = None + + source = MemoryOperationSource( + extraction_id=(archive_uri.rstrip("/").rsplit("/", 1)[-1] if archive_uri else ""), + archive_uri=archive_uri or None, + trace_id=tracer.get_trace_id() or None, + ) + operations = ResolvedOperations( + upsert_operations=[ + ResolvedOperation( + old_memory_file_content=old_file, + memory_fields=fields, + memory_type=_CASES_MEMORY_TYPE, + uris=[uri], + source=source, + ) + ], + delete_file_contents=[], + errors=[], + ) + updater = self._get_or_create_updater(registry, transaction_handle=None) + result = await updater.apply_operations( + operations, + ctx, + extract_context=extract_context, + isolation_handler=MemoryIsolationHandler( + ctx, + extract_context, + allowed_memory_types={_CASES_MEMORY_TYPE}, + ), + ) + memory_diff = None + if archive_uri: + memory_diff = await self._build_memory_diff( + result=result, + operations=operations, + viking_fs=viking_fs, + ctx=ctx, + archive_uri=archive_uri, + ) + tracer.info( + "Training CaseSpec fast path wrote case memory: " + f"case={case.name} uri={uri} written={result.written_uris} edited={result.edited_uris}" + ) + return _V3AppliedMemory(result=result, operations=operations, memory_diff=memory_diff) + + @tracer("train.compressor_v3.extract_user_memories", ignore_result=True, ignore_args=True) + async def _extract_user_memories( + self, + messages: List[Message], + user: Optional[Any] = None, + session_id: Optional[str] = None, + ctx: Optional[RequestContext] = None, + strict_extract_errors: bool = False, + latest_archive_overview: str = "", + archive_uri: Optional[str] = None, + allowed_memory_types: Optional[set[str]] = None, + allow_self_memory: bool = True, + allowed_peer_ids: Optional[set[str]] = None, + ) -> "_V3ExtractionResult": + del user + if not messages: + return _V3ExtractionResult() + if not ctx: + logger.warning("No RequestContext provided, skipping v3 memory extraction") + return _V3ExtractionResult() + + try: + viking_fs = get_viking_fs() + except Exception: + logger.warning("VikingFS unavailable, skipping v3 memory extraction", exc_info=True) + return _V3ExtractionResult() + + registry = create_default_registry() + if allow_self_memory: + await registry.initialize_memory_files(ctx) + + extract_context = ExtractContext(messages) + isolation_handler = MemoryIsolationHandler( + ctx, + extract_context, + allowed_memory_types=allowed_memory_types, + allow_self=allow_self_memory, + allowed_peer_ids=allowed_peer_ids, + ) + isolation_handler.prepare_messages() + + orchestrator = self._get_or_create_react( + ctx=ctx, + messages=messages, + latest_archive_overview=latest_archive_overview, + isolation_handler=isolation_handler, + transaction_handle=None, + ) + operations, _tools_used = await orchestrator.run() + if operations is None: + tracer.info("[v3_patch_merge] No memory operations generated") + return _V3ExtractionResult() + + extraction_id = uuid4().hex + extracted_at = datetime.now(timezone.utc).isoformat() + + updater = await get_streaming_memory_updater( + key=make_streaming_memory_updater_key(request_context=ctx), + registry=registry, + vikingdb=self.vikingdb, + config=self.streaming_memory_updater_config, + ) + update_result = await updater.submit( + MemoryUpdateRequest( + operations=operations, + messages=list(messages), + ctx=ctx, + strict_extract_errors=strict_extract_errors, + isolation_options={ + "allowed_memory_types": allowed_memory_types, + "allow_self": allow_self_memory, + "allowed_peer_ids": allowed_peer_ids, + }, + metadata={ + "source_extraction_id": extraction_id, + "session_id": session_id, + "archive_uri": archive_uri, + "trace_id": tracer.get_trace_id(), + "extracted_at": extracted_at, + }, + ) + ) + + result = update_result.apply_result + patch_operations = update_result.operations + + memory_diff = None + if archive_uri and viking_fs and result is not None: + memory_diff = await self._build_memory_diff( + result=result, + operations=patch_operations, + viking_fs=viking_fs, + ctx=ctx, + archive_uri=archive_uri, + ) + + contexts = _contexts_from_update_result(result) + canonical_cases = await _canonical_cases_from_update_result( + operations=patch_operations, + result=result, + viking_fs=viking_fs, + ctx=ctx, + ) + return _V3ExtractionResult( + contexts=contexts, + cases=canonical_cases, + memory_diff=memory_diff, + case_uri_by_name=_case_uri_by_name(canonical_cases, patch_operations, result), + ) + + @tracer("train.compressor_v3.train_from_extracted_cases", ignore_result=True, ignore_args=True) + async def train_from_extracted_cases( + self, + *, + cases: list[Case], + messages: list[Message], + ctx: Optional[RequestContext], + case_uri_by_name: dict[str, str] | None = None, + session_id: Optional[str] = None, + archive_uri: str = "", + strict_extract_errors: bool = False, + collect_memory_diff: bool = False, + ) -> dict[str, Any]: + if not messages or ctx is None: + return {"case_count": 0, "submitted": 0, "reason": "missing_messages_or_ctx"} + if not cases: + tracer.info("No commit training case memories extracted; skipping streaming train") + return {"case_count": 0, "submitted": 0} + + config = get_openviking_config() + skill_enabled = ( + config.memory.session_skill_extraction_enabled and self.skill_processor is not None + ) + + try: + viking_fs = get_viking_fs() + + # --- Experience streaming trainer --- + exp_root_uri = _experience_root_uri(ctx) + exp_policy_set = await ExperienceSetLoader(viking_fs=viking_fs).load( + exp_root_uri, + ctx=ctx, + ) + optimizer_context = PatchMergePolicyOptimizerContext(request_context=ctx) + gradient_context = ExperienceGradientContext( + request_context=ctx, + messages=list(messages), + strict_extract_errors=strict_extract_errors, + ) + analysis_context = TrajectoryAnalyzerContext( + request_context=ctx, + strict_extract_errors=strict_extract_errors, + include_session_skills=skill_enabled, + ) + exp_trainer = await get_streaming_policy_trainer( + key=make_streaming_policy_trainer_key( + policy_root_uri=exp_root_uri, + request_context=ctx, + ), + policy_set=exp_policy_set, + rollout_analyzer=self.rollout_analyzer, + gradient_estimator=ExperienceGradientEstimator( + viking_fs=viking_fs, + ), + policy_optimizer=PatchMergePolicyOptimizer( + viking_fs=viking_fs, + memory_type="experiences", + ), + policy_updater=MemoryFilePolicyUpdater(viking_fs=viking_fs, vikingdb=self.vikingdb), + context=PipelineContext( + analysis_context=analysis_context, + gradient_context=gradient_context, + optimization_context=optimizer_context, + apply_context=ctx, + ), + config=self.streaming_trainer_config, + ) + + # --- Skill streaming trainer --- + skill_trainer = None + if skill_enabled: + skill_root_uri = _skill_root_uri(ctx) + skill_policy_set = await SkillSetLoader(viking_fs=viking_fs).load( + skill_root_uri, + ctx=ctx, + ) + skill_trainer = await get_streaming_policy_trainer( + key=_skill_trainer_key(ctx), + policy_set=skill_policy_set, + rollout_analyzer=self.rollout_analyzer, + gradient_estimator=_NoopGradientEstimator(), + policy_optimizer=PatchMergePolicyOptimizer( + viking_fs=viking_fs, + memory_type="skills", + ), + policy_updater=SkillPolicyUpdater( + skill_processor=self.skill_processor, + viking_fs=viking_fs, + vikingdb=self.vikingdb, + memory_type="skills", + ), + context=PipelineContext( + analysis_context=analysis_context, + gradient_context=gradient_context, + optimization_context=optimizer_context, + apply_context=ctx, + ), + config=self.streaming_trainer_config, + ) + + submitted = 0 + skill_submitted = 0 + skill_uris: list[str] = [] + filtered_exp_gradient_count = 0 + memory_diffs: list[dict[str, Any]] = [] + policy_snapshot_id = _commit_policy_snapshot_id( + session_id=session_id, + archive_uri=archive_uri, + ) + + case_uri_map = dict(case_uri_by_name or {}) + + for case in cases: + case_uri = _case_uri_for_case(case, case_uri_map) + rollout = Rollout( + case=case, + messages=list(messages), + policy_snapshot_id=policy_snapshot_id, + ) + # Analyze once — trajectories + skill patches co-extracted + analysis = await self.rollout_analyzer.analyze(rollout, analysis_context) + + # Experience path: estimate gradients, then submit to exp trainer + exp_gradients = await ExperienceGradientEstimator( + viking_fs=viking_fs, + ).estimate(analysis, exp_trainer.policy_set, gradient_context) + exp_training_result = _trajectory_only_training_result( + analysis=analysis, + rollout=rollout, + policy_set=exp_trainer.policy_set, + ) + if exp_gradients: + exp_training_result = await exp_trainer.submit_gradients( + exp_gradients, + analysis=analysis, + rollout=rollout, + ) + if case_uri: + await self._link_case_to_training_outputs( + analysis=analysis, + case_uri=case_uri, + plan=exp_training_result.plan, + apply_result=exp_training_result.apply_result, + ctx=ctx, + viking_fs=viking_fs, + ) + # Skill path: co-extracted skill gradients go directly to skill trainer + if skill_trainer is not None and analysis.gradients: + skill_gradients = [ + g for g in analysis.gradients if _gradient_memory_type(g) == "skills" + ] + if skill_gradients: + skill_training_result = await skill_trainer.submit_gradients( + skill_gradients, + analysis=analysis, + rollout=rollout, + ) + skill_submitted += 1 + apply_result = getattr(skill_training_result, "apply_result", None) + if apply_result is not None: + for uri in getattr(apply_result, "written_uris", []) or []: + if uri: + skill_uris.append(str(uri)) + + submitted += 1 + + if collect_memory_diff: + # Build diff from the strongly typed training result returned by + # submit_gradients. Do not use exp_trainer.last_apply_result here: + # it is only a PolicyApplyResult and does not carry analyses/plan, + # so trajectory and experience diffs would be lost. + memory_diff = await self._build_training_memory_diff( + training_result=exp_training_result, + viking_fs=viking_fs, + ctx=ctx, + archive_uri=archive_uri, + ) + if _memory_diff_has_changes(memory_diff): + memory_diffs.append(memory_diff) + + tracer.info( + "Submitted commit case memories to streaming train: " + f"case_count={len(cases)} submitted={submitted} " + f"skill_submitted={skill_submitted}", + console=self.streaming_trainer_config.trace_console, + ) + response: dict[str, Any] = { + "case_count": len(cases), + "submitted": submitted, + "skill_submitted": skill_submitted, + "skill_uris": skill_uris, + "filtered_exp_gradient_count": filtered_exp_gradient_count, + } + if collect_memory_diff: + response["memory_diff"] = _merge_memory_diffs( + memory_diffs, + archive_uri=archive_uri, + ) + return response + except Exception as exc: + logger.warning("Commit streaming train failed: %s", exc, exc_info=True) + if strict_extract_errors: + raise + return {"case_count": len(cases), "submitted": 0, "error": str(exc)} + + async def _build_training_memory_diff( + self, + *, + training_result: RolloutTrainingResult, + viking_fs: Any, + ctx: RequestContext, + archive_uri: str, + ) -> dict[str, Any]: + adds: list[dict[str, Any]] = [] + updates: list[dict[str, Any]] = [] + deletes: list[dict[str, Any]] = [] + + seen_trajectory_uris: set[str] = set() + for analysis in training_result.analyses: + for trajectory in analysis.trajectories: + uri = trajectory.uri + if not uri or uri in seen_trajectory_uris: + continue + seen_trajectory_uris.add(uri) + adds.append( + { + "uri": uri, + "memory_type": "trajectories", + "after": trajectory.content, + } + ) + + applied_uris = set(training_result.apply_result.written_uris) + deleted_uris = set(training_result.apply_result.deleted_uris) + root_uri = training_result.apply_result.updated_policy_set.root_uri + if not root_uri: + raise ValueError( + "PolicyApplyResult.updated_policy_set.root_uri is required for training memory diff" + ) + source_trajectory_uris = set(seen_trajectory_uris) + + for item in training_result.plan.items: + if item.memory_type != "experiences": + continue + if source_trajectory_uris and not _plan_item_has_source_trajectory( + item, + source_trajectory_uris, + ): + continue + uri = _experience_plan_item_uri(item, root_uri) + if not uri: + continue + if item.kind == "delete": + if uri in deleted_uris: + deletes.append( + { + "uri": uri, + "memory_type": "experiences", + "deleted_content": item.before_content or "", + } + ) + continue + if item.kind != "upsert" or uri not in applied_uris: + continue + after = await _read_plain_memory_content( + viking_fs, + uri=uri, + ctx=ctx, + fallback=item.after_content or "", + ) + before = item.before_content + if before is None: + adds.append({"uri": uri, "memory_type": "experiences", "after": after}) + else: + # Filter no-op experience updates the same way as user-memory + # updates: a patch that re-serializes to identical content should + # not appear in memory_diff.json. + try: + old_file = MemoryFileUtils.read(before, uri=uri) if before else None + new_file = MemoryFileUtils.read(after, uri=uri) if after else None + except Exception: + old_file, new_file = None, None + if old_file is not None and _same_memory_file(old_file, new_file): + logger.info("Skipping unchanged experience memory in memory_diff.json: %s", uri) + continue + updates.append( + { + "uri": uri, + "memory_type": "experiences", + "before": before, + "after": after, + } + ) + + return _make_memory_diff( + archive_uri=archive_uri, + adds=adds, + updates=updates, + deletes=deletes, + ) + + async def _link_case_to_training_outputs( + self, + *, + analysis: RolloutAnalysis, + case_uri: str, + plan: PolicyUpdatePlan, + apply_result: PolicyApplyResult, + ctx: RequestContext, + viking_fs: Any, + ) -> None: + links = _case_training_links( + analysis=analysis, + case_uri=case_uri, + plan=plan, + apply_result=apply_result, + ) + if not links: + return + await _render_case_links_from_template( + case_uri=case_uri, + links=links, + ctx=ctx, + viking_fs=viking_fs, + ) + await write_stored_links(links, ctx, viking_fs, skip_uris={case_uri}) + + async def _write_final_memory_diff( + self, + *, + archive_uri: str, + ctx: Optional[RequestContext], + memory_diffs: list[Any], + ) -> None: + if not archive_uri or ctx is None: + return + merged = _merge_memory_diffs( + [diff for diff in memory_diffs if isinstance(diff, dict)], + archive_uri=archive_uri, + ) + viking_fs = get_viking_fs() + if viking_fs is None: + return + await viking_fs.write_file( + uri=f"{archive_uri.rstrip('/')}/memory_diff.json", + content=json.dumps(merged, ensure_ascii=False, indent=4), + ctx=ctx, + ) + + +@dataclass(slots=True) +class _V3ExtractionResult: + contexts: list[Context] = field(default_factory=list) + cases: list[Case] = field(default_factory=list) + memory_diff: dict[str, Any] | None = None + case_uri_by_name: dict[str, str] = field(default_factory=dict) + + +@dataclass(slots=True) +class _V3AppliedMemory: + result: Any + operations: ResolvedOperations + memory_diff: dict[str, Any] | None = None + + +def _contexts_from_update_result(result: Any) -> list[Context]: + contexts = [] + for uri in result.written_uris: + contexts.append(Context(uri=uri, category="memory_write", context_type="memory")) + for uri in result.edited_uris: + contexts.append(Context(uri=uri, category="memory_edit", context_type="memory")) + for uri in result.deleted_uris: + contexts.append(Context(uri=uri, category="memory_delete", context_type="memory")) + return contexts + + +def _training_case_from_first_message( + messages: list[Message], + allowed_memory_types: Optional[set[str]], +) -> Case | None: + """Parse a batch-training CaseSpec control message from message[0]. + + The fast path is deliberately gated by the commit memory policy so normal + user sessions never bypass user-memory extraction. Once the protocol + header is present, malformed payloads are treated as errors instead of + silently falling back to LLM extraction. + """ + if not messages or allowed_memory_types is None: + return None + if not set(allowed_memory_types).issubset(_TRAINING_FAST_PATH_MEMORY_TYPES): + return None + + payload = _training_case_spec_payload_from_message(messages[0]) + if payload is None: + return None + return _case_from_payload(payload) + + +def _training_case_spec_payload_from_message(message: Message) -> dict[str, Any] | None: + text = _message_text(message).strip() + if not text.startswith(_TRAINING_CASE_SPEC_HEADER): + return None + return _parse_training_case_spec_payload(text) + + +def _message_text(message: Message) -> str: + content = getattr(message, "content", "") + if content: + return str(content) + texts: list[str] = [] + for part in getattr(message, "parts", []) or []: + text = getattr(part, "text", None) + if text: + texts.append(str(text)) + return "\n".join(texts) + + +def _training_messages_after_case_spec(messages: list[Message]) -> list[Message]: + """Return commit messages after CaseSpec.""" + return list(messages[1:]) + + +def _parse_training_case_spec_payload(text: str) -> dict[str, Any]: + match = _JSON_FENCE_RE.search(text) + raw_payload = ( + match.group(1).strip() if match else text.removeprefix(_TRAINING_CASE_SPEC_HEADER).strip() + ) + if not raw_payload: + raise ValueError("Training CaseSpec fast path payload is empty") + try: + payload = JsonUtils.loads(raw_payload) + except Exception as exc: + raise ValueError("Training CaseSpec fast path payload is not valid JSON") from exc + if not isinstance(payload, dict): + raise ValueError("Training CaseSpec fast path payload must be a JSON object") + protocol = str(payload.get("protocol") or "") + if protocol != _TRAINING_CASE_SPEC_PROTOCOL: + raise ValueError( + "Training CaseSpec fast path protocol mismatch: " + f"expected {_TRAINING_CASE_SPEC_PROTOCOL!r}, got {protocol!r}" + ) + if not isinstance(payload.get("case"), dict): + raise ValueError("Training CaseSpec fast path payload must contain a case object") + return payload + + +def _case_from_payload(payload: dict[str, Any]) -> Case: + raw_case = payload["case"] + name = str(raw_case.get("name") or "").strip() + task_signature = str(raw_case.get("task_signature") or "").strip() + if not name: + raise ValueError("Training CaseSpec case.name is required") + if not task_signature: + raise ValueError("Training CaseSpec case.task_signature is required") + case_input = raw_case.get("input") + if not isinstance(case_input, dict): + raise ValueError("Training CaseSpec case.input must be an object") + rubric = _rubric_from_payload(raw_case.get("rubric"), fallback_name=f"{name}_rubric") + metadata = raw_case.get("metadata") if isinstance(raw_case.get("metadata"), dict) else {} + return Case( + name=name, + task_signature=task_signature, + input=dict(case_input), + rubric=rubric, + metadata={ + "source": "batch_training_case_spec", + **dict(metadata), + }, + ) + + +def _rubric_from_payload(raw_rubric: Any, *, fallback_name: str) -> Rubric: + if not isinstance(raw_rubric, dict): + raise ValueError("Training CaseSpec case.rubric must be an object") + raw_criteria = raw_rubric.get("criteria") + if not isinstance(raw_criteria, list) or not raw_criteria: + raise ValueError("Training CaseSpec case.rubric.criteria must be a non-empty list") + + criteria: list[RubricCriterion] = [] + for index, item in enumerate(raw_criteria): + if not isinstance(item, dict): + raise ValueError("Training CaseSpec rubric criteria must be objects") + name = str(item.get("name") or f"criterion_{index + 1}").strip() + description = str(item.get("description") or "").strip() + if not description: + raise ValueError("Training CaseSpec rubric criterion.description is required") + criteria.append( + RubricCriterion( + name=name, + description=description, + required=bool(item.get("required", True)), + weight=_safe_weight(item.get("weight"), default=1.0), + metadata=dict(item.get("metadata") or {}) + if isinstance(item.get("metadata"), dict) + else {}, + ) + ) + + return Rubric( + name=str(raw_rubric.get("name") or fallback_name), + description=str( + raw_rubric.get("description") or "Defines what good means for this batch training case." + ), + criteria=criteria, + metadata=dict(raw_rubric.get("metadata") or {}) + if isinstance(raw_rubric.get("metadata"), dict) + else {}, + ) + + +def _case_to_memory_fields(case: Case) -> dict[str, Any]: + return { + "case_name": case.name, + "task_signature": case.task_signature, + "input": json.dumps(case.input or {}, ensure_ascii=False, sort_keys=True), + "rubric": json.dumps( + _rubric_to_payload(case.rubric), + ensure_ascii=False, + sort_keys=True, + ), + "evidence": _case_evidence(case), + } + + +def _rubric_to_payload(rubric: Rubric) -> dict[str, Any]: + return { + "name": rubric.name, + "description": rubric.description, + "criteria": [ + { + "name": criterion.name, + "description": criterion.description, + "required": criterion.required, + "weight": criterion.weight, + } + for criterion in rubric.criteria + ], + } + + +def _case_evidence(case: Case) -> str: + raw_evidence = (case.metadata or {}).get("evidence") + if raw_evidence: + return str(raw_evidence) + return "Structured batch training CaseSpec supplied by the training pipeline." + + +def _operations_to_cases(operations: ResolvedOperations) -> list[Case]: + cases: list[Case] = [] + for op in getattr(operations, "upsert_operations", []) or []: + if getattr(op, "memory_type", None) != _CASES_MEMORY_TYPE: + continue + case = _operation_to_case(op) + if case is not None: + cases.append(case) + return cases + + +async def _canonical_cases_from_update_result( + *, + operations: ResolvedOperations, + result: Any, + viking_fs: Any, + ctx: RequestContext, +) -> list[Case]: + """Build training cases from canonical case files after patch merge/apply. + + The extractor can emit multiple case proposals, but ``StreamingMemoryUpdater`` + may merge, rename, or deduplicate them before storage. Training must follow + the actual persisted case-first state, so only cases whose canonical URI was + written/edited by this update are returned. + """ + + touched_uris = _case_result_touched_uris(result) + if not touched_uris: + return [] + + case_ops_by_uri: dict[str, ResolvedOperation] = {} + for op in getattr(operations, "upsert_operations", []) or []: + if getattr(op, "memory_type", None) != _CASES_MEMORY_TYPE: + continue + for uri in getattr(op, "uris", []) or []: + if uri in touched_uris and uri not in case_ops_by_uri: + case_ops_by_uri[uri] = op + + cases: list[Case] = [] + for uri in touched_uris: + if uri not in case_ops_by_uri: + continue + case = await _case_from_persisted_memory_file(uri=uri, viking_fs=viking_fs, ctx=ctx) + if case is None: + case = _operation_to_case( + case_ops_by_uri[uri].model_copy(update={"uris": [uri]}, deep=True) + ) + if case is not None: + cases.append(case) + return cases + + +def _case_result_touched_uris(result: Any) -> list[str]: + uris = list(getattr(result, "written_uris", []) or []) + list( + getattr(result, "edited_uris", []) or [] + ) + return [ + uri + for uri in dict.fromkeys(str(uri) for uri in uris if uri) + if _uri_is_case_memory_file(uri) + ] + + +def _uri_is_case_memory_file(uri: str) -> bool: + return "/memories/cases/" in str(uri or "") and not str(uri).rstrip("/").endswith( + ("/.overview.md", "/.abstract.md") + ) + + +async def _case_from_persisted_memory_file( + *, + uri: str, + viking_fs: Any, + ctx: RequestContext, +) -> Case | None: + try: + raw = await viking_fs.read_file(uri, ctx=ctx) + except Exception as exc: + tracer.info(f"Failed to read canonical case memory for training {uri}: {exc}") + return None + try: + memory_file = MemoryFileUtils.read(raw or "", uri=uri) + except Exception as exc: + tracer.info(f"Failed to parse canonical case memory for training {uri}: {exc}") + return None + return _memory_file_to_case(memory_file) + + +def _memory_file_to_case(memory_file: MemoryFile) -> Case | None: + fields = dict(getattr(memory_file, "extra_fields", {}) or {}) + uri = str(getattr(memory_file, "uri", "") or "") + name = str(fields.get("case_name") or fields.get("name") or _basename_case_name(uri)).strip() + task_signature = str(fields.get("task_signature") or name).strip() + if not name or not task_signature: + return None + memory_fields = dict(fields) + if uri: + memory_fields.setdefault("uri", uri) + return Case( + name=name, + task_signature=task_signature, + input=_parse_case_input(fields.get("input")), + rubric=_parse_rubric(fields.get("rubric"), fallback_name=f"{name}_rubric"), + metadata={ + "source": "session_commit_case_memory", + "case_uris": [uri] if uri else [], + "evidence": str(fields.get("evidence") or ""), + "memory_fields": memory_fields, + }, + ) + + +def _operation_to_case(op: ResolvedOperation) -> Case | None: + fields = dict(getattr(op, "memory_fields", {}) or {}) + name = str(fields.get("case_name") or fields.get("name") or _fallback_case_name(op)).strip() + task_signature = str(fields.get("task_signature") or name).strip() + if not name or not task_signature: + return None + return Case( + name=name, + task_signature=task_signature, + input=_parse_case_input(fields.get("input")), + rubric=_parse_rubric(fields.get("rubric"), fallback_name=f"{name}_rubric"), + metadata={ + "source": "session_commit_case_memory", + "case_uris": list(getattr(op, "uris", []) or []), + "evidence": str(fields.get("evidence") or ""), + "memory_fields": fields, + }, + ) + + +def _parse_case_input(raw_input: Any) -> dict[str, Any]: + parsed = _parse_jsonish(raw_input) + if isinstance(parsed, dict): + return parsed + if isinstance(parsed, list): + return {"items": parsed} + return {"summary": str(raw_input or "")} + + +def _parse_rubric(raw_rubric: Any, *, fallback_name: str) -> Rubric: + parsed = _parse_jsonish(raw_rubric) + raw = parsed if isinstance(parsed, dict) else {} + raw_criteria = raw.get("criteria") if isinstance(raw.get("criteria"), list) else [] + criteria: list[RubricCriterion] = [] + for index, item in enumerate(raw_criteria): + if not isinstance(item, dict): + continue + criteria.append( + RubricCriterion( + name=str(item.get("name") or f"criterion_{index + 1}"), + description=str(item.get("description") or "The rollout satisfies the task."), + required=bool(item.get("required", True)), + weight=_safe_weight(item.get("weight"), default=1.0), + ) + ) + if not criteria: + criteria = [ + RubricCriterion( + name="task_completed", + description="The assistant completed the user's task with verifiable outcome.", + required=True, + weight=1.0, + ) + ] + return Rubric( + name=str(raw.get("name") or fallback_name), + description=str(raw.get("description") or "Defines what good means for this commit case."), + criteria=criteria, + ) + + +def _parse_jsonish(value: Any) -> Any: + if isinstance(value, (dict, list)): + return value + if not isinstance(value, str) or not value.strip(): + return None + try: + return JsonUtils.loads(value) + except Exception: + return None + + +def _safe_weight(value: Any, *, default: float) -> float: + try: + parsed = float(value) + except (TypeError, ValueError): + return default + return parsed if parsed > 0 else default + + +def _first_uri(uris: list[str] | None) -> str | None: + return uris[0] if uris else None + + +def _fallback_case_name(op: ResolvedOperation) -> str: + uri = _first_uri(getattr(op, "uris", []) or []) + name = _basename_case_name(uri) + if name: + return name + return "commit_case" + + +def _basename_case_name(uri: str | None) -> str: + if uri: + return uri.rstrip("/").split("/")[-1].removesuffix(".md") + return "" + + +def _user_space_from_ctx(ctx: RequestContext, *, purpose: str) -> str: + user_space = getattr(getattr(ctx, "user", None), "user_id", None) + if not user_space: + raise ValueError(f"RequestContext.user.user_id is required for {purpose}") + return str(user_space) + + +def _experience_root_uri(ctx: RequestContext) -> str: + user_space = _user_space_from_ctx(ctx, purpose="experience memory root URI") + return f"viking://user/{user_space}/memories/experiences" + + +def _skill_root_uri(ctx: RequestContext) -> str: + user_space = _user_space_from_ctx(ctx, purpose="skill root URI") + return f"viking://user/{user_space}/skills" + + +def _skill_trainer_key(ctx: RequestContext) -> tuple[str, str, str]: + """Registry key for the skill streaming trainer (separate from exp trainer).""" + from openviking.session.train.components.policy_trainer import ( + make_streaming_policy_trainer_key, + ) + + return make_streaming_policy_trainer_key( + policy_root_uri=_skill_root_uri(ctx), + request_context=ctx, + ) + + +@dataclass(slots=True) +class _NoopGradientEstimator: + """GradientEstimator that returns empty gradients. + + Used for the skill trainer because skill gradients are co-extracted + during trajectory analysis and submitted directly via + ``submit_gradients``; the estimator is never called in practice but + ``StreamingPolicyTrainer`` requires one. + """ + + async def estimate( + self, + analysis: Any, + policy_set: Any, + context: Any = None, + ) -> list[Any]: + return [] + + +def _case_uri_for_case(case: Case, case_uri_by_name: dict[str, str]) -> str: + if case.name in case_uri_by_name: + return case_uri_by_name[case.name] + uris = (case.metadata or {}).get("case_uris") + if isinstance(uris, list) and uris: + return str(uris[0]) + fields = (case.metadata or {}).get("memory_fields") + if isinstance(fields, dict): + uri = fields.get("uri") + if uri: + return str(uri) + return "" + + +def _case_uri_by_name( + cases: list[Case], + operations: ResolvedOperations, + result: Any, +) -> dict[str, str]: + candidates = set( + (getattr(result, "written_uris", []) or []) + (getattr(result, "edited_uris", []) or []) + ) + mapping: dict[str, str] = {} + for op in getattr(operations, "upsert_operations", []) or []: + if getattr(op, "memory_type", None) != _CASES_MEMORY_TYPE: + continue + fields = dict(getattr(op, "memory_fields", {}) or {}) + name = str(fields.get("case_name") or fields.get("name") or "").strip() + if not name: + continue + for uri in getattr(op, "uris", []) or []: + if not candidates or uri in candidates: + mapping[name] = uri + break + for case in cases: + if case.name not in mapping: + uri = _case_uri_for_case(case, {}) + if uri: + mapping[case.name] = uri + return mapping + + +def _first_context_uri(contexts: list[Context]) -> str: + for context in contexts or []: + uri = getattr(context, "uri", "") + if uri: + return str(uri) + return "" + + +def _case_training_links( + *, + analysis: RolloutAnalysis, + case_uri: str, + plan: PolicyUpdatePlan, + apply_result: PolicyApplyResult, +) -> list[StoredLink]: + trajectory_links = _case_trajectory_links(analysis=analysis, case_uri=case_uri) + trajectory_uris = {link.to_uri for link in trajectory_links if link.to_uri} + experience_links = _case_experience_links_via_trajectories( + case_uri=case_uri, + trajectory_uris=trajectory_uris, + plan=plan, + apply_result=apply_result, + ) + return merge_link_lists([*trajectory_links, *experience_links]) + + +def _case_trajectory_links( + *, + analysis: RolloutAnalysis, + case_uri: str, +) -> list[StoredLink]: + links: list[StoredLink] = [] + for trajectory in getattr(analysis, "trajectories", []) or []: + uri = str(getattr(trajectory, "uri", "") or "") + if not uri or "/memories/trajectories/" not in uri: + continue + links.append( + _stored_link( + from_uri=case_uri, + target_uri=uri, + link_type="related_to", + description="", + ) + ) + return links + + +def _case_experience_links_via_trajectories( + *, + case_uri: str, + trajectory_uris: set[str], + plan: PolicyUpdatePlan, + apply_result: PolicyApplyResult, +) -> list[StoredLink]: + if not trajectory_uris: + return [] + touched = set(getattr(apply_result, "written_uris", []) or []) + touched.update(getattr(apply_result, "edited_uris", []) or []) + result: list[StoredLink] = [] + seen: set[str] = set() + root_uri = getattr(getattr(apply_result, "updated_policy_set", None), "root_uri", "") + if not root_uri: + raise ValueError( + "PolicyApplyResult.updated_policy_set.root_uri is required for case-to-experience links" + ) + for item in getattr(plan, "items", []) or []: + if item.memory_type != "experiences" or item.kind != "upsert": + continue + if not _plan_item_has_source_trajectory(item, trajectory_uris): + continue + uri = _experience_plan_item_uri(item, root_uri) + if touched and uri not in touched: + continue + if uri in seen: + continue + seen.add(uri) + result.append( + _stored_link( + from_uri=case_uri, + target_uri=uri, + link_type="related_to", + description="", + ) + ) + return result + + +def _plan_item_has_source_trajectory(item: PolicyPlanItem, trajectory_uris: set[str]) -> bool: + for link in getattr(item, "links", []) or []: + try: + stored = link if isinstance(link, StoredLink) else StoredLink(**dict(link)) + except Exception: + continue + if ( + stored.link_type == "derived_from" + and stored.to_uri in trajectory_uris + and "/memories/trajectories/" in str(stored.to_uri or "") + ): + return True + return False + + +def _stored_link( + *, + from_uri: str, + target_uri: str, + link_type: str, + description: str, +) -> StoredLink: + return StoredLink( + from_uri=from_uri, + to_uri=target_uri, + link_type=link_type, + weight=1.0, + match_text=None, + description=description, + created_at=datetime.now(timezone.utc).isoformat(), + ) + + +async def _render_case_links_from_template( + *, + case_uri: str, + links: list[StoredLink], + ctx: RequestContext, + viking_fs: Any, +) -> None: + if not links: + return + try: + raw = await viking_fs.read_file(case_uri, ctx=ctx) + except Exception as exc: + tracer.error(f"Failed to read case memory for link rendering {case_uri}: {exc}") + return + + mf = MemoryFileUtils.read(raw or "", uri=case_uri) + from openviking.session.memory.merge_op.link_merge import merge_links + + merged_links = merge_links(mf.links, [link.model_dump() for link in links]) + if merged_links != mf.links: + mf.links = merged_links + + schema = create_default_registry().get(_CASES_MEMORY_TYPE) + content_template = schema.content_template if schema is not None else None + await viking_fs.write_file( + case_uri, + MemoryFileUtils.write(mf, content_template=content_template), + ctx=ctx, + ) + + +def _gradient_memory_type(gradient: Any) -> str: + """Extract memory_type from a semantic gradient.""" + after_file = getattr(gradient, "after_file", None) + if after_file is not None: + mt = getattr(after_file, "memory_type", None) + if mt: + return str(mt) + metadata = dict(getattr(gradient, "metadata", {}) or {}) + if metadata.get("memory_type"): + return str(metadata["memory_type"]) + before_file = getattr(gradient, "before_file", None) + if before_file is not None: + mt = getattr(before_file, "memory_type", None) + if mt: + return str(mt) + return "experiences" + + +def _trajectory_only_training_result( + *, + analysis: RolloutAnalysis, + rollout: Rollout, + policy_set: Any, +) -> RolloutTrainingResult: + """Return a typed no-op training result that still carries trajectories. + + Some rollouts produce useful trajectory memories but no experience gradients. + The memory diff should still include those trajectory writes, so callers use + this as the baseline result and replace it only when experience training + returns a full RolloutTrainingResult. + """ + + return RolloutTrainingResult( + analyses=[analysis], + gradients=[], + plan=PolicyUpdatePlan(items=[], metadata={"no_experience_gradients": True}), + apply_result=PolicyApplyResult( + updated_policy_set=policy_set, + written_uris=[], + errors=[], + metadata={"no_experience_gradients": True}, + ), + metadata={ + "source": "trajectory_only", + "case_name": rollout.case.name, + "trajectory_count": len(analysis.trajectories), + }, + ) + + +def _applied_memory_result(value: Any) -> Any: + if isinstance(value, _V3AppliedMemory): + return value.result + result = getattr(value, "result", None) + return result if result is not None else value + + +def _applied_memory_diff(value: Any) -> dict[str, Any] | None: + if isinstance(value, _V3AppliedMemory): + return value.memory_diff + memory_diff = getattr(value, "memory_diff", None) + return memory_diff if isinstance(memory_diff, dict) else None + + +def _same_memory_file(before: Optional[MemoryFile], after: Optional[MemoryFile]) -> bool: + """Return whether two parsed memory files represent the same stored memory.""" + if before is None or after is None: + return False + # memory_type is commonly known from the operation/URI even when the raw + # memory file does not serialize it, so do not treat that metadata-only + # representation difference as a real file update. + return before.model_dump(exclude={"uri", "memory_type"}) == after.model_dump( + exclude={"uri", "memory_type"} + ) + + +def _v3_extraction_response( + *, + contexts: list[Context], + train_result: Any, + archive_uri: str, +) -> list[Context] | dict[str, Any]: + """Build the extraction response. + + Historically ``extract_long_term_memories`` returned ``list[Context]`` and + a number of direct callers still index/compare the return value as a list. + Commit orchestration now also understands the execution-memory style + ``{"contexts": ..., "session_skills": ...}`` shape so it can count + session skills. Preserve the old list shape unless there are actual + session skills to report. + """ + skill_dicts: list[dict[str, Any]] = [] + seen: set[str] = set() + if isinstance(train_result, dict): + for uri in train_result.get("skill_uris", []) or []: + uri_str = str(uri or "") + if uri_str and uri_str not in seen: + seen.add(uri_str) + skill_dicts.append({"uri": uri_str, "archive_uri": archive_uri}) + if not skill_dicts: + return contexts + return {"contexts": contexts, "session_skills": skill_dicts} + + +def _make_memory_diff( + *, + archive_uri: str, + adds: list[dict[str, Any]], + updates: list[dict[str, Any]], + deletes: list[dict[str, Any]], +) -> dict[str, Any]: + return { + "archive_uri": archive_uri, + "trace_id": tracer.get_trace_id() or None, + "extracted_at": datetime.utcnow().isoformat() + "Z", + "operations": { + "adds": list(adds), + "updates": list(updates), + "deletes": list(deletes), + }, + "summary": { + "total_adds": len(adds), + "total_updates": len(updates), + "total_deletes": len(deletes), + }, + } + + +def _merge_memory_diffs( + diffs: list[dict[str, Any]], + *, + archive_uri: str, +) -> dict[str, Any]: + adds: list[dict[str, Any]] = [] + updates: list[dict[str, Any]] = [] + deletes: list[dict[str, Any]] = [] + trace_id = tracer.get_trace_id() or None + for diff in diffs: + if not isinstance(diff, dict): + continue + if trace_id is None and diff.get("trace_id"): + trace_id = str(diff.get("trace_id")) + operations = diff.get("operations") + if not isinstance(operations, dict): + continue + adds.extend([item for item in operations.get("adds", []) if isinstance(item, dict)]) + updates.extend([item for item in operations.get("updates", []) if isinstance(item, dict)]) + deletes.extend([item for item in operations.get("deletes", []) if isinstance(item, dict)]) + merged = _make_memory_diff( + archive_uri=archive_uri, + adds=adds, + updates=updates, + deletes=deletes, + ) + merged["trace_id"] = trace_id + return merged + + +def _memory_diff_has_changes(diff: Any) -> bool: + if not isinstance(diff, dict): + return False + summary = diff.get("summary") + if not isinstance(summary, dict): + return False + return any( + int(summary.get(key) or 0) > 0 for key in ("total_adds", "total_updates", "total_deletes") + ) + + +async def _read_plain_memory_content( + viking_fs: Any, + *, + uri: str, + ctx: RequestContext, + fallback: str, +) -> str: + try: + raw = await viking_fs.read_file(uri, ctx=ctx) + return MemoryFileUtils.read(raw, uri=uri).content + except Exception: + return fallback + + +def _experience_plan_item_uri(item: PolicyPlanItem, root_uri: str) -> str: + if item.target_uri: + return item.target_uri + name = item.target_name or "new_experience" + return f"{root_uri.rstrip('/')}/{_safe_experience_filename(name)}.md" + + +_EXPERIENCE_NAME_RE = re.compile(r"[^a-zA-Z0-9_.-]+") + + +def _safe_experience_filename(name: str) -> str: + filename = _EXPERIENCE_NAME_RE.sub("_", name.strip()).strip("._-") + return filename or "new_experience" + + +def _commit_policy_snapshot_id(*, session_id: Optional[str], archive_uri: str) -> str: + if archive_uri: + return f"session-commit:{archive_uri.rstrip('/').rsplit('/', 1)[-1]}" + if session_id: + return f"session-commit:{session_id}" + return f"session-commit:{uuid4().hex}" diff --git a/openviking/session/memory/__init__.py b/openviking/session/memory/__init__.py index 1a0a506f24..5090571ab7 100644 --- a/openviking/session/memory/__init__.py +++ b/openviking/session/memory/__init__.py @@ -20,7 +20,20 @@ from openviking.session.memory.memory_type_registry import MemoryTypeRegistry from openviking.session.memory.memory_updater import MemoryUpdater, MemoryUpdateResult from openviking.session.memory.merge_op import FieldType, MergeOp +from openviking.session.memory.patch_merge_context_provider import ( + PatchMergeContextProvider, + PatchMergePatch, +) from openviking.session.memory.schema_model_generator import SchemaModelGenerator +from openviking.session.memory.streaming_memory_updater import ( + MemoryUpdateRequest, + StreamingMemoryUpdater, + StreamingMemoryUpdaterConfig, + StreamingMemoryUpdateResult, + StreamingMemoryUpdaterKey, + get_streaming_memory_updater, + make_streaming_memory_updater_key, +) from openviking.session.memory.tools import ( MemoryLsTool, MemoryReadTool, @@ -48,6 +61,8 @@ # Operations "MemoryOperations", "StructuredMemoryOperations", + "PatchMergeContextProvider", + "PatchMergePatch", # Registry "MemoryTypeRegistry", # Schema models @@ -55,6 +70,13 @@ # Updater "MemoryUpdater", "MemoryUpdateResult", + "MemoryUpdateRequest", + "StreamingMemoryUpdater", + "StreamingMemoryUpdaterConfig", + "StreamingMemoryUpdaterKey", + "StreamingMemoryUpdateResult", + "get_streaming_memory_updater", + "make_streaming_memory_updater_key", # ExtractLoop "ExtractLoop", # Tools (Tool implementations) diff --git a/openviking/session/memory/agent_experience_context_provider.py b/openviking/session/memory/agent_experience_context_provider.py index b8ac9386c2..c27239487a 100644 --- a/openviking/session/memory/agent_experience_context_provider.py +++ b/openviking/session/memory/agent_experience_context_provider.py @@ -12,6 +12,7 @@ from typing import Any, Dict, List, Optional +from openviking.pyagfs.exceptions import AGFSNotFoundError from openviking.server.identity import RequestContext from openviking.session.memory.dataclass import MemoryFile from openviking.session.memory.session_extract_context_provider import ( @@ -33,9 +34,16 @@ MAX_SOURCE_TRAJS = 3 # max trajectories to load per experience +def _is_directory_not_found_error(exc: Exception) -> bool: + message = str(exc).lower() + return "directory not found" in message or "not_found" in message + + class AgentExperienceContextProvider(SessionExtractContextProvider): """Phase 2 provider: consolidate the new trajectory into experience memories.""" + include_tool_parts_in_conversation = True + def __init__( self, messages: Any, @@ -77,7 +85,7 @@ def instruction(self) -> str: - **One experience per distinct user intent.** If a trajectory covers N different user goals (e.g., cancel + modify + add baggage), output N separate entries — never merge them into one. - **Split over merge.** When in doubt whether two patterns belong together, split them. Only merge with an existing experience when it covers the EXACT same user intent and tool sequence. - **Consistent naming language.** All `experience_name` values in one output must use the same language. -- **Do NOT use `delete_uris`** for experience operations — use `supersedes` instead. +- **Do NOT use `delete_ids`** for experience operations — use `supersedes` instead. - Follow field descriptions in the schema. - Output JSON only. Do not call any tools. @@ -187,8 +195,15 @@ async def prefetch(self) -> List[Dict]: continue fallback_uris.append(uri) candidate_uris = fallback_uris[:SEARCH_TOP_K] + except AGFSNotFoundError: + candidate_uris = [] + except FileNotFoundError: + candidate_uris = [] except Exception as e: - tracer.error(f"Failed to list experiences in {experience_dir}: {e}") + if _is_directory_not_found_error(e): + candidate_uris = [] + else: + tracer.error(f"Failed to list experiences in {experience_dir}: {e}") prefetch_messages: List[Dict[str, Any]] = [self._build_conversation_message()] add_tool_call_pair_to_messages( diff --git a/openviking/session/memory/agent_trajectory_context_provider.py b/openviking/session/memory/agent_trajectory_context_provider.py index 37253f46c6..c60015fed5 100644 --- a/openviking/session/memory/agent_trajectory_context_provider.py +++ b/openviking/session/memory/agent_trajectory_context_provider.py @@ -29,6 +29,9 @@ class AgentTrajectoryContextProvider(SessionExtractContextProvider): """Phase 1 provider: extract trajectories and optional session skills.""" + include_tool_parts_in_conversation = True + split_long_text_messages_for_extraction = False + _SHARED_SKILL_STATE = { "messages", "latest_archive_overview", diff --git a/openviking/session/memory/dataclass.py b/openviking/session/memory/dataclass.py index e28553c098..5ae47478f8 100644 --- a/openviking/session/memory/dataclass.py +++ b/openviking/session/memory/dataclass.py @@ -135,6 +135,40 @@ class StoredLink(BaseModel): created_at: str = "" +class DeleteId(BaseModel): + """Delete request by page_id, optionally remapped to a replacement page_id.""" + + delete_page_id: Annotated[Optional[int], WithJsonSchema({"type": "integer"})] = Field( + ..., description="Page_id of the memory item to delete." + ) + replacement_page_id: Annotated[ + Optional[int], + WithJsonSchema({"anyOf": [{"type": "integer"}, {"type": "null"}]}), + ] = Field( + ..., + description=( + "Replacement page_id that should inherit this deleted page's existing links/backlinks; " + "use null for a pure delete." + ), + ) + + +class MemoryOperationSource(BaseModel): + """Runtime and persisted provenance for one extracted memory operation. + + ``extraction_id`` identifies one archive/commit extraction run, not an + entire chat session. Streaming memory merge uses it to detect batches that + combine patches produced from different extraction snapshots. + """ + + extraction_id: Optional[str] = None + session_id: Optional[str] = None + archive_uri: Optional[str] = None + task_id: Optional[str] = None + trace_id: Optional[str] = None + extracted_at: Optional[str] = None + + # ============================================================================ # Memory Field and Schema Definitions # ============================================================================ @@ -166,9 +200,13 @@ class MemoryTypeSchema(BaseModel): operation_mode: str = Field( "upsert", description="Operation mode: 'upsert' (default), 'add_only', or 'update_only'" ) - agent_only: bool = Field( - False, - description="If true, only used by execution-derived extraction, not long-term memory", + stage: str = Field( + "user", + description="Extraction stage: 'user' for long-term user memory, 'agent' for execution-derived memory.", + ) + peer_enabled: bool = Field( + True, + description="Whether this memory type is stored separately under peer directories.", ) overview_template: Optional[str] = Field( None, description="Overview template for auto-generating .overview.md files" @@ -237,7 +275,12 @@ def from_parsed(cls, uri: Optional[str] = None, parsed: Dict[str, Any] = None) - def to_metadata(self) -> Dict[str, Any]: """Flatten to a dict suitable for serialize_with_metadata.""" - metadata = dict(self.extra_fields) + metadata = { + key: value + for key, value in dict(self.extra_fields).items() + if key not in {"user_id", "user_ids", "_uri"} + } + metadata.setdefault("version", 1) metadata["content"] = self.content if self.links: metadata["links"] = self.links @@ -254,6 +297,7 @@ class ResolvedOperation(BaseModel): memory_type: str # The memory type (e.g., 'tools', 'skills', 'events') uris: List[str] page_id: Optional[int] = None # Temporary page_id for link resolution (not persisted) + source: Optional[MemoryOperationSource] = None def is_edit(self): return self.old_memory_file_content is not None @@ -264,6 +308,7 @@ class ResolvedOperations(BaseModel): delete_file_contents: List[MemoryFile] errors: List[str] resolved_links: List[StoredLink] = Field(default_factory=list) + delete_replacements: Dict[str, str] = Field(default_factory=dict) def has_errors(self) -> bool: return len(self.errors) > 0 @@ -398,7 +443,7 @@ class MemoryOperationsProtocol(Protocol): reasoning: str write_uris: List[Any] edit_uris: List[Any] - delete_uris: List[str] + delete_ids: List[DeleteId] def is_empty(self) -> bool: ... @@ -423,21 +468,21 @@ class StructuredMemoryOperations(FaultTolerantBaseModel): default_factory=list, description="Edit operations with flat data format", ) - delete_uris: List[str] = Field( + delete_ids: List[DeleteId] = Field( default_factory=list, - description="Delete operations as URI strings", + description="Delete operations by page_id, with optional replacement_page_id", ) def is_empty(self) -> bool: """Check if there are any operations.""" - return len(self.write_uris) == 0 and len(self.edit_uris) == 0 and len(self.delete_uris) == 0 + return len(self.write_uris) == 0 and len(self.edit_uris) == 0 and len(self.delete_ids) == 0 def to_legacy_operations(self) -> Dict[str, Any]: """Convert to legacy format (identity for fallback).""" return { "write_uris": self.write_uris, "edit_uris": self.edit_uris, - "delete_uris": self.delete_uris, + "delete_ids": self.delete_ids, } model_config = {"extra": "ignore"} diff --git a/openviking/session/memory/extract_loop.py b/openviking/session/memory/extract_loop.py index 5e6c89bd85..27c57f9711 100644 --- a/openviking/session/memory/extract_loop.py +++ b/openviking/session/memory/extract_loop.py @@ -8,11 +8,13 @@ import asyncio import json +import re from typing import Any, Dict, List, Optional, Tuple from openviking.models.vlm.base import ToolCall, VLMBase from openviking.server.identity import RequestContext from openviking.session.memory.dataclass import ( + DeleteId, MemoryFile, ResolvedOperation, ResolvedOperations, @@ -38,6 +40,45 @@ logger = get_logger(__name__) +_CANNED_REFUSAL_RE = re.compile( + r"(抱歉|不好意思|很遗憾|sorry).{0,20}" + r"(无法|不能|未能|不给|没有找到|未找到|can't|cannot|unable|won't|not able)", + re.IGNORECASE, +) + + +def _looks_like_canned_refusal(content: str) -> bool: + """Return whether a non-JSON LLM response looks like a generic refusal.""" + + text = " ".join(str(content or "").split()) + if not text: + return False + if _CANNED_REFUSAL_RE.search(text): + return True + return any( + phrase in text + for phrase in ( + "您的问题我无法回答", + "您的问题我无法识别", + "我无法回答这个问题", + "我无法给到相关内容", + "这个问题未找到相关结果", + "没有找到相关的结果", + "I can't answer that", + "I cannot answer that", + "I can't help with that", + "I cannot help with that", + ) + ) + + +def _preview_text(content: str, *, limit: int = 240) -> str: + text = " ".join(str(content or "").split()) + if len(text) <= limit: + return text + return text[: limit - 3] + "..." + + class ExtractLoop: """ Simplified ReAct orchestrator for memory updates. @@ -58,6 +99,7 @@ def __init__( ctx: Optional[RequestContext] = None, context_provider: Optional[Any] = None, # ExtractContextProvider isolation_handler: MemoryIsolationHandler = None, + thinking: bool = False, ): """ Initialize the ExtractLoop. @@ -69,6 +111,7 @@ def __init__( max_iterations: Maximum number of ReAct iterations (default: 5) ctx: Request context context_provider: ExtractContextProvider - 必须提供(由 provider 加载 schema) + thinking: Whether to explicitly enable model thinking for this extraction loop. """ self.vlm = vlm self.viking_fs = viking_fs or get_viking_fs() @@ -76,10 +119,13 @@ def __init__( self.max_iterations = max_iterations self.ctx = ctx self.context_provider = context_provider + self.thinking = bool(thinking) # Use provided isolation_handler or create one in run() self._isolation_handler = isolation_handler # Track format error retry (max 1 retry) self._format_retry_count = 0 + self._last_llm_failure_kind: Optional[str] = None + self._last_llm_failure_content: str = "" # Schema 生成器(在 run() 中初始化) self.schema_model_generator = None @@ -134,7 +180,7 @@ async def run(self) -> Tuple[Optional[Any], List[Dict[str, Any]]]: # 预计算 expected_fields config = get_openviking_config() self._link_enabled = config.memory.link_enabled if config.memory else False - self._expected_fields = ["delete_uris"] + self._expected_fields = [] if self._link_enabled: self._expected_fields.append("links") @@ -155,8 +201,6 @@ async def run(self) -> Tuple[Optional[Any], List[Dict[str, Any]]]: json_schema = self._operations_model.model_json_schema() # Build initial messages from provider - import json - schema_str = json.dumps(json_schema, ensure_ascii=False) messages = [] page_id_rules = """ @@ -165,6 +209,8 @@ async def run(self) -> Tuple[Optional[Any], List[Dict[str, Any]]]: - For existing items, use the page_id shown in read/search results. - For new items, assign a unique page_id >= 100. - When editing an existing item, reuse its existing page_id. +- To delete an existing item, add an entry to `delete_ids` using its page_id. +- For canonical merges, set `replacement_page_id` to the surviving page that should inherit the deleted page's existing links/backlinks; for pure deletes, set `replacement_page_id` to null. """ link_rules = "" if self._link_enabled: @@ -193,8 +239,6 @@ async def run(self) -> Tuple[Optional[Any], List[Dict[str, Any]]]: } ) - await self._mark_cache_breakpoint(messages) - # Pre-fetch context via provider tool_call_messages = await self.context_provider.prefetch() messages.extend(tool_call_messages) @@ -269,22 +313,41 @@ async def run(self) -> Tuple[Optional[Any], List[Dict[str, Any]]]: continue break # If no tool calls either, continue to next iteration (don't break!) + failure_kind = self._last_llm_failure_kind or "unknown" + failure_preview = _preview_text(self._last_llm_failure_content) tracer.error( - f"LLM returned neither tool calls nor operations (iteration {iteration}/{max_iterations})" + "LLM returned neither tool calls nor operations " + f"(iteration {iteration}/{max_iterations}) " + f"failure_kind={failure_kind} response_preview={failure_preview!r}" ) # Add format error message if parse failed (max 1 retry) if self._format_retry_count == 0: self._format_retry_count += 1 max_iterations += 1 - tracer.info(f"Extended max_iterations to {max_iterations} for format retry") + retry_reason = ( + "refusal_text" if failure_kind == "refusal_text" else "format_retry" + ) + tracer.info(f"Extended max_iterations to {max_iterations} for {retry_reason}") self._add_format_error_message(messages) - # If it's the last iteration, fail instead of silently treating an - # unparseable final response as "no memory operations". + # If it's the last iteration, treat unparseable response as + # "no memory operations" rather than failing hard. if iteration >= max_iterations: - raise RuntimeError( - "Memory extraction final response could not be parsed as JSON operations" + tracer.info( + "Memory extraction final response could not be parsed as JSON operations " + f"after {max_iterations} iterations — treating as no operations " + f"failure_kind={failure_kind} response_preview={failure_preview!r}" + ) + final_operations = ResolvedOperations( + upsert_operations=[], + delete_file_contents=[], + errors=[ + "Final response could not be parsed as JSON operations " + f"after {max_iterations} iterations " + f"(failure_kind={failure_kind})" + ], ) + break self._disable_tools_for_iteration = True continue @@ -322,7 +385,23 @@ async def resolve_operations(self, operations) -> tuple[ResolvedOperations, List for item in items: item_dict = dict(item) item_dict["memory_type"] = memory_type - self._isolation_handler.fill_identity_fields(item_dict, role_scope=role_scope) + try: + self._isolation_handler.fill_identity_fields( + item_dict, + role_scope=role_scope, + memory_type_schema=schema, + ) + except TypeError as exc: + # Some tests and older custom isolation handlers implement + # the pre-schema signature ``fill_identity_fields(item, + # role_scope=None)``. Keep that extension point working + # while the real handler can still use schema peer settings. + if "memory_type_schema" not in str(exc): + raise + self._isolation_handler.fill_identity_fields( + item_dict, + role_scope=role_scope, + ) page_id = item_dict.pop("page_id", None) resolved_op = ResolvedOperation( @@ -365,20 +444,37 @@ async def resolve_operations(self, operations) -> tuple[ResolvedOperations, List upsert_operations.append(resolved_op) - delete_uris_raw = getattr(operations, "delete_uris", []) or [] - for uri_str in delete_uris_raw: - uri_str = uri_str.strip() - if not uri_str: + delete_ids = self._normalize_delete_ids(getattr(operations, "delete_ids", []) or []) + delete_replacements: dict[str, str] = {} + for delete_id in delete_ids: + if delete_id.delete_page_id is None or page_id_map is None: + continue + delete_uri = page_id_map.resolve(delete_id.delete_page_id) + if not delete_uri: + continue + old_content = self.context_provider.read_file_contents.get(delete_uri) + if not old_content: + continue + delete_file_contents.append(old_content) + + replacement_page_id = delete_id.replacement_page_id + if replacement_page_id is None: continue - old_content = self.context_provider.read_file_contents.get(uri_str) - if old_content: - delete_file_contents.append(old_content) + replacement_uri = page_id_map.resolve(replacement_page_id) + if not replacement_uri: + for op in upsert_operations: + if op.page_id == replacement_page_id and op.uris: + replacement_uri = op.uris[0] + break + if replacement_uri and replacement_uri != delete_uri: + delete_replacements[delete_uri] = replacement_uri raw_links = getattr(operations, "links", None) or [] resolved = ResolvedOperations( upsert_operations=upsert_operations, delete_file_contents=delete_file_contents, errors=errors, + delete_replacements=delete_replacements, ) for op in upsert_operations: @@ -390,6 +486,16 @@ async def resolve_operations(self, operations) -> tuple[ResolvedOperations, List return resolved, raw_links + + def _normalize_delete_ids(self, raw_delete_ids: List[Any]) -> List[DeleteId]: + delete_ids: List[DeleteId] = [] + for raw in raw_delete_ids: + try: + delete_ids.append(DeleteId.model_validate(raw)) + except Exception as e: + tracer.info(f"Skipping invalid delete_ids item: {raw}, error={e}") + return delete_ids + async def finalize_operations(self, operations: ResolvedOperations, raw_links: List) -> None: """Register new page_ids and resolve links after refetch is complete. @@ -556,7 +662,7 @@ async def execute_single_tool_call(idx: int, tool_call): action_tasks = [ execute_single_tool_call(idx, tool_call) for idx, tool_call in enumerate(tool_calls) ] - results = await self._execute_in_parallel(action_tasks) + results = await asyncio.gather(*action_tasks) has_unknown_tool = False @@ -601,9 +707,6 @@ async def _call_llm( Returns: Tuple of (tool_calls, operations) - one will be None, the other set """ - # 标记 cache breakpoint - await self._mark_cache_breakpoint(messages) - # Call LLM with tools - use tools from strategy tools = None tool_choice = None @@ -615,8 +718,11 @@ async def _call_llm( messages=messages, tools=tools, tool_choice=tool_choice, + thinking=self.thinking, ) tracer.info(f"llm_response={response}") + self._last_llm_failure_kind = None + self._last_llm_failure_content = "" # print(f'response={response}') # Log cache hit info if hasattr(response, "usage") and response.usage: @@ -677,8 +783,16 @@ async def _call_llm( ) if error is not None: - print(f"content={content}") - tracer.error(f"Failed to parse memory operations: {error}") + failure_kind = ( + "refusal_text" if _looks_like_canned_refusal(content) else "parse_error" + ) + self._last_llm_failure_kind = failure_kind + self._last_llm_failure_content = content + tracer.error( + "Failed to parse memory operations " + f"failure_kind={failure_kind} error={error} " + f"response_preview={_preview_text(content)!r}" + ) return (None, None) return (None, operations) @@ -686,16 +800,15 @@ async def _call_llm( logger.exception(f"Error parsing operations: {e}") # Case 3: No tool calls and no parsable operations - tracer.error("No tool calls or operations parsed") + self._last_llm_failure_kind = "empty_response" if not content else "parse_error" + self._last_llm_failure_content = content or "" + tracer.error( + "No tool calls or operations parsed " + f"failure_kind={self._last_llm_failure_kind} " + f"response_preview={_preview_text(self._last_llm_failure_content)!r}" + ) return (None, None) - async def _execute_in_parallel( - self, - tasks: List[Any], - ) -> List[Any]: - """Execute tasks in parallel, similar to AgentLoop.""" - return await asyncio.gather(*tasks) - async def _check_unread_existing_files(self, operations: ResolvedOperations) -> Dict: refetch_uris = {} for operation in operations.upsert_operations: @@ -731,7 +844,7 @@ def _add_format_error_message(self, messages: List[Dict[str, Any]]) -> None: def _build_final_operations_skeleton(self) -> Dict[str, List[Any]]: """Build an empty operations object matching the expected flat schema fields.""" - fields = ["delete_uris", *(self._expected_fields or [])] + fields = ["delete_ids", *(self._expected_fields or [])] return {field: [] for field in dict.fromkeys(fields)} def _build_final_operations_instruction(self) -> str: @@ -842,11 +955,3 @@ async def _add_refetch_results_to_messages( "content": "Note: The files above were automatically read because they exist and you didn't read them before deciding to write. Please consider the existing content when making write decisions. You can now output updated operations.", } ) - - async def _mark_cache_breakpoint(self, messages): - # 支持 dict 消息和 object 消息 - # last_msg = messages[-1] - # last_msg["cache_control"] = {"type": "ephemeral"} - - # 暂时注释掉,不确定对所有模型的影响 - pass diff --git a/openviking/session/memory/memory_isolation_handler.py b/openviking/session/memory/memory_isolation_handler.py index 2b5ac81a35..3a4c655bf0 100644 --- a/openviking/session/memory/memory_isolation_handler.py +++ b/openviking/session/memory/memory_isolation_handler.py @@ -5,6 +5,7 @@ from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Set +from openviking.core.peer_id import safe_peer_id from openviking.server.identity import RequestContext from openviking.session.memory.dataclass import MemoryTypeSchema, ResolvedOperation from openviking.session.memory.memory_updater import ExtractContext @@ -14,6 +15,7 @@ logger = get_logger(__name__) _INTERNAL_MEMORY_TYPES = {"session_skills"} +_SELF_PEER_ID = "__self" @dataclass @@ -26,6 +28,8 @@ class RoleScope: def peer_user_space(user_space: str, peer_id: str) -> str: """Return the user-space fragment for memory about a stable peer.""" + if peer_id == _SELF_PEER_ID: + return user_space return f"{user_space}/peers/{peer_id}" @@ -47,9 +51,14 @@ def __init__( if allowed_memory_types is not None else None ) + peer_ids = { + item + for item in (safe_peer_id(item) for item in allowed_peer_ids or set()) + if item and item != _SELF_PEER_ID + } self.allow_self = bool(allow_self) - self.allowed_peer_ids = {item for item in allowed_peer_ids or set() if item} - self.allow_peer = bool(self.allowed_peer_ids) + self.allowed_peer_ids = peer_ids + self.allow_peer = bool(peer_ids) def prepare_messages(self) -> None: """No-op hook kept for the extraction pipeline.""" @@ -59,21 +68,26 @@ def _messages(self) -> List[Any]: messages = getattr(self._extract_context, "messages", None) return messages if isinstance(messages, list) else [] - def _has_self_user_message(self) -> bool: - for msg in self._messages(): - if getattr(msg, "role", None) != "user": - continue - if not getattr(msg, "peer_id", None): - return True - return False - - def _first_peer_id_in_messages(self) -> Optional[str]: - for msg in self._messages(): - peer_id = getattr(msg, "peer_id", None) - if peer_id and self._can_write_peer(peer_id): - return peer_id + def _message_target_id(self, msg: Any) -> Optional[str]: + raw_peer_id = getattr(msg, "peer_id", None) + peer_id = safe_peer_id(raw_peer_id) + if peer_id and self._can_write_peer(peer_id): + return peer_id + if raw_peer_id in (None, "") and self.allow_self: + return _SELF_PEER_ID return None + def _first_target_id_in_messages(self) -> Optional[str]: + targets = [ + target_id + for msg in self._messages() + if (target_id := self._message_target_id(msg)) + ] + for target_id in targets: + if target_id != _SELF_PEER_ID: + return target_id + return targets[0] if targets else None + def get_read_scope(self) -> RoleScope: user_ids = set() peer_ids = set() @@ -91,14 +105,23 @@ def get_read_scope(self) -> RoleScope: peer_ids=sorted(peer_ids), ) - def fill_identity_fields(self, item_dict: Dict[str, Any], role_scope: RoleScope) -> None: + def fill_identity_fields( + self, + item_dict: Dict[str, Any], + role_scope: RoleScope, + memory_type_schema: Optional[MemoryTypeSchema] = None, + ) -> None: del role_scope if self.ctx and self.ctx.user and self.ctx.user.user_id: item_dict["user_id"] = self.ctx.user.user_id item_dict.pop("user_ids", None) - peer_id = item_dict.get("peer_id") - if peer_id: + if memory_type_schema is not None and not memory_type_schema.peer_enabled: + item_dict.pop("peer_id", None) + return + + peer_id = safe_peer_id(item_dict.get("peer_id")) + if peer_id and peer_id != _SELF_PEER_ID: item_dict["peer_id"] = peer_id else: item_dict.pop("peer_id", None) @@ -120,7 +143,7 @@ def render_schema_directories(self, memory_type_schema: MemoryTypeSchema) -> Lis user_spaces: List[str] = [] if self.allow_self: user_spaces.append(user_space) - if self.allow_peer: + if self.allow_peer and getattr(memory_type_schema, "peer_enabled", True): for peer_id in sorted(self.allowed_peer_ids): user_spaces.append(peer_user_space(user_space, peer_id)) @@ -135,26 +158,35 @@ def render_schema_directories(self, memory_type_schema: MemoryTypeSchema) -> Lis ) return directories - def _range_targets(self, ranges: Any) -> tuple[bool, List[str]]: + def _range_targets(self, ranges: Any) -> List[str]: if not ranges or not self._extract_context: - return False, [] + return [] try: msg_range = self._extract_context.read_message_ranges(str(ranges)) except Exception: - logger.warning("Failed to parse memory ranges for peer routing: %s", ranges) - return False, [] + logger.warning("Failed to parse memory ranges for peer memory: %s", ranges) + return [] - include_self = False - peer_ids = set() + target_ids = [] for msg_group in getattr(msg_range, "elements", []) or []: for msg in msg_group: - peer_id = getattr(msg, "peer_id", None) - if peer_id: - if self._can_write_peer(peer_id): - peer_ids.add(peer_id) - elif self.allow_self: - include_self = True - return include_self, sorted(peer_ids) + target_id = self._message_target_id(msg) + if target_id: + target_ids.append(target_id) + return list(dict.fromkeys(target_ids)) + + def _resolve_operation_target_id(self, raw_peer_id: Any) -> Optional[str]: + peer_id = safe_peer_id(raw_peer_id) + if peer_id == _SELF_PEER_ID and self.allow_self: + return _SELF_PEER_ID + if peer_id and self._can_write_peer(peer_id): + return peer_id + fallback_target_id = self._first_target_id_in_messages() + if fallback_target_id: + return fallback_target_id + if self.allow_self: + return _SELF_PEER_ID + return None def calculate_memory_uris( self, @@ -171,57 +203,53 @@ def calculate_memory_uris( user_id = self.ctx.user.user_id operation.memory_fields["user_id"] = user_id - peer_ids_to_write: List[str] = [] - include_self = False - - if operation.memory_fields.get("ranges") is not None: - include_self, peer_ids_to_write = self._range_targets( + target_ids: List[str] = [] + has_ranges = operation.memory_fields.get("ranges") is not None + if not getattr(memory_type_schema, "peer_enabled", True): + operation.memory_fields.pop("peer_id", None) + target_ids = [_SELF_PEER_ID] + elif operation.memory_fields.get("ranges") is not None: + target_ids = self._range_targets( operation.memory_fields.get("ranges"), ) operation.memory_fields.pop("peer_id", None) else: - peer_id = operation.memory_fields.get("peer_id") - if peer_id: - if not self._can_write_peer(peer_id): - return [] - peer_ids_to_write = [peer_id] - operation.memory_fields["peer_id"] = peer_id + target_id = self._resolve_operation_target_id( + operation.memory_fields.get("peer_id"), + ) + if target_id: + target_ids = [target_id] + if target_id == _SELF_PEER_ID: + operation.memory_fields.pop("peer_id", None) + elif target_id: + operation.memory_fields["peer_id"] = target_id else: operation.memory_fields.pop("peer_id", None) - if self.allow_self and self._has_self_user_message(): - include_self = True - else: - fallback_peer_id = self._first_peer_id_in_messages() - if fallback_peer_id: - peer_ids_to_write = [fallback_peer_id] - operation.memory_fields["peer_id"] = fallback_peer_id - elif self.allow_self: - include_self = True - - if not include_self and not peer_ids_to_write: + + if not target_ids: return [] # 文件 uris = set() user_space = user_id - if include_self: - uris.add( - generate_uri( - memory_type=memory_type_schema, - fields=operation.memory_fields, - user_space=user_space, - extract_context=extract_context, - ) - ) - for peer_id in peer_ids_to_write: - target_user_space = peer_user_space(user_space, peer_id) + base_fields = dict(operation.memory_fields) + for target_id in target_ids: + fields = dict(base_fields) + if target_id == _SELF_PEER_ID: + target_user_space = user_space + fields.pop("peer_id", None) + else: + target_user_space = peer_user_space(user_space, target_id) + fields["peer_id"] = target_id uris.add( generate_uri( memory_type=memory_type_schema, - fields=operation.memory_fields, + fields=fields, user_space=target_user_space, extract_context=extract_context, ) ) + if has_ranges: + operation.memory_fields.pop("peer_id", None) return list(uris) diff --git a/openviking/session/memory/memory_type_registry.py b/openviking/session/memory/memory_type_registry.py index a07c20d179..36abf081e2 100644 --- a/openviking/session/memory/memory_type_registry.py +++ b/openviking/session/memory/memory_type_registry.py @@ -227,7 +227,8 @@ def _parse_memory_type(self, data: dict) -> MemoryTypeSchema: directory=data.get("directory", ""), enabled=data.get("enabled", data.get("enable", True)), operation_mode=data.get("operation_mode", "upsert"), - agent_only=data.get("agent_only", False), + stage=data.get("stage", "user"), + peer_enabled=data.get("peer_enabled", True), overview_template=data.get("overview_template"), ) diff --git a/openviking/session/memory/memory_updater.py b/openviking/session/memory/memory_updater.py index 24689273c9..96a91387b3 100644 --- a/openviking/session/memory/memory_updater.py +++ b/openviking/session/memory/memory_updater.py @@ -28,7 +28,11 @@ from openviking.session.memory.memory_type_registry import MemoryTypeRegistry from openviking.session.memory.merge_op import MergeOpFactory from openviking.session.memory.page_id_map import PageIdMap -from openviking.session.memory.utils.memory_file_utils import MemoryFileUtils +from openviking.session.memory.utils.memory_file_utils import ( + MemoryFileUtils, + bump_memory_version, + next_memory_version, +) from openviking.session.memory.utils.resource_refs import ( RESOURCE_REF_SOURCE_SESSION_COMMIT, sync_memory_resource_refs, @@ -38,6 +42,7 @@ from openviking.storage.viking_fs import get_viking_fs from openviking.telemetry import tracer from openviking.telemetry.request_wait_tracker import get_request_wait_tracker +from openviking.telemetry.tracer import get_trace_id from openviking.utils.time_utils import parse_iso_datetime from openviking_cli.exceptions import NotFoundError from openviking_cli.utils import VikingURI, get_logger @@ -71,12 +76,15 @@ async def write_stored_links( ctx: RequestContext, viking_fs: Any, skip_uris: Optional[set] = None, -) -> None: +) -> List[str]: """Write StoredLinks to their endpoint files' links/backlinks fields. For each link: from_uri's ``links`` receives the forward link; to_uri's ``backlinks`` receives the reverse reference. Files listed in skip_uris are skipped (caller handles them in the same write). + Returns the endpoint URIs that were successfully rewritten. Callers can + use this to avoid reporting link-only edits for files that failed the + read/modify/write step. """ from openviking.session.memory.merge_op.link_merge import merge_links @@ -90,6 +98,7 @@ async def write_stored_links( file_links.setdefault(link.to_uri, {"links": [], "backlinks": []}) file_links[link.to_uri]["backlinks"].append(link) + updated_uris: List[str] = [] for uri, link_groups in file_links.items(): try: content = await viking_fs.read_file(uri, ctx=ctx) @@ -102,17 +111,67 @@ async def write_stored_links( mf.backlinks = merge_links( mf.backlinks, [l.model_dump() for l in link_groups["backlinks"]] ) + current_trace_id = get_trace_id() + if current_trace_id: + mf.extra_fields["last_update_trace_id"] = current_trace_id + bump_memory_version(mf) await viking_fs.write_file(uri, MemoryFileUtils.write(mf), ctx=ctx) + updated_uris.append(uri) except Exception as e: tracer.error(f"Failed to apply links to {uri}: {e}") + return updated_uris + + + +def _remap_link_dict(link: Dict[str, Any], uri_remap: Dict[str, str]) -> Dict[str, Any]: + remapped = dict(link or {}) + if remapped.get("from_uri") in uri_remap: + remapped["from_uri"] = uri_remap[remapped["from_uri"]] + if remapped.get("to_uri") in uri_remap: + remapped["to_uri"] = uri_remap[remapped["to_uri"]] + return remapped + + +def remap_stored_links(links: List[StoredLink], uri_remap: Dict[str, str]) -> List[StoredLink]: + if not links or not uri_remap: + return list(links or []) + remapped_links: List[StoredLink] = [] + for link in links: + from_uri = uri_remap.get(link.from_uri, link.from_uri) + to_uri = uri_remap.get(link.to_uri, link.to_uri) + if from_uri == to_uri: + continue + remapped_links.append(link.model_copy(update={"from_uri": from_uri, "to_uri": to_uri})) + return remapped_links + +def _operation_trace_id(op: ResolvedOperation) -> str | None: + source = getattr(op, "source", None) + trace_id = getattr(source, "trace_id", None) if source else None + if trace_id: + return str(trace_id) + fields = dict(getattr(op, "memory_fields", {}) or {}) + field_value = fields.get("last_update_trace_id") or fields.get("trace_id") + if field_value: + return str(field_value) + current_trace_id = get_trace_id() + return current_trace_id or None class ExtractContext: """Extract context for template rendering.""" - def __init__(self, messages: List[Message], chunk_meta: Optional[Dict[int, ChunkMeta]] = None): + def __init__( + self, + messages: List[Message], + chunk_meta: Optional[Dict[int, ChunkMeta]] = None, + *, + split_long_text_messages: bool = True, + ): if chunk_meta is None: - self.messages, self.chunk_meta = self._build_extraction_messages(messages) + if split_long_text_messages: + self.messages, self.chunk_meta = self._build_extraction_messages(messages) + else: + self.messages, self.chunk_meta = list(messages or []), {} else: self.messages = messages self.chunk_meta = chunk_meta @@ -539,7 +598,7 @@ def _can_merge_messages(self, previous: Message, current: Message) -> bool: ) def _format_merged_content(self, messages: List[Message]) -> str: - content = "".join((msg.content or "") for msg in messages) + content = "".join((self._message_content(msg) or "") for msg in messages) if not messages or not self._contains_chunk_message(messages): return content @@ -551,6 +610,15 @@ def _format_merged_content(self, messages: List[Message]) -> str: content = content.rstrip() + "..." return content + def _message_content(self, message: Message) -> str: + texts: List[str] = [] + for part in getattr(message, "parts", []) or []: + if isinstance(part, TextPart): + texts.append(part.text or "") + if texts: + return "".join(texts) + return getattr(message, "content", "") or "" + def _contains_chunk_message(self, messages: List[Message]) -> bool: return any(self._chunk_meta_for(msg) is not None for msg in messages) @@ -619,6 +687,17 @@ def summary(self) -> str: ) +def _same_batch_delete_conflict_key(uri: str) -> str: + """Return a conservative key for detecting same-batch upsert/delete URI conflicts. + + Some local filesystems are case-insensitive. Treat case-only URI variants as + conflicting inside one apply batch so a loser delete cannot remove a winner + upsert before vectorization. + """ + + return str(uri or "").rstrip("/").casefold() + + class MemoryUpdater: """ Applies MemoryOperations to storage. @@ -775,24 +854,38 @@ async def apply_operations( for uri in resolved_op.uris: result.add_error(uri, e) + operations.resolved_links = remap_stored_links( + list(getattr(operations, "resolved_links", []) or []), + dict(getattr(operations, "delete_replacements", {}) or {}), + ) + await self._inherit_deleted_link_relations(operations, result, ctx) + # Apply delete operations (delete_file_contents is List[MemoryFile]) # Skip deletes whose URI was just written in the same batch — this happens when the # LLM issues a Replace with the same experience_name (delete old + create same-name new), # which is semantically an Update. Executing the delete would remove the just-written file. upserted_uris = set(result.written_uris + result.edited_uris) + upserted_uri_keys = {_same_batch_delete_conflict_key(uri) for uri in upserted_uris} for file_content in operations.delete_file_contents: - if file_content.uri in upserted_uris: + delete_uri = file_content.uri + if delete_uri in upserted_uris: tracer.info( - f"[apply_operations] skipping delete for {file_content.uri}: " + f"[apply_operations] skipping delete for {delete_uri}: " "URI was upserted in the same batch (Replace-with-same-name treated as Update)" ) continue + if _same_batch_delete_conflict_key(delete_uri) in upserted_uri_keys: + tracer.info( + f"[apply_operations] skipping delete for {delete_uri}: " + "URI case-conflicts with an upserted URI in the same batch" + ) + continue try: - await self._apply_delete(file_content.uri, ctx) - result.add_deleted(file_content.uri) + await self._apply_delete(delete_uri, ctx) + result.add_deleted(delete_uri) except Exception as e: - tracer.error(f"Failed to delete memory {file_content.uri}", e) - result.add_error(file_content.uri, e) + tracer.error(f"Failed to delete memory {delete_uri}", e) + result.add_error(delete_uri, e) await self._sync_resource_refs_for_result(result, ctx) @@ -891,6 +984,13 @@ async def _apply_upsert( old_content = resolved_op.old_memory_file_content metadata: Dict[str, Any] = dict(resolved_op.memory_fields) + source = getattr(resolved_op, "source", None) + source_extraction_id = getattr(source, "extraction_id", None) if source else None + if source_extraction_id: + metadata["source_extraction_id"] = str(source_extraction_id) + source_trace_id = _operation_trace_id(resolved_op) + if source_trace_id: + metadata["last_update_trace_id"] = source_trace_id # Process fields defined in schema (apply merge_op) for field in schema.fields: if field.name in resolved_op.memory_fields: @@ -928,6 +1028,8 @@ async def _apply_upsert( if key not in schema_field_names and key not in metadata and val is not None: metadata[key] = val + metadata["version"] = next_memory_version(old_content) + # Handle links/backlinks fields: merge with existing incoming_links_by_uri = getattr(resolved_op, "_incoming_links_by_uri", {}) incoming_backlinks_by_uri = getattr(resolved_op, "_incoming_backlinks_by_uri", {}) @@ -1023,6 +1125,87 @@ async def _apply_links_to_existing_files( skip = upserted_uris | (deleted_uris or set()) | non_memory_endpoints await write_stored_links(resolved_links, ctx, viking_fs, skip_uris=skip) + + async def _inherit_deleted_link_relations( + self, + operations: ResolvedOperations, + result: MemoryUpdateResult, + ctx: RequestContext, + ) -> None: + uri_remap = dict(getattr(operations, "delete_replacements", {}) or {}) + if not uri_remap: + return + viking_fs = self._get_viking_fs() + if not viking_fs: + return + + from openviking.session.memory.merge_op.link_merge import merge_links + + inherited_by_uri: Dict[str, Dict[str, List[Dict[str, Any]]]] = {} + for deleted_uri, replacement_uri in uri_remap.items(): + if not deleted_uri or not replacement_uri or deleted_uri == replacement_uri: + continue + try: + content = await viking_fs.read_file(deleted_uri, ctx=ctx) + except Exception as e: + tracer.error(f"Failed to read deleted memory links for replacement {deleted_uri}: {e}") + continue + if not content: + continue + deleted_file = MemoryFileUtils.read(content, uri=deleted_uri) + for link in list(deleted_file.links or []): + remapped = _remap_link_dict(link, uri_remap) + if remapped.get("from_uri") == remapped.get("to_uri"): + continue + target_uri = remapped.get("from_uri") + if target_uri: + inherited_by_uri.setdefault(target_uri, {"links": [], "backlinks": []})[ + "links" + ].append(remapped) + neighbor_uri = remapped.get("to_uri") + if neighbor_uri and neighbor_uri not in uri_remap: + inherited_by_uri.setdefault(neighbor_uri, {"links": [], "backlinks": []})[ + "backlinks" + ].append(remapped) + for link in list(deleted_file.backlinks or []): + remapped = _remap_link_dict(link, uri_remap) + if remapped.get("from_uri") == remapped.get("to_uri"): + continue + target_uri = remapped.get("to_uri") + if target_uri: + inherited_by_uri.setdefault(target_uri, {"links": [], "backlinks": []})[ + "backlinks" + ].append(remapped) + neighbor_uri = remapped.get("from_uri") + if neighbor_uri and neighbor_uri not in uri_remap: + inherited_by_uri.setdefault(neighbor_uri, {"links": [], "backlinks": []})[ + "links" + ].append(remapped) + + written_or_edited = set(result.written_uris + result.edited_uris) + for uri, link_groups in inherited_by_uri.items(): + if uri in uri_remap: + continue + if uri in written_or_edited: + continue + try: + content = await viking_fs.read_file(uri, ctx=ctx) + if not content: + continue + mf = MemoryFileUtils.read(content, uri=uri) + if link_groups["links"]: + mf.links = merge_links(mf.links, link_groups["links"]) + if link_groups["backlinks"]: + mf.backlinks = merge_links(mf.backlinks, link_groups["backlinks"]) + current_trace_id = get_trace_id() + if current_trace_id: + mf.extra_fields["last_update_trace_id"] = current_trace_id + bump_memory_version(mf) + await viking_fs.write_file(uri, MemoryFileUtils.write(mf), ctx=ctx) + result.add_edited(uri) + except Exception as e: + tracer.error(f"Failed to inherit deleted memory links for {uri}: {e}") + async def _apply_delete(self, uri: str, ctx: RequestContext) -> None: """Apply delete operation (uri is already a string).""" viking_fs = self._get_viking_fs() diff --git a/openviking/session/memory/patch_merge_context_provider.py b/openviking/session/memory/patch_merge_context_provider.py new file mode 100644 index 0000000000..1122591548 --- /dev/null +++ b/openviking/session/memory/patch_merge_context_provider.py @@ -0,0 +1,398 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Context provider for merging structured memory-file patches via ExtractLoop.""" + +from __future__ import annotations + +import difflib +from dataclasses import dataclass, field +from typing import Any + +from openviking.server.identity import RequestContext +from openviking.session.memory.dataclass import MemoryFile, MemoryTypeSchema +from openviking.session.memory.session_extract_context_provider import ( + SessionExtractContextProvider, +) +from openviking.session.memory.utils.language import resolve_output_language_from_text + +_SYSTEM_HIDDEN_FIELDS = { + "source_extraction_id", + "source_extraction_ids", + "last_update_trace_id", +} +_MAX_EXTRA_CANDIDATE_FILES = 10 +_PATCH_METADATA_KEYS = ("confidence",) + + +@dataclass(slots=True) +class PatchMergePatch: + """A before/after memory-file patch rendered as field-level line diffs.""" + + before_file: MemoryFile | None + after_file: MemoryFile + metadata: dict[str, Any] = field(default_factory=dict) + + @property + def target_uri(self) -> str | None: + return self.after_file.uri or ( + self.before_file.uri if self.before_file is not None else None + ) + + @property + def memory_type(self) -> str: + return str( + self.after_file.memory_type + or (self.before_file.memory_type if self.before_file is not None else "") + or self.after_file.extra_fields.get("memory_type") + or ( + self.before_file.extra_fields.get("memory_type") + if self.before_file is not None + else "" + ) + ) + + @property + def target_name(self) -> str: + fields = self.after_file.extra_fields or {} + memory_type = self.memory_type + type_specific_key = f"{str(memory_type).rstrip('s')}_name" + name = ( + fields.get(type_specific_key) + or fields.get("experience_name") # backward compat + or fields.get("name") + ) + if name: + return str(name) + uri = self.target_uri + if uri: + # For SKILL.md-style paths, use the directory name. + if uri.endswith("/SKILL.md"): + parts = uri.rstrip("/").split("/") + if len(parts) >= 2: + return parts[-2] + return uri.rstrip("/").split("/")[-1].removesuffix(".md") + return "unknown" + + +def _resolve_patch_output_language(patches: list[PatchMergePatch]) -> str: + return resolve_output_language_from_text(_patch_language_text(patches), fallback_language="en") + + +def _patch_language_text(patches: list[PatchMergePatch]) -> str: + parts: list[str] = [] + for patch in patches: + parts.extend(_memory_file_language_text(patch.after_file)) + return "\n".join(part for part in parts if part) + + +def _memory_file_language_text(file: MemoryFile) -> list[str]: + parts: list[str] = [] + for key, value in (file.extra_fields or {}).items(): + if key in _SYSTEM_HIDDEN_FIELDS or key in {"memory_type", "version"}: + continue + parts.extend(_string_values(value)) + if file.content: + parts.append(file.content) + return parts + + +def _string_values(value: Any) -> list[str]: + if isinstance(value, str): + return [value] + if isinstance(value, list): + return [item for entry in value for item in _string_values(entry)] + if isinstance(value, dict): + return [item for entry in value.values() for item in _string_values(entry)] + return [] + + +class PatchMergeContextProvider(SessionExtractContextProvider): + """Provide original memory files and structured field diffs to ExtractLoop. + + The provider is generic: callers decide which patches to pass in; this class + only exposes original files as prefetched read results and memory-file field + diffs as compact merge context. + """ + + def __init__( + self, + *, + memory_type: str, + patches: list[PatchMergePatch], + required_file_uris: list[str] | None = None, + output_language: str | None = None, + ): + super().__init__(messages=[]) + self.memory_type = memory_type + self.required_file_uris = list(required_file_uris or []) + self.patches = list(patches) + self._output_language = output_language or _resolve_patch_output_language(self.patches) + + def instruction(self) -> str: + output_language = self._output_language + return f"""You are a memory patch merge agent. + +You are given original memory files and structured memory-file field diffs. Merge them by producing final memory operations that follow the provided JSON schema. + +Do not call tools. Output JSON only. + +All memory content must be written in {output_language}. + +Reconcile independent extraction patch proposals: merge duplicate/overlapping +memories into one canonical file patch, and keep distinct memories separate. +Normalize URI/path variants for directory/filename fields. Treat path segment +fields as stable schema identifiers, not free-form labels. Reuse existing +equivalent directories across singular/plural, synonym, or language/script +variants. For new segments, use singular snake_case for English and one concise +canonical term for Chinese; e.g. book not books, 书籍 not 书/图书. If a loser URI +is an existing file, put it in delete_ids; if it is only a new proposal, omit it. +""" + + def get_tools(self) -> list[str]: + return [] + + def get_memory_schemas(self, ctx: RequestContext) -> list[MemoryTypeSchema]: + del ctx + schema = self._get_registry().get(self.memory_type) + if schema is None or not schema.enabled: + raise ValueError(f"Memory schema not found or disabled: {self.memory_type}") + return [schema] + + async def prefetch(self) -> list[dict[str, Any]]: + messages: list[dict[str, Any]] = [] + call_id = 0 + file_uris = await self._resolve_prefetch_file_uris() + for uri in file_uris: + call_id = await self._append_structured_read_result( + messages, + call_id, + uri, + ) + + messages.append( + { + "role": "user", + "content": _render_field_diff_patches(self.patches), + } + ) + return messages + + async def _resolve_prefetch_file_uris(self) -> list[str]: + """Resolve required files plus semantic-search candidates for this merge.""" + + required_uris = _dedupe_uris(self.required_file_uris) + max_extra_candidate_files = min(_MAX_EXTRA_CANDIDATE_FILES, max(5, len(required_uris))) + search_limit = max_extra_candidate_files * 2 + candidate_uris = await self._search_candidate_file_uris(limit=search_limit) + extra_uris: list[str] = [] + required_set = set(required_uris) + for uri in candidate_uris: + if not uri or uri in required_set or uri in extra_uris: + continue + extra_uris.append(uri) + if len(extra_uris) >= max_extra_candidate_files: + break + return [*required_uris, *extra_uris] + + async def _search_candidate_file_uris(self, *, limit: int) -> list[str]: + schema = self._get_registry().get(self.memory_type) + if schema is None or not schema.directory: + return [] + search_dirs = self._render_search_directories(schema) + if not search_dirs: + return [] + query = _build_patch_search_query(self.patches) + if not query: + return [] + return await self.search_files(query=query, search_uris=search_dirs, limit=limit) + + def _render_search_directories(self, schema: MemoryTypeSchema) -> list[str]: + if self._isolation_handler: + return list(dict.fromkeys(self._isolation_handler.render_schema_directories(schema))) + + ctx = self._ctx + user = getattr(ctx, "user", None) + user_id = ( + getattr(ctx, "user_id", None) + or getattr(user, "user_id", None) + or _infer_user_space_from_uris(self.required_file_uris) + or _infer_user_space_from_uris([patch.target_uri for patch in self.patches]) + ) + if not user_id: + return [] + + from openviking.session.memory.utils.uri import render_template + + return [render_template(schema.directory, {"user_space": user_id})] + + +def _render_field_diff_patches(patches: list[PatchMergePatch]) -> str: + if not patches: + return "# Memory File Patches\n\nNo patches provided." + rendered = [ + _render_one_field_diff_patch(index, patch) for index, patch in enumerate(patches, start=1) + ] + return "# Memory File Patches\n\n" + "\n\n".join(rendered).rstrip() + + +def _render_one_field_diff_patch(index: int, patch: PatchMergePatch) -> str: + lines = [f"Patch {index}"] + if patch.metadata: + compact_metadata = _compact_patch_metadata(patch.metadata) + if compact_metadata: + lines.append(f" meta: {_compact_value(compact_metadata)}") + field_diffs = _field_diffs(patch.before_file, patch.after_file) + if not field_diffs: + lines.append(" (no changes)") + return "\n".join(lines) + for field_name, diff in field_diffs: + lines.append(f" {field_name}:") + # Strip unified diff headers (---, +++) but keep @@ hunk markers and content + for diff_line in diff.splitlines(): + if diff_line.startswith("---") or diff_line.startswith("+++"): + continue + lines.append(f" {diff_line}") + return "\n".join(lines) + + +def _field_diffs(before_file: MemoryFile | None, after_file: MemoryFile) -> list[tuple[str, str]]: + before_fields = _memory_file_fields(before_file) if before_file is not None else {} + after_fields = _memory_file_fields(after_file) + diffs: list[tuple[str, str]] = [] + for field_name in sorted(set(before_fields) | set(after_fields)): + before_value = before_fields.get(field_name) + after_value = after_fields.get(field_name) + if before_value == after_value: + continue + diff = _value_unified_diff( + field_name=field_name, + before_value=before_value, + after_value=after_value, + ) + if diff.strip(): + diffs.append((field_name, diff)) + return diffs + + +def _memory_file_fields(file: MemoryFile) -> dict[str, Any]: + fields = dict(file.extra_fields or {}) + for hidden_field in _SYSTEM_HIDDEN_FIELDS: + fields.pop(hidden_field, None) + if file.memory_type is not None: + fields["memory_type"] = file.memory_type + if file.content: + fields["content"] = file.content + if file.links: + fields["links"] = file.links + if file.backlinks: + fields["backlinks"] = file.backlinks + return fields + + +def _value_unified_diff(*, field_name: str, before_value: Any, after_value: Any) -> str: + before_lines = _value_lines(before_value) + after_lines = _value_lines(after_value) + diff_lines = difflib.unified_diff( + before_lines, + after_lines, + fromfile=f"{field_name}.before", + tofile=f"{field_name}.after", + n=1, + lineterm="", + ) + return "\n".join(diff_lines) + + +def _value_lines(value: Any) -> list[str]: + if value is None: + return [] + if isinstance(value, str): + return value.splitlines() + import json + + return json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True).splitlines() + + +def _compact_value(value: Any) -> str: + import json + + return json.dumps(value, ensure_ascii=False, sort_keys=True) + + +def _hide_system_fields(value: Any) -> Any: + if isinstance(value, dict): + return { + key: _hide_system_fields(item) + for key, item in value.items() + if key not in _SYSTEM_HIDDEN_FIELDS + } + if isinstance(value, list): + return [_hide_system_fields(item) for item in value] + return value + + +def _compact_patch_metadata(metadata: dict[str, Any]) -> dict[str, Any]: + """Keep only metadata that helps reconcile patch proposals. + + Full gradient metadata can contain large duplicated fields (links, uris, and + memory_fields). The patch body already renders the target URI and field + changes, while source links are merged outside the LLM response. Keep only + decision signals that help the merge model rank or reconcile proposals. + """ + + cleaned = _hide_system_fields(dict(metadata or {})) + result = { + key: cleaned[key] + for key in _PATCH_METADATA_KEYS + if key in cleaned and _metadata_value_is_useful(cleaned[key]) + } + + return result + + +def _metadata_value_is_useful(value: Any) -> bool: + if value is None: + return False + if value == "": + return False + if isinstance(value, (list, dict, tuple, set)) and not value: + return False + return True + + +def _dedupe_uris(uris: list[str] | None) -> list[str]: + return list(dict.fromkeys(uri for uri in (uris or []) if uri)) + + +def _build_patch_search_query(patches: list[PatchMergePatch]) -> str: + parts: list[str] = [] + for patch in patches: + if patch.target_name: + parts.append(str(patch.target_name)) + if patch.target_uri: + parts.append(str(patch.target_uri).rstrip("/").split("/")[-1].removesuffix(".md")) + content = str(patch.after_file.content or "") + if content: + parts.append(_truncate_query_text(content, 1200)) + return _truncate_query_text("\n\n".join(parts), 5000) + + +def _truncate_query_text(text: Any, max_chars: int) -> str: + normalized = " ".join(str(text or "").split()) + if len(normalized) <= max_chars: + return normalized + return normalized[: max_chars - 3].rstrip() + "..." + + +def _infer_user_space_from_uris(uris: list[str | None]) -> str | None: + for uri in uris: + if not uri: + continue + prefix = "viking://user/" + if not uri.startswith(prefix): + continue + rest = uri.removeprefix(prefix) + user_space = rest.split("/", 1)[0] + if user_space and user_space != "memories": + return user_space + return None diff --git a/openviking/session/memory/schema_model_generator.py b/openviking/session/memory/schema_model_generator.py index b69c53a868..1b2e1cea2e 100644 --- a/openviking/session/memory/schema_model_generator.py +++ b/openviking/session/memory/schema_model_generator.py @@ -13,7 +13,7 @@ from pydantic import BaseModel, Field, WithJsonSchema, create_model from pydantic.config import ConfigDict -from openviking.session.memory.dataclass import FaultTolerantBaseModel, MemoryTypeSchema, WikiLink +from openviking.session.memory.dataclass import DeleteId, FaultTolerantBaseModel, MemoryTypeSchema, WikiLink from openviking.session.memory.memory_isolation_handler import RoleScope from openviking.session.memory.merge_op import MergeOp, MergeOpFactory from openviking.session.memory.merge_op.base import FieldType, get_python_type_for_field @@ -45,6 +45,8 @@ def __init__( schemas: List[MemoryTypeSchema], template_context: Optional[Dict[str, Any]] = None, ): + if hasattr(schemas, "list_all"): + schemas = schemas.list_all() self.schemas = schemas self._template_context = dict(template_context or {}) self._model_cache: Dict[str, Type[BaseModel]] = {} @@ -101,7 +103,7 @@ def create_flat_data_model( # Skip if schema has "ranges" field (like events) - these are message-based and # their self/peer targets are derived from message ranges instead of explicit routing fields. has_ranges = any(field.name == "ranges" for field in memory_type.fields) - if has_peer_scope and not has_ranges: + if has_peer_scope and memory_type.peer_enabled and not has_ranges: peer_values = ", ".join(role_scope.peer_ids) field_definitions["peer_id"] = ( Optional[str], @@ -218,7 +220,7 @@ class MemoryDataWrapper(BaseModel): self._union_model = MemoryDataWrapper return self._union_model - def create_structured_operations_model(self, role_scope: RoleScope) -> Type[BaseModel]: + def create_structured_operations_model(self, role_scope: Optional[RoleScope] = None) -> Type[BaseModel]: """ Create a structured MemoryOperations model with type-safe write operations. @@ -261,14 +263,21 @@ def create_structured_operations_model(self, role_scope: RoleScope) -> Type[Base ), ) - # Only expose delete_uris when at least one schema supports it. + # Only expose delete_ids when at least one schema supports deletion. # add_only schemas (e.g. trajectories) never delete existing records, - # so excluding this field prevents the LLM from hallucinating fake URIs. + # so excluding this field prevents the LLM from hallucinating fake deletes. has_deletable_schema = any(mt.operation_mode != "add_only" for mt in enabled_memory_types) if has_deletable_schema: - field_definitions["delete_uris"] = ( - List[str], - Field(default_factory=list, description="Delete operations as URI strings"), + field_definitions["delete_ids"] = ( + List[DeleteId], + Field( + default_factory=list, + description=( + "Delete operations by page_id. Each item has delete_page_id and " + "replacement_page_id; set replacement_page_id to null for a pure delete, " + "or to the canonical replacement page_id so existing links/backlinks are inherited." + ), + ), ) # Add links field for link extraction (only when enabled globally) @@ -308,7 +317,7 @@ def is_empty(self) -> bool: else: # Single value (not None) return False - return len(self.delete_uris) == 0 + return len(getattr(self, "delete_ids", [])) == 0 def to_legacy_operations(self) -> Dict[str, Any]: """Convert new per-type structure to legacy write_uris/edit_uris format.""" @@ -334,7 +343,7 @@ def to_legacy_operations(self) -> Dict[str, Any]: return { "write_uris": write_uris, "edit_uris": edit_uris, - "delete_uris": self.delete_uris, + "delete_ids": self.delete_ids, } # Attach methods @@ -348,6 +357,10 @@ def to_legacy_operations(self) -> Dict[str, Any]: self._operations_model = StructuredMemoryOperations return self._operations_model + def get_llm_json_schema(self, role_scope: Optional[RoleScope] = None) -> Dict[str, Any]: + """Get the JSON schema for the structured LLM operations model.""" + return self.create_structured_operations_model(role_scope).model_json_schema() + def get_memory_data_json_schema(self) -> Dict[str, Any]: """ Get the JSON schema just for the flat memory data union. @@ -372,6 +385,8 @@ def __init__( schemas: List[MemoryTypeSchema], template_context: Optional[Dict[str, Any]] = None, ): + if hasattr(schemas, "list_all"): + schemas = schemas.list_all() self.schemas = schemas self._template_context = dict(template_context or {}) diff --git a/openviking/session/memory/session_extract_context_provider.py b/openviking/session/memory/session_extract_context_provider.py index d6d35d28b5..0a142f44f7 100644 --- a/openviking/session/memory/session_extract_context_provider.py +++ b/openviking/session/memory/session_extract_context_provider.py @@ -6,7 +6,6 @@ 从会话消息中提取记忆的实现。 """ -import json import os import re from typing import TYPE_CHECKING, Any, Dict, List, Optional @@ -57,6 +56,9 @@ class SessionExtractContextProvider(ExtractContextProvider): """会话提取 Provider - 从会话消息中提取记忆""" + include_tool_parts_in_conversation: bool = False + split_long_text_messages_for_extraction: bool = True + def __init__( self, messages: Any, @@ -108,7 +110,8 @@ def get_extract_context(self) -> "ExtractContext": if self._extract_context is None: self._extract_context = ExtractContext( - self.messages if isinstance(self.messages, list) else [] + self.messages if isinstance(self.messages, list) else [], + split_long_text_messages=self.split_long_text_messages_for_extraction, ) return self._extract_context @@ -293,36 +296,49 @@ def _assemble_conversation(self, messages: Any) -> str: conversation_sections: List[str] = [] def format_message_with_parts(msg: Message) -> str: - """Format message with text parts and ToolCall details.""" + """Format message parts for extraction. + + By default this provider extracts user/session memories, so tool + calls/results are omitted: they are execution evidence rather than + user utterances and can leak environment/database state into user + memories. Agent-scope providers enable tool evidence explicitly. + """ + from openviking.message.part import ToolPart + parts = getattr(msg, "parts", []) formatted_parts: List[str] = [] for part in parts: if hasattr(part, "text") and part.text: formatted_parts.append(part.text) - elif isinstance(part, ToolPart): - tool_info = { - "type": "tool_call", - "tool_name": part.tool_name, - "tool_input": part.tool_input, - "tool_output": part.tool_output[:500] if part.tool_output else "", - "tool_status": part.tool_status, - "duration_ms": part.duration_ms, - } + elif self.include_tool_parts_in_conversation and isinstance(part, ToolPart): + fields = [f"tool_name={part.tool_name}"] + if part.tool_status: + fields.append(f"status={part.tool_status}") + if part.tool_input: + fields.append(f"input={part.tool_input}") + if part.tool_output: + fields.append(f"output={part.tool_output[:500]}") + if part.duration_ms is not None: + fields.append(f"duration_ms={part.duration_ms}") if part.skill_uri: - tool_info["skill_name"] = part.skill_uri.rstrip("/").split("/")[-1] - formatted_parts.append( - f"[ToolCall] {json.dumps(tool_info, ensure_ascii=False)}" - ) - return "\n".join(formatted_parts) if formatted_parts else msg.content + fields.append(f"skill={part.skill_uri.rstrip('/').split('/')[-1]}") + formatted_parts.append("ToolCall: " + "; ".join(fields)) + return "\n".join(formatted_parts) - def format_message_header(msg: Message, idx: int) -> str: + def format_message_header(msg: Message, idx: int) -> str | None: """Format message header with role and stable interaction peer when present.""" + body = format_message_with_parts(msg) + if not body.strip(): + return None speaker = msg.peer_id or msg.role - return f"[{idx}][{msg.role}][{speaker}]: {format_message_with_parts(msg)}" + return f"[{idx}][{msg.role}][{speaker}]: {body}" - conversation_sections.append( - "\n".join([format_message_header(msg, idx) for idx, msg in enumerate(messages)]) - ) + formatted_messages = [ + formatted + for idx, msg in enumerate(messages) + if (formatted := format_message_header(msg, idx)) is not None + ] + conversation_sections.append("\n".join(formatted_messages)) return "\n\n".join(section for section in conversation_sections if section) @@ -332,33 +348,6 @@ def _truncate_prefetch_query_text(self, text: Any, max_chars: int) -> str: return normalized return normalized[: max_chars - 3].rstrip() + "..." - def _format_tool_part_for_search(self, part: ToolPart) -> str: - fields = [] - if part.tool_name: - fields.append(f"tool_name={part.tool_name}") - if part.skill_uri: - skill_name = part.skill_uri.rstrip("/").split("/")[-1] - fields.append(f"skill_name={skill_name}") - if part.tool_status: - fields.append(f"status={part.tool_status}") - if part.tool_input: - fields.append( - "input=" - + self._truncate_prefetch_query_text( - json.dumps(part.tool_input, ensure_ascii=False), - _PREFETCH_SEARCH_TOOL_FIELD_MAX_CHARS, - ) - ) - if part.tool_output and part.tool_status == "error": - fields.append( - "error=" - + self._truncate_prefetch_query_text( - part.tool_output, - _PREFETCH_SEARCH_TOOL_FIELD_MAX_CHARS, - ) - ) - return "ToolCall: " + "; ".join(fields) - def _build_prefetch_search_query(self) -> str: """Build a compact semantic query from raw conversation messages. @@ -378,7 +367,6 @@ def _build_prefetch_search_query(self) -> str: parts = getattr(msg, "parts", []) text_parts: List[str] = [] - tool_parts: List[str] = [] for part in parts: if hasattr(part, "text") and part.text: @@ -388,11 +376,6 @@ def _build_prefetch_search_query(self) -> str: else _PREFETCH_SEARCH_ASSISTANT_TEXT_PART_MAX_CHARS ) text_parts.append(self._truncate_prefetch_query_text(part.text, limit)) - elif isinstance(part, ToolPart): - tool_part = self._format_tool_part_for_search(part) - if tool_part != "ToolCall: ": - tool_parts.append(tool_part) - if text_parts: section = f"{speaker}: " + "\n".join(text_parts) if role == "user": @@ -400,9 +383,6 @@ def _build_prefetch_search_query(self) -> str: else: supporting_sections.append(section) - if tool_parts: - supporting_sections.append(f"{speaker}: " + "\n".join(tool_parts)) - query = "\n\n".join(primary_sections + supporting_sections) if not query.strip(): query = self._assemble_conversation(self.messages) @@ -421,6 +401,13 @@ def create_tool_context(self, default_search_uris=[]): ) return tool_ctx + @staticmethod + def _is_expected_read_not_found(error: Any) -> bool: + if not error: + return False + error_text = str(error) + return error_text == "not_found" or error_text.startswith("File not found") + async def read_file(self, uri: str) -> Optional[Dict]: """Read a file via MemoryReadTool (auto-registers page_id, fills read_file_contents).""" read_tool = get_tool("read") @@ -429,11 +416,13 @@ async def read_file(self, uri: str) -> Optional[Dict]: try: result = await read_tool.execute(self.create_tool_context(), uri=uri) if isinstance(result, dict) and "error" in result: - tracer.info(f"Failed to read {uri}: {result['error']}") + if not self._is_expected_read_not_found(result["error"]): + tracer.info(f"Failed to read {uri}: {result['error']}") return None return result except Exception as e: - tracer.error(f"Failed to read {uri}: {e}") + if not self._is_expected_read_not_found(e): + tracer.error(f"Failed to read {uri}: {e}") return None async def search_files( @@ -494,11 +483,11 @@ async def prefetch(self) -> List[Dict]: pre_fetch_messages = [] pre_fetch_messages.append(self._build_conversation_message()) - # 触发 registry 加载,过滤掉 agent_only 的 schema(trajectory/experience 由执行提取处理) + # 触发 registry 加载,过滤掉 agent stage 的 schema(trajectory/experience 由执行提取处理) schemas = [ s for s in self._get_registry().list_all(include_disabled=False) - if not getattr(s, "agent_only", False) + if getattr(s, "stage", "user") == "user" ] if self._isolation_handler: schemas = [s for s in schemas if self._isolation_handler.allows_schema(s)] @@ -595,7 +584,6 @@ async def prefetch(self) -> List[Dict]: return pre_fetch_messages - @tracer("execute_tool", ignore_result=False) async def execute_tool( self, tool_call, @@ -603,8 +591,15 @@ async def execute_tool( tool = get_tool(tool_call.name) if not tool: return {"error": f"Unknown tool: {tool_call.name}"} - tracer.info(f"tool_call.arguments={tool_call.arguments}") result = await tool.execute(self.create_tool_context(), **tool_call.arguments) + is_expected_read_not_found = ( + tool_call.name == "read" + and isinstance(result, dict) + and result.get("error") + and self._is_expected_read_not_found(result["error"]) + ) + if not is_expected_read_not_found: + tracer.info(f"tool_call.arguments={tool_call.arguments}") return result def get_tools(self) -> List[str]: @@ -619,7 +614,7 @@ def get_memory_schemas(self, ctx: RequestContext) -> List[Any]: schemas = [ s for s in self._get_registry().list_all(include_disabled=False) - if not getattr(s, "agent_only", False) + if getattr(s, "stage", "user") == "user" ] if self._isolation_handler: schemas = [s for s in schemas if self._isolation_handler.allows_schema(s)] diff --git a/openviking/session/memory/streaming_memory_updater.py b/openviking/session/memory/streaming_memory_updater.py new file mode 100644 index 0000000000..f554aab863 --- /dev/null +++ b/openviking/session/memory/streaming_memory_updater.py @@ -0,0 +1,1851 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Streaming updater for ordinary user memories. + +This module provides a realtime batching layer for session user-memory writes. +Multiple concurrent commits can submit resolved memory operations; the updater +buffers them for a small count/time window, merges patches with the generic +PatchMergeContextProvider, then applies the merged operations with MemoryUpdater. +""" + +from __future__ import annotations + +import asyncio +import re +import threading +from dataclasses import dataclass, field +from typing import Any, Hashable + +from openviking.core.peer_id import safe_peer_id +from openviking.message import Message +from openviking.server.identity import RequestContext +from openviking.session.memory.dataclass import ( + MemoryFile, + MemoryOperationSource, + MemoryTypeSchema, + ResolvedOperation, + ResolvedOperations, + StoredLink, +) +from openviking.session.memory.extract_loop import ExtractLoop +from openviking.session.memory.memory_isolation_handler import MemoryIsolationHandler +from openviking.session.memory.memory_type_registry import ( + MemoryTypeRegistry, + create_default_registry, +) +from openviking.session.memory.memory_updater import ( + ExtractContext, + MemoryUpdater, + MemoryUpdateResult, + remap_stored_links, + write_stored_links, +) +from openviking.session.memory.merge_op import MergeOpFactory +from openviking.session.memory.patch_merge_context_provider import ( + PatchMergeContextProvider, + PatchMergePatch, +) +from openviking.session.memory.session_extract_context_provider import SessionExtractContextProvider +from openviking.session.memory.utils.memory_file_utils import MemoryFileUtils, next_memory_version +from openviking.session.memory.utils.streaming_batcher import ( + StreamingBatcher, + StreamingBatcherConfig, +) +from openviking.storage.viking_fs import get_viking_fs +from openviking.telemetry import tracer +from openviking.telemetry.tracer import get_trace_id +from openviking_cli.utils import get_logger +from openviking_cli.utils.config import get_openviking_config + +logger = get_logger(__name__) + + +@dataclass(slots=True) +class StreamingMemoryUpdaterConfig: + """Configuration for automatic streaming ordinary-memory updates.""" + + max_operations_per_update: int = 8 + max_wait_seconds: float = 10.0 + timer_check_interval_seconds: float = 1.0 + trace_console: bool = False + + def __post_init__(self) -> None: + if self.max_operations_per_update <= 0: + raise ValueError("max_operations_per_update must be > 0") + if self.max_wait_seconds <= 0: + raise ValueError("max_wait_seconds must be > 0") + if self.timer_check_interval_seconds <= 0: + raise ValueError("timer_check_interval_seconds must be > 0") + + +@dataclass(frozen=True, slots=True) +class StreamingMemoryUpdaterKey: + """Process-local registry key for one shared user-memory updater.""" + + account_id: str + user_id: str + + +@dataclass(frozen=True, slots=True) +class MemoryMergeGroupKey: + """Per-scope/type batching key for second-stage memory merges.""" + + peer_id: str | None + memory_type: str + + +@dataclass(slots=True) +class MemoryUpdateRequest: + """One commit's resolved user-memory update request.""" + + operations: ResolvedOperations + messages: list[Message] + ctx: RequestContext + strict_extract_errors: bool = False + isolation_options: dict[str, Any] = field(default_factory=dict) + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class StreamingMemoryUpdateResult: + """Result returned when a submit triggers a flush.""" + + operations: ResolvedOperations + apply_result: MemoryUpdateResult + request_count: int + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class StreamingMemoryUpdater: + """Long-lived ordinary-memory updater with count/time window batching.""" + + registry: MemoryTypeRegistry | None = None + vikingdb: Any = None + config: StreamingMemoryUpdaterConfig = field(default_factory=StreamingMemoryUpdaterConfig) + _group_batchers: dict[ + MemoryMergeGroupKey, + StreamingBatcher[MemoryUpdateRequest, StreamingMemoryUpdateResult], + ] = field(init=False, repr=False) + _group_batchers_lock: asyncio.Lock = field(init=False, repr=False) + _apply_lock: asyncio.Lock = field(init=False, repr=False) + _last_result: StreamingMemoryUpdateResult | None = field(init=False, default=None, repr=False) + _closed: bool = field(init=False, default=False, repr=False) + + def __post_init__(self) -> None: + self.registry = self.registry or create_default_registry() + self._group_batchers = {} + self._group_batchers_lock = asyncio.Lock() + self._apply_lock = asyncio.Lock() + self._last_result = None + self._closed = False + + @property + def closed(self) -> bool: + return self._closed + + @property + def last_result(self) -> StreamingMemoryUpdateResult | None: + return self._last_result + + async def get_buffered_operation_count(self) -> int: + async with self._group_batchers_lock: + batchers = list(self._group_batchers.values()) + sizes = await asyncio.gather(*(batcher.get_buffered_size() for batcher in batchers)) + return sum(sizes) + + async def close(self) -> StreamingMemoryUpdateResult | None: + if self._closed: + return None + self._closed = True + async with self._group_batchers_lock: + batchers = list(self._group_batchers.values()) + self._group_batchers = {} + results = await asyncio.gather(*(batcher.close() for batcher in batchers)) + return combine_streaming_memory_results(*results) + + @tracer("memory.streaming_updater.submit", ignore_result=True, ignore_args=True) + async def submit(self, request: MemoryUpdateRequest) -> StreamingMemoryUpdateResult: + """Submit one resolved update request. + + The request is buffered and flushed by the shared count/time window. + ``submit`` waits until the batch containing this request is merged and + applied, preserving session.commit's "write is visible on return" + semantics while still allowing concurrent commits to batch together. + """ + + if self._closed: + raise RuntimeError("StreamingMemoryUpdater is closed") + if request.ctx is None: + raise ValueError("MemoryUpdateRequest.ctx is required") + attach_source_to_request_operations(request) + append_only_request, merge_request = self._split_append_only_request(request) + append_result = ( + await self._apply_append_only_request_now(append_only_request) + if append_only_request is not None + else None + ) + merge_result = ( + await self._submit_grouped_merge_request(merge_request) + if merge_request is not None + else None + ) + result = combine_streaming_memory_results( + append_result, + merge_result, + fallback_request_count=1, + ) + self._last_result = result + scoped_result = scope_memory_update_result_to_submitter(result, request) + tracer.info( + "StreamingMemoryUpdater submit finished " + f"batch_id={result.metadata.get('batch_id')} " + f"batch_trace_id={result.metadata.get('batch_trace_id')} " + f"flush_reason={result.metadata.get('flush_reason')} " + f"request_count={result.request_count} " + f"operation_count={result.metadata.get('operation_count')} " + f"written_uris={scoped_result.apply_result.written_uris} " + f"edited_uris={scoped_result.apply_result.edited_uris} " + f"deleted_uris={scoped_result.apply_result.deleted_uris} " + f"errors={scoped_result.apply_result.errors}", + console=self.config.trace_console, + ) + return scoped_result + + async def _submit_grouped_merge_request( + self, + request: MemoryUpdateRequest, + ) -> StreamingMemoryUpdateResult | None: + grouped_requests = split_request_by_merge_group(request) + if not grouped_requests: + return None + submissions = [ + (await self._get_group_batcher(group_key)).submit(group_request) + for group_key, group_request in grouped_requests + ] + group_results = list(await asyncio.gather(*submissions)) + result = combine_streaming_memory_results(*group_results, fallback_request_count=1) + await self._apply_post_group_links(request, result) + return result + + async def _apply_post_group_links( + self, + request: MemoryUpdateRequest, + result: StreamingMemoryUpdateResult, + ) -> None: + links = merge_link_lists(list(getattr(request.operations, "resolved_links", []) or [])) + if not links: + return + links = remap_stored_links( + links, dict(getattr(result.operations, "delete_replacements", {}) or {}) + ) + valid_links = await filter_valid_links( + links, + upsert_operations=result.operations.upsert_operations, + delete_file_contents=result.operations.delete_file_contents, + ctx=request.ctx, + trace_console=self.config.trace_console, + ) + if not valid_links: + return + viking_fs = safe_get_viking_fs() + if viking_fs is not None: + updated_uris = await write_stored_links(valid_links, request.ctx, viking_fs) + for uri in dict.fromkeys(updated_uris): + result.apply_result.add_edited(uri) + result.operations.resolved_links = merge_link_lists( + list(getattr(result.operations, "resolved_links", []) or []), + valid_links, + ) + + async def _get_group_batcher( + self, + group_key: MemoryMergeGroupKey, + ) -> StreamingBatcher[MemoryUpdateRequest, StreamingMemoryUpdateResult]: + async with self._group_batchers_lock: + batcher = self._group_batchers.get(group_key) + if batcher is not None: + return batcher + + batcher = self._create_group_batcher(group_key) + self._group_batchers[group_key] = batcher + return batcher + + def _create_group_batcher( + self, + group_key: MemoryMergeGroupKey, + ) -> StreamingBatcher[MemoryUpdateRequest, StreamingMemoryUpdateResult]: + async def process_batch( + requests: list[MemoryUpdateRequest], + reason: str, + ) -> StreamingMemoryUpdateResult: + return await self._process_batch(group_key, requests, reason) + + batcher = StreamingBatcher( + name=( + "openviking-streaming-memory-updater:" + f"{group_key.peer_id or 'self'}:{group_key.memory_type}" + ), + process_batch=process_batch, + config=StreamingBatcherConfig( + max_items_per_batch=self.config.max_operations_per_update, + max_wait_seconds=self.config.max_wait_seconds, + timer_check_interval_seconds=self.config.timer_check_interval_seconds, + ), + item_size=lambda request: _operation_count(request.operations), + result_metadata=lambda result: result.metadata, + ) + return batcher + + def _split_append_only_request( + self, request: MemoryUpdateRequest + ) -> tuple[MemoryUpdateRequest | None, MemoryUpdateRequest | None]: + operations = request.operations + registry = self.registry or create_default_registry() + append_ops: list[ResolvedOperation] = [] + merge_ops: list[ResolvedOperation] = [] + for op in list(operations.upsert_operations or []): + schema = registry.get(op.memory_type) + if op.uris and getattr(schema, "operation_mode", None) == "add_only": + append_ops.append(op) + else: + merge_ops.append(op) + + append_links, merge_links = split_links_for_append_only_ops( + list(getattr(operations, "resolved_links", []) or []), + append_ops=append_ops, + merge_ops=merge_ops, + ) + append_request = None + if append_ops: + append_request = clone_memory_update_request( + request, + operations=ResolvedOperations( + upsert_operations=append_ops, + delete_file_contents=[], + errors=[], + resolved_links=append_links, + ), + ) + + merge_request = None + if merge_ops or operations.delete_file_contents or operations.errors: + merge_request = clone_memory_update_request( + request, + operations=ResolvedOperations( + upsert_operations=merge_ops, + delete_file_contents=list(operations.delete_file_contents or []), + errors=list(operations.errors or []), + resolved_links=merge_links, + delete_replacements=dict(getattr(operations, "delete_replacements", {}) or {}), + ), + ) + return append_request, merge_request + + async def _apply_append_only_request_now( + self, + request: MemoryUpdateRequest, + ) -> StreamingMemoryUpdateResult: + tracer.info( + "StreamingMemoryUpdater fast path started " + f"reason=append_only operation_count={_operation_count(request.operations)}", + console=self.config.trace_console, + ) + operations = request.operations.model_copy(deep=True) + operations.resolved_links = await filter_valid_links( + merge_link_lists(list(getattr(operations, "resolved_links", []) or [])), + upsert_operations=operations.upsert_operations, + delete_file_contents=operations.delete_file_contents, + ctx=request.ctx, + trace_console=self.config.trace_console, + ) + apply_result = await self._apply_operations( + operations=operations, + request=request, + messages=request.messages, + ) + result = StreamingMemoryUpdateResult( + operations=operations, + apply_result=apply_result, + request_count=1, + metadata={ + "flush_reason": "append_only_fast_path", + "operation_count": _operation_count(operations), + "fast_path": True, + "append_only_operation_count": _operation_count(operations), + }, + ) + tracer.info( + "StreamingMemoryUpdater fast path finished " + f"written_uris={apply_result.written_uris} " + f"edited_uris={apply_result.edited_uris} " + f"deleted_uris={apply_result.deleted_uris} " + f"errors={apply_result.errors}", + console=self.config.trace_console, + ) + return result + + async def _process_batch( + self, + group_key: MemoryMergeGroupKey, + requests: list[MemoryUpdateRequest], + reason: str, + ) -> StreamingMemoryUpdateResult: + input_operations = sum(_operation_count(request.operations) for request in requests) + input_patches = sum( + len(getattr(request.operations, "upsert_operations", []) or []) for request in requests + ) + input_deletes = sum( + len(getattr(request.operations, "delete_file_contents", []) or []) + for request in requests + ) + tracer.info( + "StreamingMemoryUpdater flush started " + f"group={group_key} reason={reason} request_count={len(requests)} " + f"input_operations={input_operations} " + f"input_patches={input_patches} " + f"input_deletes={input_deletes}", + console=self.config.trace_console, + ) + merged_operations = await self._merge_requests(requests) + first_request = requests[0] + apply_result = await self._apply_operations( + operations=merged_operations, + request=first_request, + messages=_combined_request_messages(requests), + ) + result = StreamingMemoryUpdateResult( + operations=merged_operations, + apply_result=apply_result, + request_count=len(requests), + metadata={ + "flush_reason": reason, + "operation_count": _operation_count(merged_operations), + "merge_group": _merge_group_key_label(group_key), + }, + ) + self._last_result = result + tracer.info( + "StreamingMemoryUpdater flush finished " + f"group={group_key} reason={reason} request_count={len(requests)} " + f"written_uris={apply_result.written_uris} " + f"edited_uris={apply_result.edited_uris} " + f"deleted_uris={apply_result.deleted_uris} " + f"errors={apply_result.errors}", + console=self.config.trace_console, + ) + return result + + async def _apply_operations( + self, + *, + operations: ResolvedOperations, + request: MemoryUpdateRequest, + messages: list[Message], + ) -> MemoryUpdateResult: + updater = MemoryUpdater( + registry=self.registry, + vikingdb=self.vikingdb, + transaction_handle=None, + ) + extract_context = ExtractContext(messages) + isolation_handler = _make_isolation_handler(request, extract_context) + async with self._apply_lock: + return await updater.apply_operations( + operations, + request.ctx, + extract_context=extract_context, + isolation_handler=isolation_handler, + ) + + async def _merge_requests(self, requests: list[MemoryUpdateRequest]) -> ResolvedOperations: + all_ops = ResolvedOperations( + upsert_operations=[], + delete_file_contents=[], + errors=[], + resolved_links=[], + delete_replacements={}, + ) + for request in requests: + ops = request.operations + all_ops.upsert_operations.extend(list(ops.upsert_operations or [])) + all_ops.delete_file_contents.extend(list(ops.delete_file_contents or [])) + all_ops.errors.extend(list(ops.errors or [])) + all_ops.resolved_links.extend(list(getattr(ops, "resolved_links", []) or [])) + all_ops.delete_replacements.update(dict(getattr(ops, "delete_replacements", {}) or {})) + return await merge_memory_operations( + operations=all_ops, + messages=_combined_request_messages(requests), + ctx=requests[0].ctx, + registry=self.registry or create_default_registry(), + strict_extract_errors=any(request.strict_extract_errors for request in requests), + trace_console=self.config.trace_console, + ) + + +def split_request_by_merge_group( + request: MemoryUpdateRequest, +) -> list[tuple[MemoryMergeGroupKey, MemoryUpdateRequest]]: + """Split one commit request into per-(peer_id, memory_type) merge requests. + + A submit/session.commit awaits all returned group requests, so commits touching + multiple memory types still return only after every affected group is merged + and applied. + """ + operations = request.operations + upsert_groups: dict[MemoryMergeGroupKey, list[ResolvedOperation]] = {} + delete_groups: dict[MemoryMergeGroupKey, list[MemoryFile]] = {} + passthrough_upserts: list[ResolvedOperation] = [] + + for op in list(operations.upsert_operations or []): + if not op.uris: + passthrough_upserts.append(op) + continue + peer_id = _peer_id_for_operation(op) + for uri in op.uris: + single_uri_op = clone_operation_for_uri(op, uri) + group_key = MemoryMergeGroupKey(peer_id=peer_id, memory_type=single_uri_op.memory_type) + upsert_groups.setdefault(group_key, []).append(single_uri_op) + + for file in list(operations.delete_file_contents or []): + group_key = MemoryMergeGroupKey( + peer_id=_peer_id_for_memory_file(file), + memory_type=file.memory_type or "", + ) + delete_groups.setdefault(group_key, []).append(file) + + group_keys = list(dict.fromkeys(list(upsert_groups.keys()) + list(delete_groups.keys()))) + grouped_requests: list[tuple[MemoryMergeGroupKey, MemoryUpdateRequest]] = [] + for group_key in group_keys: + group_upserts = upsert_groups.get(group_key, []) + group_deletes = delete_groups.get(group_key, []) + grouped_requests.append( + ( + group_key, + clone_memory_update_request( + request, + operations=ResolvedOperations( + upsert_operations=group_upserts, + delete_file_contents=group_deletes, + errors=list(operations.errors or []), + resolved_links=[], + delete_replacements={ + file.uri: replacement_uri + for file in group_deletes + if file.uri + if ( + replacement_uri := ( + getattr(operations, "delete_replacements", {}) or {} + ).get(file.uri) + ) + }, + ), + ), + ) + ) + + if passthrough_upserts: + group_key = MemoryMergeGroupKey(peer_id=None, memory_type="") + grouped_requests.append( + ( + group_key, + clone_memory_update_request( + request, + operations=ResolvedOperations( + upsert_operations=passthrough_upserts, + delete_file_contents=[], + errors=list(operations.errors or []), + resolved_links=[], + delete_replacements={}, + ), + ), + ) + ) + return grouped_requests + + +def _merge_group_key_label(group_key: MemoryMergeGroupKey) -> str: + peer_label = group_key.peer_id or "self" + memory_type = group_key.memory_type or "unknown" + return f"peer={peer_label},memory_type={memory_type}" + + +async def merge_memory_operations( + *, + operations: ResolvedOperations, + messages: list[Message], + ctx: RequestContext, + registry: MemoryTypeRegistry | None = None, + strict_extract_errors: bool = False, + trace_console: bool = False, +) -> ResolvedOperations: + """Merge resolved memory operations by memory type/URI using patch context.""" + + if operations.has_errors(): + tracer.info( + "[streaming_memory_updater] merge skipped reason=operation_errors " + f"error_count={len(operations.errors)} " + f"patch_count={len(operations.upsert_operations or [])} " + f"delete_count={len(operations.delete_file_contents or [])}", + console=trace_console, + ) + return operations + + # Group by (peer_id, memory_type) — peer_id is None for self memories. + # Upserts get peer_id from memory_fields; deletes get it from extra_fields. + # Types with ranges (e.g. events) pop peer_id from memory_fields, but those are + # add_only and skip merge entirely, so they never reach this grouping. + upsert_groups: dict[tuple[str | None, str], list[ResolvedOperation]] = {} + delete_groups: dict[tuple[str | None, str], list[MemoryFile]] = {} + passthrough_upserts: list[ResolvedOperation] = [] + for op in operations.upsert_operations: + if not op.uris: + passthrough_upserts.append(op) + continue + peer_id = _peer_id_for_operation(op) + for uri in op.uris: + single_uri_op = clone_operation_for_uri(op, uri) + upsert_groups.setdefault((peer_id, single_uri_op.memory_type), []).append(single_uri_op) + for df in operations.delete_file_contents: + peer_id = _peer_id_for_memory_file(df) + memory_type = df.memory_type or "" + delete_groups.setdefault((peer_id, memory_type), []).append(df) + + # Union all group keys from both upserts and deletes + all_group_keys = list(dict.fromkeys(list(upsert_groups.keys()) + list(delete_groups.keys()))) + + tracer.info( + "[streaming_memory_updater] merge batch " + f"patch_count={len(operations.upsert_operations or [])} " + f"delete_count={len(operations.delete_file_contents or [])} " + f"passthrough_upserts={len(passthrough_upserts)} " + f"group_count={len(all_group_keys)} " + f"groups={sorted(str(k) for k in all_group_keys)}", + console=trace_console, + ) + + merged_upserts = list(passthrough_upserts) + merged_deletes: list[MemoryFile] = [] + merged_delete_replacements: dict[str, str] = {} + merged_links = merge_link_lists(list(getattr(operations, "resolved_links", []) or [])) + registry = registry or create_default_registry() + merge_results = await asyncio.gather( + *[ + merge_one_memory_type_operations( + memory_type=memory_type, + operations=upsert_groups.get((peer_id, memory_type), []), + delete_files=delete_groups.get((peer_id, memory_type), []), + messages=messages, + ctx=ctx, + registry=registry, + peer_id=peer_id, + trace_console=trace_console, + ) + for (peer_id, memory_type) in all_group_keys + ], + return_exceptions=True, + ) + + for (peer_id, memory_type), group_key, merge_result in zip( + all_group_keys, all_group_keys, merge_results, strict=True + ): + ops_list = upsert_groups.get(group_key, []) + if not isinstance(merge_result, Exception): + merged = merge_result + enforce_merge_group_peer_id( + merged.upsert_operations, + peer_id=peer_id, + memory_type=memory_type, + registry=registry, + ctx=ctx, + ) + _inherit_source_metadata_to_merged_operations(ops_list, merged.upsert_operations) + merged_upserts.extend(merged.upsert_operations) + merged_deletes.extend(merged.delete_file_contents) + merged_delete_replacements.update( + dict(getattr(merged, "delete_replacements", {}) or {}) + ) + merged_links = merge_link_lists( + merged_links, + list(getattr(merged, "resolved_links", []) or []), + ) + continue + + peer_label = f"peer={peer_id}" if peer_id else "peer=self" + tracer.info( + "[streaming_memory_updater] merge fallback " + f"memory_type={memory_type} {peer_label} mode=fallback_original " + f"reason=llm_merge_failed patch_count={len(ops_list)} " + f"target_count={len(_unique_operation_uris(ops_list))} error={merge_result}", + console=trace_console, + ) + logger.warning( + "[streaming_memory_updater] merge failed for %s (%s): %s", + memory_type, + peer_label, + merge_result, + ) + if strict_extract_errors or is_cross_extraction_group(ops_list): + raise merge_result + # Fallback: keep original operations and delete files for this group + merged_upserts.extend(ops_list) + fallback_deletes = delete_groups.get(group_key, []) + merged_deletes.extend(fallback_deletes) + for delete_file in fallback_deletes: + replacement_uri = dict(getattr(operations, "delete_replacements", {}) or {}).get( + delete_file.uri + ) + if replacement_uri: + merged_delete_replacements[delete_file.uri] = replacement_uri + + merged_links = await filter_valid_links( + merged_links, + upsert_operations=merged_upserts, + delete_file_contents=merged_deletes, + ctx=ctx, + trace_console=trace_console, + ) + return ResolvedOperations( + upsert_operations=merged_upserts, + delete_file_contents=merged_deletes, + errors=list(operations.errors), + resolved_links=merged_links, + delete_replacements=merged_delete_replacements, + ) + + +async def merge_one_memory_type_operations( + *, + memory_type: str, + operations: list[ResolvedOperation], + delete_files: list[MemoryFile] | None = None, + messages: list[Message], + ctx: RequestContext, + registry: MemoryTypeRegistry | None = None, + peer_id: str | None = None, + trace_console: bool = False, +) -> ResolvedOperations: + registry = registry or create_default_registry() + schema = registry.get(memory_type) + delete_files = list(delete_files or []) + patch_count = len(operations) + target_uris = _unique_operation_uris(operations) + target_count = len(target_uris) + existing_file_count = sum( + 1 for op in operations if getattr(op, "old_memory_file_content", None) is not None + ) + delete_count = len(delete_files) + duplicate_target_count = patch_count - target_count + operation_mode = ( + getattr(schema, "operation_mode", "unknown") if schema is not None else "unknown" + ) + + # Fast path: no upserts, only deletes — passthrough directly + if not operations and delete_files: + tracer.info( + "[streaming_memory_updater] memory_type merge decision " + f"memory_type={memory_type} mode=no_merge " + f"reason=delete_only delete_count={delete_count}", + console=trace_console, + ) + return ResolvedOperations( + upsert_operations=[], + delete_file_contents=list(delete_files), + errors=[], + resolved_links=[], + delete_replacements={}, + ) + if operation_mode == "add_only": + tracer.info( + "[streaming_memory_updater] memory_type merge decision " + f"memory_type={memory_type} mode=no_merge " + f"reason=add_only operation_mode={operation_mode} " + f"patch_count={patch_count} target_count={target_count} " + f"duplicate_target_count={duplicate_target_count} " + f"existing_file_count={existing_file_count}", + console=trace_console, + ) + return ResolvedOperations( + upsert_operations=list(operations), + delete_file_contents=[], + errors=[], + resolved_links=[], + delete_replacements={}, + ) + + fast_path, fast_path_reason = classify_memory_merge_mode(operations, schema=schema) + if fast_path: + tracer.info( + "[streaming_memory_updater] memory_type merge decision " + f"memory_type={memory_type} mode=no_merge " + f"reason={fast_path_reason} operation_mode={operation_mode} " + f"patch_count={patch_count} target_count={target_count} " + f"duplicate_target_count={duplicate_target_count} " + f"existing_file_count={existing_file_count}", + console=trace_console, + ) + return ResolvedOperations( + upsert_operations=list(operations), + delete_file_contents=[], + errors=[], + resolved_links=[], + delete_replacements={}, + ) + + tracer.info( + "[streaming_memory_updater] memory_type merge decision " + f"memory_type={memory_type} mode=llm_merge " + f"reason={fast_path_reason} operation_mode={operation_mode} " + f"patch_count={patch_count} delete_count={delete_count} " + f"target_count={target_count} " + f"duplicate_target_count={duplicate_target_count} " + f"existing_file_count={existing_file_count}", + console=trace_console, + ) + + if schema is None: + raise ValueError(f"Memory schema not found: {memory_type}") + + extract_context = ExtractContext(messages) + # Existing files: both upsert old_content and delete files count as "existing" + required_file_uris = list( + dict.fromkeys( + [ + uri + for op in operations + for uri in op.uris + if getattr(op, "old_memory_file_content", None) is not None + ] + + [df.uri for df in delete_files if df.uri] + ) + ) + patches = [ + operation_to_patch(op, schema=schema, extract_context=extract_context) for op in operations + ] + [ + memory_file_to_delete_patch(df, schema=schema, extract_context=extract_context) + for df in delete_files + ] + provider = PatchMergeContextProvider( + memory_type=memory_type, + required_file_uris=required_file_uris, + patches=patches, + output_language=merge_output_language_from_messages(messages), + ) + provider._ctx = ctx + provider._viking_fs = safe_get_viking_fs() + provider._extract_context = extract_context + # Build isolation handler matching this group's peer scope. + # peer_id=None → self scope; peer_id set → peer-only scope. + if peer_id: + isolation_handler = MemoryIsolationHandler( + ctx, + extract_context, + allowed_memory_types={memory_type}, + allow_self=False, + allowed_peer_ids={peer_id}, + ) + else: + isolation_handler = MemoryIsolationHandler( + ctx, + extract_context, + allowed_memory_types={memory_type}, + allow_self=True, + ) + isolation_handler.prepare_messages() + provider._isolation_handler = isolation_handler + seed_patch_merge_read_contents(provider, operations) + # Also seed delete files into read_contents so LLM can see their content + for df in delete_files: + if df.uri: + provider.read_file_contents[df.uri] = df + prefetch_messages = await provider.prefetch() + + async def _prefetch(): + return list(prefetch_messages) + + provider.prefetch = _prefetch + vlm = get_openviking_config().vlm.get_vlm_instance() + tracer.info( + "[streaming_memory_updater] llm merge input " + f"memory_type={memory_type} required_file_count={len(required_file_uris)} " + f"required_files={required_file_uris} patch_count={len(patches)} " + f"target_count={target_count}", + console=trace_console, + ) + orchestrator = ExtractLoop( + vlm=vlm, + viking_fs=safe_get_viking_fs(), + ctx=ctx, + context_provider=provider, + isolation_handler=isolation_handler, + max_iterations=1, + ) + merged, _ = await orchestrator.run() + merged = merged or ResolvedOperations(upsert_operations=[], delete_file_contents=[], errors=[]) + tracer.info( + "[streaming_memory_updater] llm merge output " + f"memory_type={memory_type} upserts={len(merged.upsert_operations)} " + f"deletes={len(merged.delete_file_contents)} errors={len(merged.errors)}", + console=trace_console, + ) + return merged + + +def merge_output_language_from_messages(messages: list[Message]) -> str | None: + if not any( + getattr(part, "text", None) + for message in messages or [] + for part in getattr(message, "parts", []) + ): + return None + return SessionExtractContextProvider(messages=messages).get_output_language() + + +def clone_operation_for_uri(op: ResolvedOperation, uri: str) -> ResolvedOperation: + old_file = getattr(op, "old_memory_file_content", None) + if old_file is not None and getattr(old_file, "uri", None) not in (None, uri): + old_file = None + return op.model_copy( + update={ + "uris": [uri], + "memory_fields": dict(getattr(op, "memory_fields", {}) or {}), + "old_memory_file_content": old_file, + "source": getattr(op, "source", None), + }, + deep=True, + ) + + +def memory_file_to_delete_patch( + mf: MemoryFile, + *, + schema: MemoryTypeSchema, + extract_context: ExtractContext, +) -> PatchMergePatch: + """Convert a delete-file MemoryFile to a PatchMergePatch. + + The before_file is the original content; after_file is empty content, + representing a deletion proposal. The merge LLM should put deleted files + in delete_ids. + """ + after_file = MemoryFile( + uri=mf.uri, + memory_type=mf.memory_type, + content="", + extra_fields=dict(mf.extra_fields or {}), + ) + return PatchMergePatch( + before_file=mf, + after_file=after_file, + ) + + +def operation_to_patch( + op: ResolvedOperation, + *, + schema: MemoryTypeSchema, + extract_context: ExtractContext, +) -> PatchMergePatch: + old_file = getattr(op, "old_memory_file_content", None) + after_file = render_operation_after_file( + op, + schema=schema, + extract_context=extract_context, + ) + return PatchMergePatch( + before_file=old_file, + after_file=after_file, + ) + + +def render_operation_after_file( + op: ResolvedOperation, + *, + schema: MemoryTypeSchema, + extract_context: ExtractContext, +) -> MemoryFile: + after_content = render_operation_after_file_content( + op, + schema=schema, + extract_context=extract_context, + ) + return MemoryFileUtils.read(after_content, uri=_first_uri(getattr(op, "uris", []) or [])) + + +def render_operation_after_file_content( + op: ResolvedOperation, + *, + schema: MemoryTypeSchema, + extract_context: ExtractContext, +) -> str: + old_content = getattr(op, "old_memory_file_content", None) + metadata: dict[str, Any] = dict(getattr(op, "memory_fields", {}) or {}) + source_extraction_id = source_extraction_id_for_operation(op) + if source_extraction_id: + metadata["source_extraction_id"] = source_extraction_id + source_trace_id = source_trace_id_for_operation(op) + if source_trace_id: + metadata["last_update_trace_id"] = source_trace_id + for field_def in schema.fields: + if field_def.name not in metadata: + continue + if old_content is None: + current_value = None + elif field_def.name == "content": + current_value = old_content.plain_content() + else: + current_value = old_content.extra_fields.get(field_def.name) + try: + metadata[field_def.name] = MergeOpFactory.from_field(field_def).apply( + current_value, + metadata[field_def.name], + ) + except Exception as exc: + logger.debug( + "Failed to preview memory patch field: memory_type=%s field=%s", + op.memory_type, + field_def.name, + exc_info=True, + ) + tracer.info( + "[streaming_memory_updater] skipping preview field update after merge_op failure " + f"memory_type={op.memory_type} field={field_def.name} error={exc}" + ) + if current_value is None: + metadata.pop(field_def.name, None) + else: + metadata[field_def.name] = current_value + + if old_content and old_content.extra_fields: + schema_field_names = {field.name for field in schema.fields} | {"content", "memory_type"} + for key, value in old_content.extra_fields.items(): + if key not in schema_field_names and key not in metadata and value is not None: + metadata[key] = value + metadata["version"] = next_memory_version(old_content) + metadata.setdefault("memory_type", op.memory_type) + mf = MemoryFile.from_parsed(uri=_first_uri(op.uris), parsed=dict(metadata)) + return MemoryFileUtils.write( + mf, + content_template=schema.content_template, + extract_context=extract_context, + ) + + +def classify_memory_merge_mode( + operations: list[ResolvedOperation], + *, + schema: MemoryTypeSchema | None = None, +) -> tuple[bool, str]: + if not operations: + return True, "empty_batch" + + uris = [_first_uri(op.uris) for op in operations] + unique_uri_count = len(set(uris)) + duplicate_target_count = len(uris) - unique_uri_count + all_new_files = all(getattr(op, "old_memory_file_content", None) is None for op in operations) + operation_mode = getattr(schema, "operation_mode", "") if schema is not None else "" + + if operation_mode == "add_only": + return True, "add_only" + if is_cross_extraction_group(operations): + return False, "cross_extraction_batch" + # Multi-patch batches always go through LLM merge even if all files are new and + # URIs are unique — the LLM handles semantic deduplication and directory name + # normalization (e.g. activity vs activities, art_form vs art_forms). + if len(operations) > 1: + return False, "multi_patch_semantic_merge" + if all_new_files and duplicate_target_count == 0: + return True, "unique_new_files" + + op = operations[0] + old_file = getattr(op, "old_memory_file_content", None) + if old_file is None: + return True, "single_new_file" + fields = dict(getattr(op, "memory_fields", {}) or {}) + if "content" not in fields: + return False, "single_existing_non_content_patch" + old_plain_content = old_file.plain_content().strip() + if schema is not None: + try: + after_content = render_operation_after_file_content( + op, + schema=schema, + extract_context=ExtractContext([]), + ) + after_file = MemoryFileUtils.read( + after_content, uri=_first_uri(getattr(op, "uris", []) or []) + ) + if old_plain_content == after_file.plain_content().strip(): + return True, "single_existing_content_unchanged" + except Exception as exc: + logger.debug( + "Failed to render memory patch preview for merge-mode classification: " + "memory_type=%s", + getattr(op, "memory_type", None), + exc_info=True, + ) + tracer.info( + "[streaming_memory_updater] merge-mode preview failed; falling back to " + f"raw content comparison memory_type={getattr(op, 'memory_type', None)} " + f"error={exc}" + ) + if old_plain_content == str(fields.get("content") or "").strip(): + return True, "single_existing_content_unchanged" + return False, "single_existing_content_changed" + + +def _inherit_source_metadata_to_merged_operations( + input_operations: list[ResolvedOperation], + merged_operations: list[ResolvedOperation], +) -> None: + """Best-effort provenance restore after patch-merge LLM output. + + Patch merge hides system provenance fields from the model, so generated + operations can lose source_extraction_id. Reattach it by exact URI match + where possible. If a merged output has no URI match but only one input + source exists, copy that source; otherwise record all input source IDs as an + ambiguous multi-source operation. + """ + + input_by_uri: dict[str, list[ResolvedOperation]] = {} + all_source_ids: set[str] = set() + for input_op in input_operations or []: + op_source_ids = _operation_source_extraction_ids(input_op) + all_source_ids.update(op_source_ids) + for uri in list(getattr(input_op, "uris", []) or []): + if uri: + input_by_uri.setdefault(uri, []).append(input_op) + + if not all_source_ids: + return + + for merged_op in merged_operations or []: + if _operation_source_extraction_ids(merged_op): + continue + matched_inputs: list[ResolvedOperation] = [] + for uri in list(getattr(merged_op, "uris", []) or []): + matched_inputs.extend(input_by_uri.get(uri, [])) + matched_ids = { + source_id + for input_op in matched_inputs + for source_id in _operation_source_extraction_ids(input_op) + } + if len(matched_ids) == 1: + _set_operation_source_extraction_id(merged_op, next(iter(matched_ids))) + elif len(matched_ids) > 1: + merged_op.memory_fields["source_extraction_ids"] = sorted(matched_ids) + elif len(all_source_ids) == 1: + _set_operation_source_extraction_id(merged_op, next(iter(all_source_ids))) + else: + merged_op.memory_fields["source_extraction_ids"] = sorted(all_source_ids) + + +def _set_operation_source_extraction_id(op: ResolvedOperation, extraction_id: str) -> None: + op.memory_fields["source_extraction_id"] = extraction_id + source = getattr(op, "source", None) + if source is None: + op.source = MemoryOperationSource(extraction_id=extraction_id) + elif not getattr(source, "extraction_id", None): + source.extraction_id = extraction_id + + +def enforce_merge_group_peer_id( + operations: list[ResolvedOperation], + *, + peer_id: str | None, + memory_type: str, + registry: MemoryTypeRegistry, + ctx: RequestContext, +) -> None: + """Pin merged operations to the peer scope selected by group-by. + + The second-stage merge LLM may omit or hallucinate peer_id. The group key is + authoritative because it is decided before merge from the original request + routing; all merged upserts must therefore be rewritten to that scope. + """ + schema = registry.get(memory_type) + effective_peer_id = peer_id if getattr(schema, "peer_enabled", True) else None + for op in operations or []: + if op.memory_type != memory_type: + continue + if effective_peer_id: + op.memory_fields["peer_id"] = effective_peer_id + else: + op.memory_fields.pop("peer_id", None) + if schema is not None: + op.uris = _uris_for_merge_group_operation( + op, + schema=schema, + ctx=ctx, + peer_id=effective_peer_id, + ) + + +def _uris_for_merge_group_operation( + op: ResolvedOperation, + *, + schema: MemoryTypeSchema, + ctx: RequestContext, + peer_id: str | None, +) -> list[str]: + fields = dict(op.memory_fields or {}) + user_id = getattr(getattr(ctx, "user", None), "user_id", None) or fields.get("user_id") + if not user_id: + return list(op.uris or []) + fields["user_id"] = user_id + if peer_id: + fields["peer_id"] = peer_id + user_space = f"{user_id}/peers/{peer_id}" + else: + fields.pop("peer_id", None) + user_space = user_id + try: + from openviking.session.memory.utils.uri import generate_uri + + return [ + generate_uri( + memory_type=schema, + fields=fields, + user_space=user_space, + ) + ] + except Exception as exc: + tracer.info( + "[streaming_memory_updater] failed to enforce merge group uri " + f"memory_type={op.memory_type} peer_id={peer_id} old_uris={op.uris} error={exc}" + ) + return list(op.uris or []) + + +def _peer_id_from_uri(uri: str | None) -> str | None: + if not uri: + return None + match = re.search(r"/peers/([^/]+)/memories/", uri) + if not match: + return None + return safe_peer_id(match.group(1)) + + +def _peer_id_for_operation(op: ResolvedOperation) -> str | None: + """Get peer_id from a resolved operation, falling back to peer URI scope. + + Returns None for self (user-level) memories. + """ + fields = dict(getattr(op, "memory_fields", {}) or {}) + peer_id = safe_peer_id(fields.get("peer_id")) + if peer_id: + return peer_id + old_file = getattr(op, "old_memory_file_content", None) + if old_file is not None: + old_peer_id = safe_peer_id((old_file.extra_fields or {}).get("peer_id")) + if old_peer_id: + return old_peer_id + old_uri_peer_id = _peer_id_from_uri(getattr(old_file, "uri", None)) + if old_uri_peer_id: + return old_uri_peer_id + for uri in getattr(op, "uris", []) or []: + uri_peer_id = _peer_id_from_uri(uri) + if uri_peer_id: + return uri_peer_id + return None + + +def _peer_id_for_memory_file(mf: MemoryFile) -> str | None: + """Get peer_id from a MemoryFile, falling back to peer URI scope. + + Returns None for self (user-level) memories. + """ + peer_id = safe_peer_id((mf.extra_fields or {}).get("peer_id")) + return peer_id or _peer_id_from_uri(mf.uri) + + +def _unique_operation_uris(operations: list[ResolvedOperation]) -> list[str]: + return list(dict.fromkeys(uri for op in operations for uri in (op.uris or []) if uri)) + + +def attach_source_to_request_operations(request: MemoryUpdateRequest) -> None: + source = memory_operation_source_from_request(request) + if source is None: + return + for op in list(getattr(request.operations, "upsert_operations", []) or []): + if getattr(op, "source", None) is None: + op.source = source + source_extraction_id = getattr(op.source, "extraction_id", None) + if source_extraction_id: + op.memory_fields.setdefault("source_extraction_id", source_extraction_id) + source_trace_id = getattr(op.source, "trace_id", None) + if source_trace_id: + op.memory_fields.setdefault("last_update_trace_id", source_trace_id) + + +def memory_operation_source_from_request( + request: MemoryUpdateRequest, +) -> MemoryOperationSource | None: + metadata = dict(getattr(request, "metadata", {}) or {}) + extraction_id = metadata.get("source_extraction_id") or metadata.get("extraction_id") + if not extraction_id: + return None + return MemoryOperationSource( + extraction_id=str(extraction_id), + session_id=_optional_str(metadata.get("session_id")), + archive_uri=_optional_str(metadata.get("archive_uri")), + task_id=_optional_str(metadata.get("task_id")), + trace_id=_optional_str(metadata.get("trace_id")), + extracted_at=_optional_str(metadata.get("extracted_at")), + ) + + +def source_extraction_id_for_operation(op: ResolvedOperation) -> str | None: + source = getattr(op, "source", None) + extraction_id = getattr(source, "extraction_id", None) if source is not None else None + if extraction_id: + return str(extraction_id) + fields = dict(getattr(op, "memory_fields", {}) or {}) + field_value = fields.get("source_extraction_id") + return str(field_value) if field_value else None + + +def source_trace_id_for_operation(op: ResolvedOperation) -> str | None: + source = getattr(op, "source", None) + trace_id = getattr(source, "trace_id", None) if source is not None else None + if trace_id: + return str(trace_id) + fields = dict(getattr(op, "memory_fields", {}) or {}) + field_value = fields.get("last_update_trace_id") or fields.get("trace_id") + if field_value: + return str(field_value) + current_trace_id = get_trace_id() + return current_trace_id or None + + +def is_cross_extraction_group(operations: list[ResolvedOperation]) -> bool: + extraction_ids = { + extraction_id + for extraction_id in (source_extraction_id_for_operation(op) for op in operations) + if extraction_id + } + return len(extraction_ids) > 1 + + +def _optional_str(value: Any) -> str | None: + if value is None: + return None + text = str(value) + return text if text else None + + +def seed_patch_merge_read_contents( + provider: PatchMergeContextProvider, operations: list[ResolvedOperation] +) -> None: + for op in operations: + old_file = getattr(op, "old_memory_file_content", None) + uri = _first_uri(getattr(op, "uris", []) or []) + if old_file is not None and uri: + provider.read_file_contents[uri] = old_file + + +def safe_get_viking_fs() -> Any | None: + try: + return get_viking_fs() + except Exception: + return None + + +def merge_link_lists(*link_lists: list[StoredLink]) -> list[StoredLink]: + """Merge links by endpoint/type/anchor, preferring stronger metadata.""" + + merged: dict[tuple[str, str, str, str | None], StoredLink] = {} + for links in link_lists: + for link in links or []: + key = (link.from_uri, link.to_uri, link.link_type, link.match_text) + current = merged.get(key) + if current is None: + merged[key] = link + continue + current_weight = float(current.weight or 0.0) + new_weight = float(link.weight or 0.0) + if new_weight > current_weight: + current.weight = link.weight + if len(link.description or "") > len(current.description or ""): + current.description = link.description + if not current.created_at and link.created_at: + current.created_at = link.created_at + return list(merged.values()) + + +async def filter_valid_links( + links: list[StoredLink], + *, + upsert_operations: list[ResolvedOperation], + delete_file_contents: list[MemoryFile], + ctx: RequestContext, + trace_console: bool = False, +) -> list[StoredLink]: + """Drop links whose endpoints are deleted or missing from storage.""" + + if not links: + return [] + upsert_uris = {uri for op in upsert_operations for uri in (op.uris or []) if uri} + deleted_uris = {file.uri for file in delete_file_contents if getattr(file, "uri", None)} + viking_fs = safe_get_viking_fs() + endpoint_exists_cache: dict[str, bool] = {} + + async def _endpoint_exists(uri: str) -> bool: + if not uri or uri in deleted_uris: + return False + if uri in upsert_uris: + return True + if uri in endpoint_exists_cache: + return endpoint_exists_cache[uri] + if viking_fs is None: + endpoint_exists_cache[uri] = False + return False + try: + content = await viking_fs.read_file(uri, ctx=ctx) + exists = bool(content) + except Exception: + exists = False + endpoint_exists_cache[uri] = exists + return exists + + valid_links: list[StoredLink] = [] + dropped = 0 + for link in merge_link_lists(links): + if await _endpoint_exists(link.from_uri) and await _endpoint_exists(link.to_uri): + valid_links.append(link) + else: + dropped += 1 + + tracer.info( + "[streaming_memory_updater] links filtered " + f"input_links={len(links)} output_links={len(valid_links)} dropped_links={dropped}", + console=trace_console, + ) + return valid_links + + +def split_links_for_append_only_ops( + links: list[StoredLink], + *, + append_ops: list[ResolvedOperation], + merge_ops: list[ResolvedOperation], +) -> tuple[list[StoredLink], list[StoredLink]]: + append_uris = {uri for op in append_ops for uri in (op.uris or []) if uri} + merge_uris = {uri for op in merge_ops for uri in (op.uris or []) if uri} + append_links: list[StoredLink] = [] + merge_links: list[StoredLink] = [] + for link in links: + touches_append = link.from_uri in append_uris or link.to_uri in append_uris + touches_merge = link.from_uri in merge_uris or link.to_uri in merge_uris + if touches_append and not touches_merge: + append_links.append(link) + else: + merge_links.append(link) + return append_links, merge_links + + +def clone_memory_update_request( + request: MemoryUpdateRequest, + *, + operations: ResolvedOperations, +) -> MemoryUpdateRequest: + return MemoryUpdateRequest( + operations=operations, + messages=list(request.messages or []), + ctx=request.ctx, + strict_extract_errors=request.strict_extract_errors, + isolation_options=dict(request.isolation_options or {}), + metadata=dict(request.metadata or {}), + ) + + +def scope_memory_update_result_to_submitter( + result: StreamingMemoryUpdateResult, + request: MemoryUpdateRequest, +) -> StreamingMemoryUpdateResult: + """Return the submitting request's view of a shared streaming flush. + + StreamingBatcher intentionally resolves every waiter in one flush with the + same aggregate batch result. Per-session consumers (archive memory_diff, + contexts, case URI mapping) must not see writes/deletes that were produced + by other concurrently flushed commits. + """ + + scope = _memory_submitter_scope_from_request(request) + if scope.is_empty: + return result + + scoped_operations = _scope_operations_to_submitter(result.operations, scope=scope) + operation_uris = _operation_uri_set(scoped_operations) + submitter_uris = _request_uri_set(request) + scoped_link_uris = _link_endpoint_uri_set( + getattr(scoped_operations, "resolved_links", []) or [] + ) + scoped_uris = operation_uris | submitter_uris | scoped_link_uris + + scoped_apply_result = _scope_apply_result_to_uris( + result.apply_result, + scoped_uris=scoped_uris, + ) + metadata = dict(result.metadata or {}) + metadata.update( + { + "batch_request_count": result.request_count, + "batch_operation_count": metadata.get("operation_count"), + "request_count": 1, + "operation_count": _operation_count(scoped_operations), + "source": "streaming_memory_scoped", + "scoped_to_submitter": True, + } + ) + if scope.extraction_id: + metadata["scoped_to_source_extraction_id"] = scope.extraction_id + if scope.archive_uri: + metadata["scoped_to_archive_uri"] = scope.archive_uri + if scope.session_id: + metadata["scoped_to_session_id"] = scope.session_id + metadata["unscoped_written_uris"] = list(getattr(result.apply_result, "written_uris", []) or []) + metadata["unscoped_edited_uris"] = list(getattr(result.apply_result, "edited_uris", []) or []) + metadata["unscoped_deleted_uris"] = list(getattr(result.apply_result, "deleted_uris", []) or []) + + return StreamingMemoryUpdateResult( + operations=scoped_operations, + apply_result=scoped_apply_result, + request_count=1, + metadata=metadata, + ) + + +@dataclass(frozen=True, slots=True) +class _MemorySubmitterScope: + extraction_id: str | None = None + session_id: str | None = None + archive_uri: str | None = None + request_uris: frozenset[str] = frozenset() + + @property + def is_empty(self) -> bool: + return ( + not self.extraction_id + and not self.session_id + and not self.archive_uri + and not self.request_uris + ) + + +def _memory_submitter_scope_from_request(request: MemoryUpdateRequest) -> _MemorySubmitterScope: + metadata = dict(getattr(request, "metadata", {}) or {}) + source = memory_operation_source_from_request(request) + extraction_id = _optional_str( + metadata.get("source_extraction_id") + or metadata.get("extraction_id") + or getattr(source, "extraction_id", None) + ) + return _MemorySubmitterScope( + extraction_id=extraction_id, + session_id=_optional_str(metadata.get("session_id") or getattr(source, "session_id", None)), + archive_uri=_optional_str( + metadata.get("archive_uri") or getattr(source, "archive_uri", None) + ), + request_uris=frozenset(_request_uri_set(request)), + ) + + +def _scope_operations_to_submitter( + operations: ResolvedOperations, + *, + scope: _MemorySubmitterScope, +) -> ResolvedOperations: + upserts = [ + op + for op in list(getattr(operations, "upsert_operations", []) or []) + if _operation_matches_scope(op, scope=scope) + ] + deletes = [ + file + for file in list(getattr(operations, "delete_file_contents", []) or []) + if _memory_file_matches_scope(file, scope=scope) + ] + kept_uris = _operation_uri_set( + ResolvedOperations(upsert_operations=upserts, delete_file_contents=deletes, errors=[]) + ) + request_uris = set(scope.request_uris) + return ResolvedOperations( + upsert_operations=upserts, + delete_file_contents=deletes, + errors=list(getattr(operations, "errors", []) or []), + resolved_links=[ + link + for link in list(getattr(operations, "resolved_links", []) or []) + if _link_matches_scoped_uris(link, scoped_uris=kept_uris, request_uris=request_uris) + ], + delete_replacements={ + str(deleted_uri): str(replacement_uri) + for deleted_uri, replacement_uri in dict( + getattr(operations, "delete_replacements", {}) or {} + ).items() + if str(deleted_uri) in kept_uris or str(replacement_uri) in kept_uris + }, + ) + + +def _scope_apply_result_to_uris( + apply_result: MemoryUpdateResult, + *, + scoped_uris: set[str], +) -> MemoryUpdateResult: + scoped = MemoryUpdateResult() + scoped.written_uris = [ + uri for uri in list(getattr(apply_result, "written_uris", []) or []) if uri in scoped_uris + ] + scoped.edited_uris = [ + uri for uri in list(getattr(apply_result, "edited_uris", []) or []) if uri in scoped_uris + ] + scoped.deleted_uris = [ + uri for uri in list(getattr(apply_result, "deleted_uris", []) or []) if uri in scoped_uris + ] + scoped.errors = [ + error + for error in list(getattr(apply_result, "errors", []) or []) + if _apply_error_matches_scoped_uris(error, scoped_uris=scoped_uris) + ] + return scoped + + +def _operation_matches_scope(op: ResolvedOperation, *, scope: _MemorySubmitterScope) -> bool: + if scope.extraction_id and scope.extraction_id in _operation_source_extraction_ids(op): + return True + source = getattr(op, "source", None) + if ( + scope.archive_uri + and _optional_str(getattr(source, "archive_uri", None)) == scope.archive_uri + ): + return True + if scope.session_id and _optional_str(getattr(source, "session_id", None)) == scope.session_id: + return True + if scope.request_uris and any( + uri in scope.request_uris for uri in list(getattr(op, "uris", []) or []) + ): + return True + return False + + +def _memory_file_matches_scope(file: MemoryFile, *, scope: _MemorySubmitterScope) -> bool: + fields = dict(getattr(file, "extra_fields", {}) or {}) + if scope.extraction_id and scope.extraction_id in _source_extraction_ids_from_fields(fields): + return True + uri = getattr(file, "uri", None) + return bool(uri and uri in scope.request_uris) + + +def _operation_source_extraction_ids(op: ResolvedOperation) -> set[str]: + fields = dict(getattr(op, "memory_fields", {}) or {}) + ids = _source_extraction_ids_from_fields(fields) + source_id = source_extraction_id_for_operation(op) + if source_id: + ids.add(source_id) + return ids + + +def _source_extraction_ids_from_fields(fields: dict[str, Any]) -> set[str]: + ids: set[str] = set() + value = fields.get("source_extraction_id") + if value: + ids.add(str(value)) + values = fields.get("source_extraction_ids") + if isinstance(values, (list, tuple, set)): + ids.update(str(item) for item in values if item) + elif values: + ids.add(str(values)) + return ids + + +def _request_uri_set(request: MemoryUpdateRequest) -> set[str]: + return _operation_uri_set(getattr(request, "operations", None)) + + +def _operation_uri_set(operations: ResolvedOperations | None) -> set[str]: + if operations is None: + return set() + uris = { + uri + for op in list(getattr(operations, "upsert_operations", []) or []) + for uri in list(getattr(op, "uris", []) or []) + if uri + } + uris.update( + str(file.uri) + for file in list(getattr(operations, "delete_file_contents", []) or []) + if getattr(file, "uri", None) + ) + return uris + + +def _link_endpoint_uri_set(links: list[StoredLink]) -> set[str]: + uris: set[str] = set() + for link in links or []: + from_uri = str(getattr(link, "from_uri", "") or "") + to_uri = str(getattr(link, "to_uri", "") or "") + if from_uri: + uris.add(from_uri) + if to_uri: + uris.add(to_uri) + return uris + + +def _link_matches_scoped_uris( + link: StoredLink, + *, + scoped_uris: set[str], + request_uris: set[str], +) -> bool: + if not scoped_uris: + return False + from_uri = str(getattr(link, "from_uri", "") or "") + to_uri = str(getattr(link, "to_uri", "") or "") + # Keep links between this submitter's touched files and their neighbors, but + # do not leak links that only connect other submitters' files. + return ( + from_uri in scoped_uris + or to_uri in scoped_uris + or from_uri in request_uris + or to_uri in request_uris + ) + + +def _apply_error_matches_scoped_uris(error: Any, *, scoped_uris: set[str]) -> bool: + if not scoped_uris: + return False + try: + uri = error[0] + except Exception: + return True + return str(uri) in scoped_uris or str(uri) == "unknown" + + +def combine_streaming_memory_results( + *results: StreamingMemoryUpdateResult | None, + fallback_request_count: int = 0, +) -> StreamingMemoryUpdateResult: + present_results = [result for result in results if result is not None] + if not present_results: + return StreamingMemoryUpdateResult( + operations=ResolvedOperations(upsert_operations=[], delete_file_contents=[], errors=[]), + apply_result=MemoryUpdateResult(), + request_count=fallback_request_count, + metadata={"flush_reason": "empty", "operation_count": 0}, + ) + if len(present_results) == 1: + return present_results[0] + + combined_operations = ResolvedOperations( + upsert_operations=[], + delete_file_contents=[], + errors=[], + resolved_links=[], + delete_replacements={}, + ) + combined_apply_result = MemoryUpdateResult() + metadata: dict[str, Any] = { + "flush_reason": "+".join( + str(result.metadata.get("flush_reason", "unknown")) for result in present_results + ), + "combined_result": True, + } + request_count = 0 + for result in present_results: + request_count += result.request_count + combined_operations.upsert_operations.extend(result.operations.upsert_operations or []) + combined_operations.delete_file_contents.extend( + result.operations.delete_file_contents or [] + ) + combined_operations.errors.extend(result.operations.errors or []) + combined_operations.resolved_links = merge_link_lists( + combined_operations.resolved_links, + list(getattr(result.operations, "resolved_links", []) or []), + ) + combined_operations.delete_replacements.update( + dict(getattr(result.operations, "delete_replacements", {}) or {}) + ) + combined_apply_result.written_uris.extend(result.apply_result.written_uris) + combined_apply_result.edited_uris.extend(result.apply_result.edited_uris) + combined_apply_result.deleted_uris.extend(result.apply_result.deleted_uris) + combined_apply_result.errors.extend(result.apply_result.errors) + for key in ("batch_id", "batch_trace_id"): + if result.metadata.get(key): + metadata.setdefault(key, result.metadata.get(key)) + if result.metadata.get("fast_path"): + metadata["fast_path"] = True + metadata["operation_count"] = _operation_count(combined_operations) + return StreamingMemoryUpdateResult( + operations=combined_operations, + apply_result=combined_apply_result, + request_count=request_count or fallback_request_count, + metadata=metadata, + ) + + +def _combined_request_messages(items: list[MemoryUpdateRequest]) -> list[Message]: + messages: list[Message] = [] + for item in items: + messages.extend(item.messages) + return messages + + +def _make_isolation_handler( + request: MemoryUpdateRequest, + extract_context: ExtractContext, +) -> MemoryIsolationHandler: + options = dict(request.isolation_options or {}) + return MemoryIsolationHandler( + request.ctx, + extract_context, + allowed_memory_types=options.get("allowed_memory_types"), + allow_self=options.get("allow_self", True), + allowed_peer_ids=options.get("allowed_peer_ids"), + ) + + +def _operation_count(operations: ResolvedOperations) -> int: + return len(operations.upsert_operations or []) + len(operations.delete_file_contents or []) + + +def _first_uri(uris: list[str] | None) -> str | None: + return uris[0] if uris else None + + +_streaming_memory_updater_registry: dict[Hashable, StreamingMemoryUpdater] = {} +_streaming_memory_updater_registry_lock = threading.RLock() + + +async def get_streaming_memory_updater( + *, + key: StreamingMemoryUpdaterKey | Hashable, + registry: MemoryTypeRegistry | None = None, + vikingdb: Any = None, + config: StreamingMemoryUpdaterConfig | None = None, +) -> StreamingMemoryUpdater: + """Get or create the process-global streaming updater for one user key.""" + + with _streaming_memory_updater_registry_lock: + existing = _streaming_memory_updater_registry.get(key) + if existing is not None: + return existing + updater = StreamingMemoryUpdater( + registry=registry, + vikingdb=vikingdb, + config=config or StreamingMemoryUpdaterConfig(), + ) + _streaming_memory_updater_registry[key] = updater + return updater + + +def make_streaming_memory_updater_key(*, request_context: Any) -> StreamingMemoryUpdaterKey: + user = getattr(request_context, "user", None) + account_id = ( + getattr(request_context, "account_id", None) + or getattr(user, "account_id", None) + or "default" + ) + user_id = getattr(request_context, "user_id", None) or getattr(user, "user_id", None) or "" + return StreamingMemoryUpdaterKey(account_id=str(account_id), user_id=str(user_id)) diff --git a/openviking/session/memory/tools.py b/openviking/session/memory/tools.py index 74a159c7db..550bbcc116 100644 --- a/openviking/session/memory/tools.py +++ b/openviking/session/memory/tools.py @@ -21,6 +21,12 @@ logger = get_logger(__name__) +_LLM_HIDDEN_MEMORY_FIELDS = { + "source_extraction_id", + "source_extraction_ids", + "last_update_trace_id", +} + def optimize_search_result(result: Any, limit: int = 10) -> Any: """优化搜索结果以减少 Token 消耗,并过滤掉抽象文件。""" @@ -192,6 +198,8 @@ async def execute( llm_result = mf.to_metadata() llm_result.pop("links", None) llm_result.pop("backlinks", None) + for hidden_field in _LLM_HIDDEN_MEMORY_FIELDS: + llm_result.pop(hidden_field, None) # Annotate with page_id for link extraction if ctx and ctx.page_id_map: page_id = ctx.page_id_map.get_page_id(uri) @@ -212,7 +220,6 @@ async def execute( ) return llm_result except NotFoundError as e: - tracer.info(f"read not found: {uri}") return {"error": str(e)} except Exception as e: tracer.error(f"Failed to execute read: {e}") diff --git a/openviking/session/memory/utils/__init__.py b/openviking/session/memory/utils/__init__.py index 2cbb32792e..cd031e1149 100644 --- a/openviking/session/memory/utils/__init__.py +++ b/openviking/session/memory/utils/__init__.py @@ -19,6 +19,7 @@ detect_language_from_conversation, resolve_output_language, resolve_output_language_from_conversation, + resolve_output_language_from_text, resolve_with_override, strip_language_detection_noise, ) @@ -47,6 +48,7 @@ "detect_language_from_conversation", "resolve_output_language", "resolve_output_language_from_conversation", + "resolve_output_language_from_text", "resolve_with_override", "strip_language_detection_noise", "add_line_numbers", diff --git a/openviking/session/memory/utils/language.py b/openviking/session/memory/utils/language.py index 3830b695de..7b57c668af 100644 --- a/openviking/session/memory/utils/language.py +++ b/openviking/session/memory/utils/language.py @@ -356,6 +356,22 @@ def resolve_with_override(config, detect: Callable[[], str]) -> str: return detect() +def resolve_output_language_from_text( + text: str, + config=None, + *, + fallback_language: str = "en", +) -> str: + """Resolve output language from text with an explicit fallback language. + + Unlike ``resolve_output_language``, this helper does not consult locale or + timezone. Use it when an empty or low-signal text should not inherit the + runtime environment language. + """ + fallback = (fallback_language or "en").strip() or "en" + return resolve_with_override(config, lambda: _detect_language_from_text(text, fallback)) + + def strip_language_detection_noise(text: str) -> str: """Remove URI-like machine tokens that should not affect output language.""" return _URI_LANGUAGE_NOISE_RE.sub(" ", text or "") @@ -364,7 +380,7 @@ def strip_language_detection_noise(text: str) -> str: def resolve_output_language(text: str, config=None) -> str: """Resolve output language from text, honoring config override before detection.""" fallback = _resolve_system_fallback_language("en") - return resolve_with_override(config, lambda: _detect_language_from_text(text, fallback)) + return resolve_output_language_from_text(text, config=config, fallback_language=fallback) def resolve_output_language_from_conversation(conversation: str, config=None) -> str: diff --git a/openviking/session/memory/utils/memory_file_utils.py b/openviking/session/memory/utils/memory_file_utils.py index 8a38146399..2377787cd9 100644 --- a/openviking/session/memory/utils/memory_file_utils.py +++ b/openviking/session/memory/utils/memory_file_utils.py @@ -1,13 +1,17 @@ import json +import logging import re from datetime import datetime from typing import Any, Dict, Optional from openviking.session.memory.dataclass import MemoryFile +from openviking.session.memory.utils.link_renderer import LinkRenderer from openviking.session.memory.utils.messages import parse_memory_file_with_fields from openviking.session.memory.utils.uri import render_template from openviking.utils.time_utils import parse_iso_datetime +logger = logging.getLogger(__name__) + # Regex patterns for MEMORY_FIELDS HTML comment _MEMORY_FIELDS_PATTERN = re.compile(r"\n\n", re.DOTALL) _MEMORY_FIELDS_PATTERN_END = re.compile(r"$", re.DOTALL) @@ -15,6 +19,29 @@ DEFAULT_TRUNCATE_MAX_CHARS = 1000 +def memory_version_from_fields(fields: Optional[Dict[str, Any]], *, default: int = 1) -> int: + """Return a positive MEMORY_FIELDS version, falling back to ``default``.""" + try: + version = int((fields or {}).get("version")) + except (TypeError, ValueError): + return default + return version if version > 0 else default + + +def next_memory_version(old_file: Optional[MemoryFile]) -> int: + """Return the next persisted MEMORY_FIELDS version for a write.""" + if old_file is None: + return 1 + return memory_version_from_fields(old_file.extra_fields, default=1) + 1 + + +def bump_memory_version(memory_file: MemoryFile) -> None: + """Increment a MemoryFile's persisted MEMORY_FIELDS version in-place.""" + memory_file.extra_fields["version"] = memory_version_from_fields( + memory_file.extra_fields, default=1 + ) + 1 + + def _serialize_datetime(obj: Any) -> Any: if isinstance(obj, datetime): return obj.isoformat() @@ -32,10 +59,23 @@ def _deserialize_datetime(metadata: Dict[str, Any]) -> Dict[str, Any]: return result + + +def _uri_basename(uri: str) -> str: + name = str(uri or "").rstrip("/").rsplit("/", 1)[-1] + return name.removesuffix(".md") + + +def _template_link_target(source_uri: Optional[str], target_uri: str) -> str: + if source_uri and target_uri: + return LinkRenderer.relative_path(str(source_uri), str(target_uri)) or str(target_uri) + return str(target_uri or "") + def _serialize_with_metadata( metadata: Dict[str, Any], content_template: str = None, extract_context: Any = None, + source_uri: Optional[str] = None, ) -> str: content = metadata.pop("content", "") or "" @@ -43,15 +83,27 @@ def _serialize_with_metadata( try: template_vars = metadata.copy() template_vars["content"] = content + template_vars.setdefault("links", []) + template_vars.setdefault("backlinks", []) + template_vars["source_uri"] = source_uri or "" + template_vars["uri_basename"] = _uri_basename + template_vars["link_target"] = lambda target_uri: _template_link_target(source_uri, target_uri) content = render_template(content_template, template_vars, extract_context) except Exception: - pass + logger.exception( + "Failed to render memory content template; using plain content fallback" + ) clean_metadata = {k: v for k, v in metadata.items() if v is not None} if not clean_metadata: return content + clean_metadata.pop("_uri", None) + links = clean_metadata.get("links") + if isinstance(links, list) and source_uri: + content = LinkRenderer.render_links(content, str(source_uri), links) + metadata_json = json.dumps( clean_metadata, indent=2, default=_serialize_datetime, ensure_ascii=False ) @@ -91,6 +143,7 @@ def write( metadata, content_template=content_template, extract_context=extract_context, + source_uri=memory_file.uri, ) @staticmethod diff --git a/openviking/session/memory/utils/streaming_batcher.py b/openviking/session/memory/utils/streaming_batcher.py new file mode 100644 index 0000000000..d91966d3e0 --- /dev/null +++ b/openviking/session/memory/utils/streaming_batcher.py @@ -0,0 +1,234 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Shared async count/time window batcher for streaming session updates.""" + +from __future__ import annotations + +import asyncio +import contextlib +import time +import uuid +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from typing import Any, Generic, TypeVar + +from openviking.telemetry import tracer +from openviking_cli.utils import get_logger + +logger = get_logger(__name__) + +T = TypeVar("T") +R = TypeVar("R") + + +@dataclass(slots=True) +class StreamingBatcherConfig: + """Count/time window configuration shared by streaming updaters.""" + + max_items_per_batch: int = 8 + max_wait_seconds: float = 10.0 + timer_check_interval_seconds: float = 1.0 + + def __post_init__(self) -> None: + if self.max_items_per_batch <= 0: + raise ValueError("max_items_per_batch must be > 0") + if self.max_wait_seconds <= 0: + raise ValueError("max_wait_seconds must be > 0") + if self.timer_check_interval_seconds <= 0: + raise ValueError("timer_check_interval_seconds must be > 0") + + +@dataclass(slots=True) +class StreamingBatcher(Generic[T, R]): + """A reusable async batcher whose submit waits for its batch result. + + Items are buffered until either the buffered size reaches + ``max_items_per_batch`` or the oldest item waits for ``max_wait_seconds``. + Flush is performed by background tasks/timer; each ``submit`` awaits the + Future attached to its own batch item, so callers only return after the + batch containing their item has been processed. + """ + + name: str + process_batch: Callable[[list[T], str], Awaitable[R]] + config: StreamingBatcherConfig = field(default_factory=StreamingBatcherConfig) + item_size: Callable[[T], int] | None = None + result_metadata: Callable[[R], dict[str, Any] | None] | None = None + _buffer: list[_PendingBatchItem[T, R]] = field(init=False, repr=False) + _buffer_lock: asyncio.Lock = field(init=False, repr=False) + _flush_lock: asyncio.Lock = field(init=False, repr=False) + _timer_task: asyncio.Task[None] | None = field(init=False, default=None, repr=False) + _closed: bool = field(init=False, default=False, repr=False) + _last_result: R | None = field(init=False, default=None, repr=False) + + def __post_init__(self) -> None: + self._buffer = [] + self._buffer_lock = asyncio.Lock() + self._flush_lock = asyncio.Lock() + self._timer_task = None + self._closed = False + self._last_result = None + + @property + def closed(self) -> bool: + return self._closed + + @property + def last_result(self) -> R | None: + return self._last_result + + async def get_buffered_size(self) -> int: + async with self._buffer_lock: + return sum(self._item_size(item.payload) for item in self._buffer) + + async def get_buffered_item_count(self) -> int: + async with self._buffer_lock: + return len(self._buffer) + + async def submit(self, payload: T) -> R: + if self._closed: + raise RuntimeError(f"{self.name} is closed") + + self._ensure_timer_task() + loop = asyncio.get_running_loop() + future: asyncio.Future[R] = loop.create_future() + should_flush = False + async with self._buffer_lock: + self._buffer.append( + _PendingBatchItem( + payload=payload, + submitted_at=time.monotonic(), + future=future, + ) + ) + should_flush = self._buffered_size_unlocked() >= self.config.max_items_per_batch + + if should_flush: + self._trigger_background_flush("count") + + return await future + + async def close(self) -> R | None: + if self._closed: + return None + self._closed = True + await self._stop_timer_task() + return await self.flush("close") + + async def flush(self, reason: str) -> R | None: + async with self._flush_lock: + async with self._buffer_lock: + if not self._buffer: + return None + items = self._buffer + self._buffer = [] + + batch_id = uuid.uuid4().hex + batch_trace_id = uuid.uuid4().hex + try: + with tracer.start_as_current_span( + name=f"{self.name}.flush", + trace_id=batch_trace_id, + ): + tracer.set("batch_id", batch_id) + tracer.set("flush_reason", reason) + tracer.set("request_count", len(items)) + tracer.set("input_size", sum(self._item_size(item.payload) for item in items)) + result = await self.process_batch([item.payload for item in items], reason) + metadata = self._get_result_metadata(result) + if metadata is not None: + metadata.setdefault("batch_id", batch_id) + metadata.setdefault("batch_trace_id", batch_trace_id) + except Exception as exc: + for item in items: + if not item.future.done(): + item.future.set_exception(exc) + raise + + self._last_result = result + for item in items: + if not item.future.done(): + item.future.set_result(result) + return result + + def _ensure_timer_task(self) -> None: + if self._timer_task is not None and not self._timer_task.done(): + return + try: + loop = asyncio.get_running_loop() + except RuntimeError: + logger.warning("[%s] timer loop not started: no running event loop", self.name) + self._timer_task = None + return + self._timer_task = loop.create_task( + self._run_timer_loop(), + name=f"{self.name}-flush-loop", + ) + + async def _stop_timer_task(self) -> None: + task = self._timer_task + if task is None: + return + self._timer_task = None + if task.done(): + with contextlib.suppress(asyncio.CancelledError): + await task + return + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + async def _run_timer_loop(self) -> None: + while True: + try: + await asyncio.sleep(self.config.timer_check_interval_seconds) + if await self._should_flush_by_time_or_count(): + await self.flush("time") + except asyncio.CancelledError: + raise + except Exception as exc: + logger.warning("[%s] timer flush iteration failed: %s", self.name, exc) + + async def _should_flush_by_time_or_count(self) -> bool: + async with self._buffer_lock: + if not self._buffer: + return False + if self._buffered_size_unlocked() >= self.config.max_items_per_batch: + return True + oldest = min(item.submitted_at for item in self._buffer) + return (time.monotonic() - oldest) >= self.config.max_wait_seconds + + def _trigger_background_flush(self, reason: str) -> None: + try: + loop = asyncio.get_running_loop() + except RuntimeError: + return + + async def _runner() -> None: + try: + await self.flush(reason) + except Exception as exc: + logger.warning("[%s] background flush failed: %s", self.name, exc) + + loop.create_task(_runner(), name=f"{self.name}-flush-{reason}") + + def _buffered_size_unlocked(self) -> int: + return sum(self._item_size(item.payload) for item in self._buffer) + + def _item_size(self, payload: T) -> int: + if self.item_size is None: + return 1 + return max(0, int(self.item_size(payload))) + + def _get_result_metadata(self, result: R) -> dict[str, Any] | None: + if self.result_metadata is not None: + return self.result_metadata(result) + metadata = getattr(result, "metadata", None) + return metadata if isinstance(metadata, dict) else None + + +@dataclass(slots=True) +class _PendingBatchItem(Generic[T, R]): + payload: T + submitted_at: float + future: asyncio.Future[R] diff --git a/openviking/session/memory_policy.py b/openviking/session/memory_policy.py index 44a2f66bb1..03629d3342 100644 --- a/openviking/session/memory_policy.py +++ b/openviking/session/memory_policy.py @@ -9,18 +9,30 @@ from openviking_cli.exceptions import InvalidArgumentError -_POLICY_KEYS = {"self", "peer", "memory_types"} +_POLICY_KEYS = {"self", "peer", "memory_types", "working_memory"} _TARGET_KEYS = {"enabled"} -def _target_enabled(data: Any, *, default_enabled: bool) -> bool: +def _memory_policy_shape_error(key: str) -> str: + if key == "working_memory": + return "memory_policy.working_memory must be an object" + return "memory_policy target must be an object" + + +def _memory_policy_keys_error(key: str) -> str: + if key == "working_memory": + return "memory_policy.working_memory supports only: enabled" + return "memory_policy target supports only: enabled" + + +def _target_enabled(data: Any, *, default_enabled: bool, key: str = "target") -> bool: if data is None: return default_enabled if not isinstance(data, dict): - raise InvalidArgumentError("memory_policy target must be an object") + raise InvalidArgumentError(_memory_policy_shape_error(key)) extra_keys = set(data) - _TARGET_KEYS if extra_keys: - raise InvalidArgumentError("memory_policy target supports only: enabled") + raise InvalidArgumentError(_memory_policy_keys_error(key)) return bool(data.get("enabled", default_enabled)) @@ -44,6 +56,7 @@ class MemoryPolicy: self_enabled: bool = True peer_enabled: bool = True memory_types: Optional[set[str]] = None + working_memory_enabled: bool = True @classmethod def default(cls) -> "MemoryPolicy": @@ -63,9 +76,12 @@ def from_dict(cls, data: Any) -> "MemoryPolicy": "memory_policy supports only: " + ", ".join(sorted(_POLICY_KEYS)) ) return cls( - self_enabled=_target_enabled(data.get("self"), default_enabled=True), - peer_enabled=_target_enabled(data.get("peer"), default_enabled=True), + self_enabled=_target_enabled(data.get("self"), default_enabled=True, key="self"), + peer_enabled=_target_enabled(data.get("peer"), default_enabled=True, key="peer"), memory_types=_parse_memory_types(data.get("memory_types")), + working_memory_enabled=_target_enabled( + data.get("working_memory"), default_enabled=True, key="working_memory" + ), ) def validate_memory_types(self, known_memory_types: set[str]) -> None: @@ -82,6 +98,8 @@ def to_dict(self) -> dict[str, Any]: "self": {"enabled": self.self_enabled}, "peer": {"enabled": self.peer_enabled}, } + if not self.working_memory_enabled: + data["working_memory"] = {"enabled": False} if self.memory_types is not None: data["memory_types"] = sorted(self.memory_types) return data diff --git a/openviking/session/session.py b/openviking/session/session.py index 121ed87714..aa63f47ddc 100644 --- a/openviking/session/session.py +++ b/openviking/session/session.py @@ -1068,14 +1068,26 @@ async def list_tool_results( return {"tool_results": []} return await store.list(tool_name=tool_name, limit=limit) - def commit(self, keep_recent_count: int = 0) -> Dict[str, Any]: + def commit( + self, + keep_recent_count: int = 0, + *, + memory_policy: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: """Sync wrapper for commit_async().""" - return run_async(self.commit_async(keep_recent_count=keep_recent_count)) + return run_async( + self.commit_async( + keep_recent_count=keep_recent_count, + memory_policy=memory_policy, + ) + ) @tracer("session.commit.phase1") async def commit_async( self, keep_recent_count: int = 0, + *, + memory_policy: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: """Async commit session: archive immediately, extract memories in background. @@ -1100,7 +1112,9 @@ async def commit_async( trace_id = tracer.get_trace_id() keep_recent_count = max(0, int(keep_recent_count or 0)) - effective_policy = MemoryPolicy.from_dict(self._meta.memory_policy) + effective_policy = MemoryPolicy.from_dict( + memory_policy if memory_policy is not None else self._meta.memory_policy + ) _validate_memory_policy_types(effective_policy) effective_memory_policy = effective_policy.to_dict() logger.info( @@ -1308,12 +1322,25 @@ async def _run_memory_extraction( }, ) - latest_archive_overview = await self._get_latest_completed_archive_overview( - exclude_archive_uri=archive_uri + ov_config = get_openviking_config() + effective_policy = MemoryPolicy.from_dict(memory_policy) + working_memory_enabled = effective_policy.working_memory_enabled + latest_archive_overview = ( + await self._get_latest_completed_archive_overview( + exclude_archive_uri=archive_uri + ) + if working_memory_enabled + else "" ) extraction_messages = await self._hydrate_tool_outputs_for_extraction(messages) async def _run_archive_summary() -> None: + if not working_memory_enabled: + logger.info( + "Working Memory summary skipped " + "(memory_policy.working_memory.enabled=false)" + ) + return summary = await self._generate_archive_summary_async( extraction_messages, latest_archive_overview=latest_archive_overview, @@ -1361,12 +1388,10 @@ async def _run_retryable_phase2_step( ) # Summary, long-term memory, and execution-derived memory run concurrently. - ov_config = get_openviking_config() memory_extraction_enabled = ov_config.memory.extraction_enabled config_session_skill_extraction_enabled = ( ov_config.memory.session_skill_extraction_enabled ) - effective_policy = MemoryPolicy.from_dict(memory_policy) extraction_scope = _resolve_memory_extraction_scope( self.ctx, effective_policy, @@ -1407,10 +1432,13 @@ async def _run_retryable_phase2_step( self._session_compressor, "extract_execution_memories" ) - extraction_tasks: List[Any] = [ - _run_retryable_phase2_step("archive_summary", _run_archive_summary) - ] - extraction_labels = ["archive_summary"] + extraction_tasks: List[Any] = [] + extraction_labels: List[str] = [] + if working_memory_enabled: + extraction_tasks.append( + _run_retryable_phase2_step("archive_summary", _run_archive_summary) + ) + extraction_labels.append("archive_summary") if long_term_has_work: @@ -1495,7 +1523,9 @@ async def _run_execution_memory_extraction() -> Any: memory_diff_uri = candidate_memory_diff_uri total_extracted = 0 - for label, result in zip(extraction_labels[1:], _results[1:], strict=True): + for label, result in zip(extraction_labels, _results, strict=True): + if label == "archive_summary": + continue if isinstance(result, dict): target_contexts = list(result.get("contexts", [])) target_skills = list(result.get("session_skills", [])) @@ -1528,7 +1558,12 @@ async def _run_execution_memory_extraction() -> Any: "Memory and session skill extraction skipped " "(disabled by config or memory_policy)" ) - await _run_retryable_phase2_step("archive_summary", _run_archive_summary) + if working_memory_enabled: + await _run_retryable_phase2_step( + "archive_summary", _run_archive_summary + ) + else: + await _run_archive_summary() # Write relations (using snapshot, not self._usage_records) if self._viking_fs: diff --git a/openviking/session/train/__init__.py b/openviking/session/train/__init__.py new file mode 100644 index 0000000000..1ec9beab36 --- /dev/null +++ b/openviking/session/train/__init__.py @@ -0,0 +1,190 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Session training framework for trajectory/experience policy optimization.""" + +from openviking.session.train.batch_runner import ( + BatchTrainEvalConfig, + BatchTrainEvalReport, + run_batch_train_eval, +) +from openviking.session.train.components.case_loader import ( + ListCaseLoader, + TrialCaseLoader, + make_trial_case_loader, +) +from openviking.session.train.components.dataset_service import create_dataset_service_app +from openviking.session.train.components.event_recorder import ( + CompositeEventRecorder, + JsonlEventRecorder, + JsonlPipelineEventHook, +) +from openviking.session.train.components.gradient_estimator import ( + ExperienceGradientContext, + ExperienceGradientEstimator, +) +from openviking.session.train.components.memory_store import ExperienceSetLoader, SkillSetLoader +from openviking.session.train.components.policy_optimizer import ( + PatchMergePolicyOptimizer, + PatchMergePolicyOptimizerContext, +) +from openviking.session.train.components.policy_trainer import ( + BatchPolicyTrainer, + StreamingPolicyTrainer, + StreamingPolicyTrainerConfig, + StreamingPolicyTrainerKey, + get_streaming_policy_trainer, + make_streaming_policy_trainer_key, +) +from openviking.session.train.components.policy_updater import ( + DryRunPolicyUpdater, + MemoryFilePolicyUpdater, +) +from openviking.session.train.components.report_builder import ( + PipelineReportBuilder, + PipelineReportHook, +) +from openviking.session.train.components.reporter import ( + ConsolePipelineReporter, + NoopPipelineLifecycleHook, + PipelineLifecycleHook, + emit_run_summary, +) +from openviking.session.train.components.rollout_artifact_recorder import ( + RolloutArtifactEventRecorder, + RolloutArtifactIndex, + RolloutArtifactRecorder, +) +from openviking.session.train.components.rollout_executor import ( + SingleTurnLLMRolloutExecutor, + default_single_turn_prompt, +) +from openviking.session.train.components.session_commit import SessionCommitPolicyTrainer +from openviking.session.train.components.skill_policy_updater import SkillPolicyUpdater +from openviking.session.train.components.snapshotter import ContentHashPolicySnapshotter +from openviking.session.train.components.trajectory_analyzer import ( + TrajectoryAnalyzerContext, + TrajectoryRolloutAnalyzer, +) +from openviking.session.train.context import ( + ExecutionContext, + PipelineContext, + PipelineHookDecision, +) +from openviking.session.train.domain import ( + Case, + CriterionResult, + Experience, + ExperienceSet, + PipelineEpochResult, + PipelineEvaluationResult, + PipelineResult, + PolicyApplyResult, + PolicyPlanItem, + PolicyPlanItemKind, + PolicyStatus, + PolicyUpdatePlan, + Rollout, + RolloutAnalysis, + RolloutTrainingResult, + Rubric, + RubricCriterion, + RubricEvaluation, + ScopedRolloutTrainingResult, + Trajectory, + TrajectoryOutcome, +) +from openviking.session.train.gradients import PatchSemanticGradient +from openviking.session.train.interfaces import ( + CaseLoader, + GradientEstimator, + PolicyOptimizationPipeline, + PolicyOptimizer, + PolicySnapshotter, + PolicyTrainer, + PolicyUpdater, + RolloutAnalyzer, + RolloutEvaluator, + RolloutExecutor, + SemanticGradient, +) +from openviking.session.train.pipeline import OfflinePolicyOptimizationPipeline + +__all__ = [ + "create_dataset_service_app", + "run_batch_train_eval", + "BatchTrainEvalReport", + "BatchTrainEvalConfig", + "RolloutArtifactIndex", + "RolloutArtifactRecorder", + "RolloutArtifactEventRecorder", + "CompositeEventRecorder", + "JsonlEventRecorder", + "JsonlPipelineEventHook", + "make_streaming_policy_trainer_key", + "get_streaming_policy_trainer", + "StreamingPolicyTrainerKey", + "StreamingPolicyTrainerConfig", + "StreamingPolicyTrainer", + "BatchPolicyTrainer", + "ExperienceGradientEstimator", + "ExperienceGradientContext", + "TrajectoryRolloutAnalyzer", + "TrajectoryAnalyzerContext", + "PatchMergePolicyOptimizer", + "PatchMergePolicyOptimizerContext", + "PolicyTrainer", + "PipelineLifecycleHook", + "PipelineReportBuilder", + "PipelineReportHook", + "ConsolePipelineReporter", + "NoopPipelineLifecycleHook", + "emit_run_summary", + "SessionCommitPolicyTrainer", + "ExperienceSetLoader", + "SkillSetLoader", + "SkillPolicyUpdater", + "DryRunPolicyUpdater", + "MemoryFilePolicyUpdater", + "SingleTurnLLMRolloutExecutor", + "default_single_turn_prompt", + "ContentHashPolicySnapshotter", + "PolicyApplyResult", + "Case", + "CaseLoader", + "CriterionResult", + "OfflinePolicyOptimizationPipeline", + "ExecutionContext", + "Experience", + "ExperienceSet", + "GradientEstimator", + "ListCaseLoader", + "TrialCaseLoader", + "make_trial_case_loader", + "PatchSemanticGradient", + "PipelineContext", + "PipelineHookDecision", + "PipelineEvaluationResult", + "PipelineEpochResult", + "PipelineResult", + "PolicyPlanItem", + "PolicyPlanItemKind", + "PolicyOptimizationPipeline", + "PolicyOptimizer", + "PolicySnapshotter", + "PolicyStatus", + "PolicyUpdatePlan", + "PolicyUpdater", + "Rollout", + "RolloutAnalysis", + "RolloutAnalyzer", + "RolloutExecutor", + "RolloutEvaluator", + "RolloutTrainingResult", + "ScopedRolloutTrainingResult", + "Rubric", + "RubricCriterion", + "RubricEvaluation", + "SemanticGradient", + "Trajectory", + "TrajectoryOutcome", +] diff --git a/openviking/session/train/batch_runner.py b/openviking/session/train/batch_runner.py new file mode 100644 index 0000000000..f3c65b6540 --- /dev/null +++ b/openviking/session/train/batch_runner.py @@ -0,0 +1,1118 @@ +"""Generic remote benchmark batch train/eval orchestration.""" + +from __future__ import annotations + +import inspect +import json +import os +import shutil +from dataclasses import dataclass, field +from datetime import datetime +from hashlib import sha256 +from pathlib import Path +from typing import Any + +from openviking.server.config import load_server_config +from openviking.server.identity import AuthMode +from openviking.session.train.components.dataset_service import rollout_from_dict, rollout_to_dict +from openviking.session.train.components.event_recorder import ( + CompositeEventRecorder, + JsonlEventRecorder, + JsonlPipelineEventHook, +) +from openviking.session.train.components.progress import format_label, label_style +from openviking.session.train.components.remote import RemoteCaseLoader, RemoteRolloutExecutor +from openviking.session.train.components.report_builder import PipelineReportBuilder +from openviking.session.train.components.reporter import ( + _accuracy_style, + _style_plain, + emit_run_summary, + fmt_percent, + fmt_percentage_point_abs, +) +from openviking.session.train.components.rollout_artifact_recorder import ( + RolloutArtifactEventRecorder, + RolloutArtifactRecorder, +) +from openviking.session.train.components.session_commit import SessionCommitPolicyTrainer +from openviking.session.train.components.snapshotter import ContentHashPolicySnapshotter +from openviking.session.train.context import ExecutionContext, PipelineContext +from openviking.session.train.domain import ( + Case, + ExperienceSet, + PolicyUpdatePlan, + Rollout, + RolloutAnalysis, +) +from openviking.session.train.pipeline import OfflinePolicyOptimizationPipeline +from openviking.telemetry import tracer +from openviking_cli.client.http import AsyncHTTPClient +from openviking_cli.utils.config.open_viking_config import OpenVikingConfigSingleton + + +@dataclass(slots=True) +class BatchTrainEvalConfig: + """Configuration for one remote benchmark batch train/eval run.""" + + domain: str + dataset: str + epochs: int = 1 + batch_size: int | None = None + concurrency: int = 200 + config_path: str | None = None + output_path: str | None = None + keep_default_tools: bool = True + max_iterations: int = 30 + server_url: str | None = None + api_key: str | None = None + account_id: str = "default" + user_id: str = "default" + commit_keep_recent_count: int = 0 + commit_poll_interval_seconds: float = 2.0 + commit_timeout_seconds: float | None = None + commit_concurrency: int = 200 + train_index: int | str | list[int] | tuple[int, ...] | None = None + eval_index: int | str | list[int] | tuple[int, ...] | None = None + benchmark_service_url: str | None = None + baseline_force_recompute: bool = False + skip_baseline_eval: bool = False + eval_each_epoch: bool = False + eval_split: str | None = "test" + skip_final_eval: bool = False + trials: int = 8 + train_trials: int = 1 + reuse_train_rollout_cache: bool = False + clean_result: bool = True + keep_recent_results: int = 5 + events_path: str | None = None + result_dir_name: str = "train" + run_timestamp: str = field(default_factory=lambda: datetime.now().strftime("%Y%m%d_%H%M%S")) + + def __post_init__(self) -> None: + if not self.dataset: + raise ValueError("dataset is required") + if not self.domain: + raise ValueError("domain is required") + if self.epochs < 0: + raise ValueError("epochs must be >= 0") + if self.batch_size is not None and self.batch_size <= 0: + raise ValueError("batch_size must be > 0") + if self.concurrency <= 0: + raise ValueError("concurrency must be > 0") + if self.max_iterations <= 0: + raise ValueError("max_iterations must be > 0") + if self.commit_poll_interval_seconds <= 0: + raise ValueError("commit_poll_interval_seconds must be > 0") + if self.commit_timeout_seconds is not None and self.commit_timeout_seconds <= 0: + raise ValueError("commit_timeout_seconds must be > 0") + if self.commit_concurrency <= 0: + raise ValueError("commit_concurrency must be > 0") + self.train_index = _normalize_index_filter(self.train_index, label="train_index") + self.eval_index = _normalize_index_filter(self.eval_index, label="eval_index") + if self.eval_split is not None: + normalized_eval_split = str(self.eval_split).strip().lower() + if normalized_eval_split in {"", "none"}: + self.eval_split = None + elif normalized_eval_split not in {"train", "test"}: + raise ValueError("eval_split must be train, test, or none") + else: + self.eval_split = normalized_eval_split + if self.trials <= 0: + raise ValueError("trials must be > 0") + if self.train_trials <= 0: + raise ValueError("train_trials must be > 0") + if self.benchmark_service_url is not None and not self.benchmark_service_url.strip(): + raise ValueError("benchmark_service_url must not be empty") + if self.keep_recent_results < 0: + raise ValueError("keep_recent_results must be >= 0") + if not str(self.result_dir_name or "").strip(): + raise ValueError("result_dir_name must not be empty") + + +def _normalize_index_filter(value: Any, *, label: str) -> list[int] | None: + if value is None: + return None + raw_items: list[Any] + if isinstance(value, str): + raw_items = [item.strip() for item in value.split(",") if item.strip()] + elif isinstance(value, int): + raw_items = [value] + else: + raw_items = [] + try: + iterable = list(value) + except TypeError as exc: + raise ValueError(f"{label} must be an integer or comma-separated integers") from exc + for item in iterable: + if isinstance(item, str) and "," in item: + raw_items.extend(part.strip() for part in item.split(",") if part.strip()) + else: + raw_items.append(item) + result: list[int] = [] + for item in raw_items: + try: + index = int(item) + except (TypeError, ValueError) as exc: + raise ValueError(f"{label} must be an integer or comma-separated integers") from exc + if index < 0: + raise ValueError(f"{label} must be >= 0") + if index not in result: + result.append(index) + if not result: + raise ValueError(f"{label} must not be empty") + return result + + +@dataclass(slots=True) +class BatchTrainEvalReport: + """Serializable report for remote benchmark batch train/eval.""" + + dataset: str + domain: str + epochs: int + batch_size: int | None + concurrency: int + commit_concurrency: int + train_index: int | list[int] | None + eval_index: int | list[int] | None + policy_root_uri: str + baseline_eval: dict[str, Any] | None + train_epochs: list[dict[str, Any]] = field(default_factory=list) + epoch_evals: list[dict[str, Any]] = field(default_factory=list) + final_eval: dict[str, Any] | None = None + accuracy_delta: float | None = None + output_path: str | None = None + trace_id: str | None = None + run_id: str = "" + server_url: str = "" + benchmark_service_url: str | None = None + eval_each_epoch: bool = False + eval_split: str | None = "test" + skip_baseline_eval: bool = False + trials: int = 8 + train_trials: int = 1 + reuse_train_rollout_cache: bool = False + rollouts_root: str | None = None + rollouts_index_path: str | None = None + latest_failed_rollout: str | None = None + clean_result: bool = True + keep_recent_results: int = 5 + events_path: str | None = None + result_dir_name: str = "train" + baseline_cache_path: str | None = None + baseline_cache_hit: bool = False + baseline_force_recompute: bool = False + skip_final_eval: bool = False + final_eval_source: str | None = None + + def to_dict(self) -> dict[str, Any]: + return { + "dataset": self.dataset, + "domain": self.domain, + "epochs": self.epochs, + "batch_size": self.batch_size, + "concurrency": self.concurrency, + "commit_concurrency": self.commit_concurrency, + "train_index": self.train_index, + "eval_index": self.eval_index, + "policy_root_uri": self.policy_root_uri, + "baseline_eval": self.baseline_eval, + "train_epochs": self.train_epochs, + "epoch_evals": self.epoch_evals, + "final_eval": self.final_eval, + "accuracy_delta": self.accuracy_delta, + "output_path": self.output_path, + "trace_id": self.trace_id, + "run_id": self.run_id, + "server_url": self.server_url, + "benchmark_service_url": self.benchmark_service_url, + "eval_each_epoch": self.eval_each_epoch, + "eval_split": self.eval_split, + "skip_baseline_eval": self.skip_baseline_eval, + "trials": self.trials, + "train_trials": self.train_trials, + "reuse_train_rollout_cache": self.reuse_train_rollout_cache, + "rollouts_root": self.rollouts_root, + "rollouts_index_path": self.rollouts_index_path, + "latest_failed_rollout": self.latest_failed_rollout, + "clean_result": self.clean_result, + "keep_recent_results": self.keep_recent_results, + "events_path": self.events_path, + "result_dir_name": self.result_dir_name, + "baseline_cache_path": self.baseline_cache_path, + "baseline_cache_hit": self.baseline_cache_hit, + "baseline_force_recompute": self.baseline_force_recompute, + "skip_final_eval": self.skip_final_eval, + "final_eval_source": self.final_eval_source, + } + + +@tracer("train.batch_train_eval.run", ignore_result=True, ignore_args=True) +async def run_batch_train_eval(config: BatchTrainEvalConfig) -> BatchTrainEvalReport: + """Run baseline eval, commit-based train epochs, and final eval for one dataset/domain.""" + + _configure_openviking_config(config.config_path) + _clean_result_dir(config) + client = _build_http_client(config) + await client.initialize() + try: + policy_root_uri = "viking://user/memories/experiences" + policy_set = ExperienceSet( + root_uri=policy_root_uri, + policies=[], + metadata=_policy_set_metadata(config, client), + ) + run_dir = _run_output_dir(config) + event_recorder = JsonlEventRecorder( + path=_events_path(config), + default_fields={ + "dataset": config.dataset, + "domain": config.domain, + "run_timestamp": config.run_timestamp, + }, + ) + await event_recorder.record( + "run_start", + stage="run_start", + epochs=config.epochs, + concurrency=config.concurrency, + commit_concurrency=config.commit_concurrency, + train_index=_index_payload(config.train_index), + eval_index=_index_payload(config.eval_index), + trials=config.trials, + train_trials=config.train_trials, + reuse_train_rollout_cache=config.reuse_train_rollout_cache, + clean_result=config.clean_result, + keep_recent_results=config.keep_recent_results, + result_dir_name=config.result_dir_name, + baseline_force_recompute=config.baseline_force_recompute, + skip_baseline_eval=config.skip_baseline_eval, + eval_split=config.eval_split, + skip_final_eval=config.skip_final_eval, + baseline_cache_path=( + None + if config.skip_baseline_eval or config.eval_split is None + else str(_baseline_cache_path(config)) + ), + ) + policy_trainer = SessionCommitPolicyTrainer( + client=client, + keep_recent_count=config.commit_keep_recent_count, + poll_interval_seconds=config.commit_poll_interval_seconds, + timeout_seconds=config.commit_timeout_seconds, + commit_concurrency=config.commit_concurrency, + show_progress=True, + progress_label="train", + ) + event_recorder.default_fields["run_id"] = policy_trainer.run_id + pipeline = _build_pipeline(config, policy_trainer) + rollout_artifact_recorder = RolloutArtifactRecorder( + run_dir=run_dir, + client=client, + latest_pointer_path=_latest_rollouts_path(config), + ) + remote_executor = getattr(pipeline, "rollout_executor", None) + if isinstance(remote_executor, (RemoteRolloutExecutor, CachedEpochZeroTrainRolloutExecutor)): + remote_executor.on_rollout_complete = ( + rollout_artifact_recorder.record_rollout_completion + ) + policy_trainer.event_recorder = CompositeEventRecorder( + (event_recorder, RolloutArtifactEventRecorder(rollout_artifact_recorder)) + ) + + baseline_eval: dict[str, Any] | None = None + baseline_cache_hit = False + baseline_cache_path = ( + None + if config.skip_baseline_eval or config.eval_split is None + else _baseline_cache_path(config) + ) + final_eval: dict[str, Any] | None = None + report_builder = PipelineReportBuilder(trial_index_key="eval_trial") + + eval_loader = ( + None + if config.eval_split is None + else _case_loader( + config, + split=config.eval_split, + sample_index=config.eval_index, + ) + ) + if ( + eval_loader is not None + and not config.skip_baseline_eval + and await eval_loader.split_exists() + ): + baseline_result, baseline_cache_hit = await _load_or_run_baseline_eval( + config=config, + pipeline=pipeline, + case_loader=eval_loader, + policy_set=policy_set, + report_builder=report_builder, + event_recorder=event_recorder, + ) + if baseline_result is not None: + rollout_artifact_recorder.record_eval( + label=_eval_rollout_stage("baseline", config.eval_split), + epoch=-1, + analyses=baseline_result.analyses, + ) + baseline_eval = baseline_result.metadata["report"] + else: + assert baseline_cache_path is not None + baseline_eval = _load_baseline_cache(baseline_cache_path) + if baseline_eval is not None: + _print_baseline_cache_hit(baseline_eval, baseline_cache_path) + + train_loader = _case_loader( + config, + split="train", + sample_index=config.train_index, + ) + + train_context = _pipeline_context( + epoch=0, + training=True, + max_epochs=config.epochs, + eval_each_epoch_case_loader=( + eval_loader + if ( + eval_loader is not None + and config.eval_each_epoch + and await eval_loader.split_exists() + ) + else None + ), + rollout_stage=( + _eval_rollout_stage("epoch", config.eval_split) + if config.eval_split is not None + else None + ), + eval_split=config.eval_split, + eval_trials=config.trials, + train_trials=config.train_trials, + trial_index_key="eval_trial", + report_builder=report_builder, + event_recorder=event_recorder, + ) + # Register rollout artifact recorder as a lifecycle hook so rollouts + # are written incrementally after each epoch/eval, instead of waiting + # for the full run to finish. + train_context.lifecycle_hooks = list(train_context.lifecycle_hooks) + [ + rollout_artifact_recorder + ] + train_result = await pipeline.train( + case_loader=train_loader, + policy_set=policy_set, + context=train_context, + ) + policy_set = train_result.apply_result.updated_policy_set + # Note: per-epoch rollout artifacts are written incrementally via the + # rollout_artifact_recorder lifecycle hook registered on train_context. + + epoch_eval_reports = _epoch_eval_reports(train_result) + final_eval_source: str | None = None + if config.skip_final_eval: + if epoch_eval_reports: + final_eval = dict(epoch_eval_reports[-1]) + final_eval_source = "last_epoch_eval" + elif eval_loader is not None and await eval_loader.split_exists(): + final_result = await pipeline.eval( + case_loader=eval_loader, + policy_set=policy_set, + context=_pipeline_context( + epoch=config.epochs, + training=False, + max_epochs=1, + rollout_stage=_eval_rollout_stage("final", config.eval_split), + eval_split=config.eval_split, + eval_trials=config.trials, + trial_index_key="eval_trial", + report_builder=report_builder, + event_recorder=event_recorder, + ), + ) + rollout_artifact_recorder.record_eval( + label=_eval_rollout_stage("final", config.eval_split), + epoch=config.epochs, + analyses=final_result.analyses, + ) + final_eval = final_result.metadata["report"] + final_eval_source = f"final_{config.eval_split}" + + accuracy_delta = report_builder.accuracy_delta(baseline_eval, final_eval) + rollout_artifact_index = rollout_artifact_recorder.finalize() + report = BatchTrainEvalReport( + dataset=config.dataset, + domain=config.domain, + epochs=config.epochs, + batch_size=config.batch_size, + concurrency=config.concurrency, + commit_concurrency=config.commit_concurrency, + train_index=_index_payload(config.train_index), + eval_index=_index_payload(config.eval_index), + policy_root_uri=policy_root_uri, + baseline_eval=baseline_eval, + train_epochs=list(train_result.metadata.get("train_reports", [])), + epoch_evals=epoch_eval_reports, + final_eval=final_eval, + accuracy_delta=accuracy_delta, + output_path=_default_output_path(config), + trace_id=tracer.get_trace_id() or None, + run_id=policy_trainer.run_id, + server_url=client_url(client), + benchmark_service_url=config.benchmark_service_url, + eval_each_epoch=config.eval_each_epoch, + eval_split=config.eval_split, + skip_baseline_eval=config.skip_baseline_eval, + trials=config.trials, + train_trials=config.train_trials, + reuse_train_rollout_cache=config.reuse_train_rollout_cache, + rollouts_root=rollout_artifact_index.rollouts_root, + rollouts_index_path=str(run_dir / "rollouts_index.json"), + latest_failed_rollout=rollout_artifact_index.latest_failed_rollout, + clean_result=config.clean_result, + keep_recent_results=config.keep_recent_results, + events_path=str(_events_path(config)), + result_dir_name=config.result_dir_name, + baseline_cache_path=( + str(baseline_cache_path) if baseline_cache_path is not None else None + ), + baseline_cache_hit=baseline_cache_hit, + baseline_force_recompute=config.baseline_force_recompute, + skip_final_eval=config.skip_final_eval, + final_eval_source=final_eval_source, + ) + _write_report(report, config) + await event_recorder.record( + "run_result", + stage="run_result", + trace_id=report.trace_id, + output_path=report.output_path, + rollouts_root=report.rollouts_root, + rollouts_index_path=report.rollouts_index_path, + latest_failed_rollout=report.latest_failed_rollout, + accuracy_delta=report.accuracy_delta, + baseline_cache_path=report.baseline_cache_path, + baseline_cache_hit=report.baseline_cache_hit, + ) + await emit_run_summary( + train_context, + title="batch train/eval", + fields={ + "dataset": config.dataset, + "domain": config.domain, + "epochs": config.epochs, + "trials": config.trials, + "train_trials": config.train_trials, + "reuse_train_rollout_cache": config.reuse_train_rollout_cache, + "run_id": policy_trainer.run_id, + "trace_id": report.trace_id, + "baseline_cache_hit": report.baseline_cache_hit, + "skip_baseline_eval": report.skip_baseline_eval, + "eval_split": report.eval_split, + "skip_final_eval": report.skip_final_eval, + "final_eval_source": report.final_eval_source, + }, + baseline_eval=baseline_eval, + final_eval=final_eval, + accuracy_delta=accuracy_delta, + output_path=report.output_path, + rollouts_root=report.rollouts_root, + rollouts_index_path=report.rollouts_index_path, + latest_failed_rollout=report.latest_failed_rollout, + ) + return report + finally: + await client.close() + + +def _configure_openviking_config(config_path: str | None) -> None: + if config_path: + os.environ["OPENVIKING_CONFIG_FILE"] = str(Path(config_path).expanduser()) + OpenVikingConfigSingleton.reset_instance() + + +def _build_http_client(config: BatchTrainEvalConfig) -> AsyncHTTPClient: + server_url = config.server_url + api_key = config.api_key + auth_mode: AuthMode | None = None + if config.config_path or server_url is None or api_key is None: + server_config = load_server_config(config.config_path) + auth_mode = server_config.get_effective_auth_mode() + server_url = server_url or f"http://{server_config.host}:{server_config.port}" + api_key = api_key or server_config.root_api_key + if auth_mode is None: + auth_mode = AuthMode.API_KEY + + # Trusted mode uses X-API-Key as the gateway/root key and takes identity from + # X-OpenViking-Account/User. In api_key/dev modes, user API keys already pin + # identity and account/user assertion headers must not be sent. + account = config.account_id if auth_mode == AuthMode.TRUSTED else None + user = config.user_id if auth_mode == AuthMode.TRUSTED else None + return AsyncHTTPClient( + url=server_url, + api_key=api_key, + account=account, + user=user, + profile_enabled=False, + timeout=max(60.0, (config.commit_timeout_seconds or 600.0) + 30.0), + ) + + +def _epoch_eval_reports(train_result: Any) -> list[dict[str, Any]]: + reports: list[dict[str, Any]] = [] + for evaluation in getattr(train_result, "evaluation_passes", []) or []: + report = getattr(evaluation, "metadata", {}).get("report") + if isinstance(report, dict): + reports.append(dict(report)) + return reports + + +def _policy_set_metadata(config: BatchTrainEvalConfig, client: AsyncHTTPClient) -> dict[str, Any]: + return { + "source": "remote_session_commit", + "openviking_url": client_url(client), + "openviking_api_key": getattr(client, "_api_key", None), + "openviking_account": config.account_id, + "openviking_user": config.user_id, + } + + +def client_url(client: AsyncHTTPClient) -> str: + return str(getattr(client, "_url", "")) + + +async def _load_or_run_baseline_eval( + *, + config: BatchTrainEvalConfig, + pipeline: OfflinePolicyOptimizationPipeline, + case_loader: RemoteCaseLoader, + policy_set: ExperienceSet, + report_builder: PipelineReportBuilder, + event_recorder: JsonlEventRecorder, +) -> tuple[Any | None, bool]: + cache_path = _baseline_cache_path(config) + if not config.baseline_force_recompute: + cached_report = _load_baseline_cache(cache_path) + if cached_report is not None: + await event_recorder.record( + "baseline_cache_hit", + stage="baseline_cache", + baseline_cache_path=str(cache_path), + eval_split=config.eval_split, + ) + return None, True + + await event_recorder.record( + "baseline_cache_recompute" if cache_path.exists() else "baseline_cache_miss", + stage="baseline_cache", + baseline_cache_path=str(cache_path), + eval_split=config.eval_split, + ) + baseline_result = await pipeline.eval( + case_loader=case_loader, + policy_set=policy_set, + context=_pipeline_context( + epoch=-1, + training=False, + max_epochs=1, + rollout_stage=_eval_rollout_stage("baseline", config.eval_split), + eval_split=config.eval_split, + eval_trials=config.trials, + trial_index_key="eval_trial", + report_builder=report_builder, + event_recorder=event_recorder, + ), + ) + _write_baseline_cache(cache_path, baseline_result.metadata["report"], config=config) + await event_recorder.record( + "baseline_cache_write", + stage="baseline_cache", + baseline_cache_path=str(cache_path), + eval_split=config.eval_split, + ) + return baseline_result, False + + +def _write_baseline_cache( + path: Path, + report: dict[str, Any], + *, + config: BatchTrainEvalConfig, +) -> None: + payload = { + "cache_version": 1, + "cache_key": _baseline_cache_key(config), + "dataset": config.dataset, + "domain": config.domain, + "split": config.eval_split, + "eval_index": _index_payload(config.eval_index), + "trials": config.trials, + "max_iterations": config.max_iterations, + "keep_default_tools": config.keep_default_tools, + "created_at": datetime.now().isoformat(), + "report": report, + } + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8") + + +def _load_baseline_cache(path: Path) -> dict[str, Any] | None: + if not path.exists(): + return None + payload = json.loads(path.read_text(encoding="utf-8")) + if payload.get("cache_version") != 1: + raise ValueError(f"unsupported baseline cache version in {path}") + report = payload.get("report") + if not isinstance(report, dict): + raise ValueError(f"baseline cache file has no report: {path}") + return { + **report, + "baseline_cache_hit": True, + "baseline_cache_path": str(path), + } + + +def _print_baseline_cache_hit(report: dict[str, Any], cache_path: Path) -> None: + """Print cached baseline info before training starts so users see it immediately.""" + label = str(report.get("rollout_stage") or "baseline_test_rollout") + trial_count = int(report.get("trial_count") or 1) + cache_info = f"(from cache: {cache_path.name})" + label_text = _style_plain(format_label(label), label_style(label)) + if trial_count > 1: + accuracy_mean = report.get("accuracy_mean") + accuracy_std = report.get("accuracy_std") + cases_per_trial = report.get("case_count_per_trial") or "varies" + print( + f"{label_text} baseline_cache_hit=1 " + f"accuracy={_style_plain(fmt_percent(accuracy_mean), _accuracy_style(accuracy_mean))} " + f"± {_style_plain(fmt_percentage_point_abs(accuracy_std), 'yellow')} " + f"trials={trial_count} cases_per_trial={cases_per_trial} " + f"{cache_info}" + ) + return + accuracy = report.get("accuracy") + passed = report.get("passed_count") + total = report.get("case_count") + print( + f"{label_text} baseline_cache_hit=1 " + f"accuracy={_style_plain(fmt_percent(accuracy), _accuracy_style(accuracy))} " + f"passed={passed}/{total} " + f"{cache_info}" + ) + + +def _eval_rollout_stage(kind: str, split: str | None) -> str: + eval_split = str(split or "test") + if kind == "epoch" and eval_split == "train": + return "eval_train_rollout" + return f"{kind}_{eval_split}_rollout" + + +def _build_pipeline( + config: BatchTrainEvalConfig, + policy_trainer: SessionCommitPolicyTrainer, +) -> OfflinePolicyOptimizationPipeline: + rollout_executor: Any = RemoteRolloutExecutor( + service_url=_require_benchmark_service_url(config), + concurrency=config.concurrency, + show_progress=True, + progress_label="rollout", + options={ + "config_path": config.config_path, + "keep_default_tools": config.keep_default_tools, + "max_iterations": config.max_iterations, + }, + ) + if config.reuse_train_rollout_cache: + rollout_executor = CachedEpochZeroTrainRolloutExecutor( + delegate=rollout_executor, + cache_dir=_train_rollout_cache_dir(config), + cache_key_prefix=_train_rollout_cache_key_prefix(config), + ) + return OfflinePolicyOptimizationPipeline( + snapshotter=ContentHashPolicySnapshotter(prefix=f"{config.dataset}-policy-snapshot"), + rollout_executor=rollout_executor, + rollout_analyzer=UnusedRolloutAnalyzer(), + gradient_estimator=UnusedGradientEstimator(), + policy_optimizer=UnusedPolicyOptimizer(), + policy_updater=UnusedPolicyUpdater(), + policy_trainer=policy_trainer, + ) + + +def _pipeline_context( + *, + epoch: int, + training: bool, + max_epochs: int = 1, + rollout_stage: str | None = None, + eval_split: str | None = None, + eval_each_epoch_case_loader: Any = None, + eval_trials: int = 1, + train_trials: int = 1, + trial_index_key: str = "trial", + report_builder: Any = None, + event_recorder: JsonlEventRecorder | None = None, +) -> PipelineContext: + execution_metadata = {"epoch": epoch, "training": training} + if rollout_stage is not None: + execution_metadata["rollout_stage"] = rollout_stage + if eval_split is not None: + execution_metadata["eval_split"] = eval_split + hooks = None + if event_recorder is not None: + from openviking.session.train.components.report_builder import PipelineReportHook + from openviking.session.train.components.reporter import ConsolePipelineReporter + + hooks = [ + PipelineReportHook(), + JsonlPipelineEventHook(event_recorder), + ConsolePipelineReporter(), + ] + return PipelineContext( + analysis_context={"epoch": epoch}, + execution_metadata=execution_metadata, + max_epochs=max_epochs, + eval_each_epoch_case_loader=eval_each_epoch_case_loader, + eval_trials=eval_trials, + train_trials=train_trials, + trial_index_key=trial_index_key, + report_builder=report_builder, + **({"lifecycle_hooks": hooks} if hooks is not None else {}), + ) + + +def _case_loader( + config: BatchTrainEvalConfig, + *, + split: str, + sample_index: list[int] | None, +) -> RemoteCaseLoader: + filters: dict[str, Any] = {} + if sample_index is not None: + filters["task_indices"] = list(sample_index) + return RemoteCaseLoader( + service_url=_require_benchmark_service_url(config), + dataset=config.dataset, + domain=config.domain, + split=split, + batch_size=config.batch_size, + filters=filters, + ) + + +@dataclass(slots=True) +class CachedEpochZeroTrainRolloutExecutor: + """Reuse cached train rollouts only for epoch 0. + + The first training epoch is the no-new-memory rollout used to generate the + initial training signal. Later epochs intentionally execute again so they + can observe policy/memory updates from earlier commits. + """ + + delegate: Any + cache_dir: Path + cache_key_prefix: str + on_rollout_complete: Any | None = None + + async def execute( + self, + cases: list[Case], + policy_set: ExperienceSet, + context: ExecutionContext, + ) -> list[Rollout]: + metadata = dict(context.metadata or {}) + if not bool(metadata.get("training")) or int(metadata.get("epoch", 0) or 0) != 0: + return await self.delegate.execute(cases, policy_set, context) + + case_list = list(cases) + results: list[Rollout | None] = [None] * len(case_list) + misses: list[tuple[int, Case]] = [] + for index, case in enumerate(case_list): + cached = self._load(case) + if cached is None: + misses.append((index, case)) + else: + cached.policy_snapshot_id = context.policy_snapshot_id + metadata = dict(cached.metadata or {}) + metadata["train_rollout_cache_hit"] = True + metadata["train_rollout_cache_path"] = str(self._path(case)) + cached.metadata = metadata + results[index] = cached + await self._emit_rollout_complete( + rollout=cached, + index=index, + context=context, + ) + + if misses: + miss_rollouts = await self.delegate.execute( + [case for _, case in misses], + policy_set, + context, + ) + for (index, case), rollout in zip(misses, miss_rollouts, strict=True): + self._write(case, rollout) + results[index] = rollout + await self._emit_rollout_complete( + rollout=rollout, + index=index, + context=context, + ) + + return [rollout for rollout in results if rollout is not None] + + def _path(self, case: Case) -> Path: + payload = { + "prefix": self.cache_key_prefix, + "case": { + "name": case.name, + "task_signature": case.task_signature, + "input": case.input, + "metadata": case.metadata, + }, + } + stable = json.dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=False) + digest = sha256(stable.encode("utf-8")).hexdigest()[:24] + return self.cache_dir / f"{_cache_slug(case.name)}_{digest}.json" + + def _load(self, case: Case) -> Rollout | None: + path = self._path(case) + if not path.exists(): + return None + payload = json.loads(path.read_text(encoding="utf-8")) + if payload.get("cache_version") != 1: + raise ValueError(f"unsupported train rollout cache version in {path}") + rollout_data = payload.get("rollout") + if not isinstance(rollout_data, dict): + raise ValueError(f"train rollout cache file has no rollout: {path}") + return rollout_from_dict(rollout_data) + + def _write(self, case: Case, rollout: Rollout) -> None: + path = self._path(case) + payload = { + "cache_version": 1, + "cache_key_prefix": self.cache_key_prefix, + "case_name": case.name, + "task_signature": case.task_signature, + "created_at": datetime.now().isoformat(), + "rollout": rollout_to_dict(rollout), + } + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8") + + async def _emit_rollout_complete( + self, + *, + rollout: Rollout, + index: int, + context: ExecutionContext, + ) -> None: + if self.on_rollout_complete is None: + return + result = self.on_rollout_complete( + rollout=rollout, + index=index, + context=context, + ) + if inspect.isawaitable(result): + await result + + +def _require_benchmark_service_url(config: BatchTrainEvalConfig) -> str: + if not config.benchmark_service_url: + raise ValueError( + "benchmark_service_url is required; start benchmark service and pass " + "--benchmark-service-url" + ) + return config.benchmark_service_url + + +class UnusedRolloutAnalyzer: + async def analyze(self, rollout: Rollout, context: Any = None) -> RolloutAnalysis: + raise RuntimeError("eval uses rollout.evaluation; training is handled by policy_trainer") + + +class UnusedGradientEstimator: + async def estimate( + self, analysis: RolloutAnalysis, experience_set: ExperienceSet, context: Any + ): + raise RuntimeError("policy_trainer handles training; gradient estimator must not run") + + +class UnusedPolicyOptimizer: + async def plan(self, gradients: list[Any], policy_set: ExperienceSet, context: Any): + raise RuntimeError("policy_trainer handles training; policy optimizer must not run") + + +class UnusedPolicyUpdater: + async def apply( + self, + plan: PolicyUpdatePlan, + policy_set: ExperienceSet, + context: Any, + *, + transaction_handle: Any = None, + ): + del transaction_handle + raise RuntimeError("policy_trainer handles training; policy updater must not run") + + +def _write_report(report: BatchTrainEvalReport, config: BatchTrainEvalConfig) -> None: + output_path = Path(_default_output_path(config)).expanduser() + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text( + json.dumps(report.to_dict(), ensure_ascii=False, indent=2), + encoding="utf-8", + ) + report.output_path = str(output_path) + + +def _events_path(config: BatchTrainEvalConfig) -> Path: + if config.events_path: + return Path(config.events_path).expanduser() + return _run_output_dir(config) / "events.jsonl" + + +def _default_output_path(config: BatchTrainEvalConfig) -> str: + if config.output_path: + return str(Path(config.output_path).expanduser()) + return str(_run_output_dir(config) / "report.json") + + +def _baseline_cache_path(config: BatchTrainEvalConfig) -> Path: + return _result_base_dir(config) / "cache" / "baseline" / f"{_baseline_cache_key(config)}.json" + + +def _train_rollout_cache_dir(config: BatchTrainEvalConfig) -> Path: + return ( + _result_base_dir(config) + / "cache" + / "train_rollouts" + / _train_rollout_cache_key_prefix(config) + ) + + +def _train_rollout_cache_key_prefix(config: BatchTrainEvalConfig) -> str: + payload = { + "dataset": config.dataset, + "domain": config.domain, + "split": "train", + "train_index": _index_payload(config.train_index), + "train_trials": config.train_trials, + "max_iterations": config.max_iterations, + "keep_default_tools": config.keep_default_tools, + } + stable = json.dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=False) + digest = sha256(stable.encode("utf-8")).hexdigest()[:16] + index = _index_label(config.train_index) + return f"{_cache_slug(config.domain)}_train_index-{index}_trials-{config.train_trials}_{digest}" + + +def _baseline_cache_key(config: BatchTrainEvalConfig) -> str: + payload = { + "dataset": config.dataset, + "domain": config.domain, + "split": config.eval_split, + "eval_index": _index_payload(config.eval_index), + "trials": config.trials, + "max_iterations": config.max_iterations, + "keep_default_tools": config.keep_default_tools, + } + stable = json.dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=False) + digest = sha256(stable.encode("utf-8")).hexdigest()[:16] + index = _index_label(config.eval_index) + split = _cache_slug(str(config.eval_split or "none")) + return f"{_cache_slug(config.domain)}_{split}_index-{index}_trials-{config.trials}_{digest}" + + +def _index_payload(indices: list[int] | None) -> int | list[int] | None: + if indices is None: + return None + return indices[0] if len(indices) == 1 else list(indices) + + +def _index_label(indices: list[int] | None) -> str: + if indices is None: + return "all" + if len(indices) == 1: + return str(indices[0]) + return "multi-" + "-".join(str(item) for item in indices) + + +def _cache_slug(value: str) -> str: + return ( + "".join(ch if ch.isalnum() or ch in ("-", "_") else "-" for ch in value).strip("-") + or "default" + ) + + +def _clean_result_dir(config: BatchTrainEvalConfig) -> None: + if not config.clean_result: + return + if config.output_path: + output_path = Path(config.output_path).expanduser() + if output_path.exists() and output_path.is_file(): + output_path.unlink() + return + + result_dir = _result_base_dir(config) + result_dir.mkdir(parents=True, exist_ok=True) + + protected_names = {_run_dir_name(config)} + run_dirs = [] + removed = 0 + if result_dir.exists(): + for child in result_dir.iterdir(): + if child.name in protected_names: + continue + if child.is_dir() and not child.is_symlink(): + if _is_default_run_dir(child, config): + run_dirs.append(child) + continue + + run_dirs.sort(key=lambda path: (path.stat().st_mtime, path.name), reverse=True) + keep_count = config.keep_recent_results + for stale_dir in run_dirs[keep_count:]: + shutil.rmtree(stale_dir) + removed += 1 + + print( + f"[batch-train-eval] clean_result=1 path={result_dir} " + f"keep_recent_results={keep_count} removed={removed}", + flush=True, + ) + + +def _is_default_run_dir(path: Path, config: BatchTrainEvalConfig) -> bool: + prefix = f"run_{config.domain}_" + if not path.name.startswith(prefix): + return False + suffix = path.name[len(prefix) :] + return len(suffix) == 15 and suffix[8] == "_" and suffix[:8].isdigit() and suffix[9:].isdigit() + + +def _run_dir_name(config: BatchTrainEvalConfig) -> str: + return f"run_{config.domain}_{config.run_timestamp}" + + +def _run_output_dir(config: BatchTrainEvalConfig) -> Path: + if config.output_path: + output_path = Path(config.output_path).expanduser() + return output_path.parent + return _result_base_dir(config) / _run_dir_name(config) + + +def _result_base_dir(config: BatchTrainEvalConfig) -> Path: + return _repo_root() / "result" / config.dataset / config.result_dir_name + + +def _latest_rollouts_path(config: BatchTrainEvalConfig) -> Path: + return _repo_root() / "result" / config.dataset / config.result_dir_name / "latest_rollouts" + + +def _repo_root() -> Path: + return Path(__file__).resolve().parents[3] diff --git a/openviking/session/train/components/__init__.py b/openviking/session/train/components/__init__.py new file mode 100644 index 0000000000..0cfc6bd2c0 --- /dev/null +++ b/openviking/session/train/components/__init__.py @@ -0,0 +1,65 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Default replaceable components for the session train framework.""" + +from openviking.session.train.components.case_loader import ListCaseLoader +from openviking.session.train.components.gradient_estimator import ( + ExperienceGradientContext, + ExperienceGradientEstimator, +) +from openviking.session.train.components.memory_store import ExperienceSetLoader, SkillSetLoader +from openviking.session.train.components.policy_optimizer import ( + PatchMergePolicyOptimizer, + PatchMergePolicyOptimizerContext, +) +from openviking.session.train.components.policy_trainer import ( + BatchPolicyTrainer, + StreamingPolicyTrainer, + StreamingPolicyTrainerConfig, + StreamingPolicyTrainerKey, + get_streaming_policy_trainer, + make_streaming_policy_trainer_key, +) +from openviking.session.train.components.policy_updater import ( + DryRunPolicyUpdater, + MemoryFilePolicyUpdater, +) +from openviking.session.train.components.remote import RemoteCaseLoader, RemoteRolloutExecutor +from openviking.session.train.components.rollout_executor import ( + SingleTurnLLMRolloutExecutor, + default_single_turn_prompt, +) +from openviking.session.train.components.session_commit import SessionCommitPolicyTrainer +from openviking.session.train.components.skill_policy_updater import SkillPolicyUpdater +from openviking.session.train.components.snapshotter import ContentHashPolicySnapshotter +from openviking.session.train.components.trajectory_analyzer import ( + TrajectoryAnalyzerContext, + TrajectoryRolloutAnalyzer, +) + +__all__ = [ + "ContentHashPolicySnapshotter", + "make_streaming_policy_trainer_key", + "get_streaming_policy_trainer", + "StreamingPolicyTrainerKey", + "StreamingPolicyTrainerConfig", + "StreamingPolicyTrainer", + "BatchPolicyTrainer", + "PatchMergePolicyOptimizerContext", + "PatchMergePolicyOptimizer", + "ListCaseLoader", + "ExperienceGradientEstimator", + "ExperienceGradientContext", + "TrajectoryRolloutAnalyzer", + "TrajectoryAnalyzerContext", + "DryRunPolicyUpdater", + "MemoryFilePolicyUpdater", + "SingleTurnLLMRolloutExecutor", + "default_single_turn_prompt", + "ExperienceSetLoader", + "SkillSetLoader", + "SkillPolicyUpdater", + "SessionCommitPolicyTrainer", + "RemoteRolloutExecutor", + "RemoteCaseLoader", +] diff --git a/openviking/session/train/components/case_loader.py b/openviking/session/train/components/case_loader.py new file mode 100644 index 0000000000..10e5f3978f --- /dev/null +++ b/openviking/session/train/components/case_loader.py @@ -0,0 +1,105 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Case loader implementations for session training.""" + +from __future__ import annotations + +from collections.abc import AsyncIterator +from dataclasses import dataclass, field +from typing import Any + +from openviking.session.train.domain import Case +from openviking.session.train.interfaces import CaseLoader +from openviking.telemetry import tracer + + +@dataclass(slots=True) +class ListCaseLoader: + """Simple in-memory CaseLoader implementation.""" + + cases: list[Case] + batch_size: int | None = None + + @tracer("train.case_loader.list.batches", ignore_result=True, ignore_args=True) + async def batches(self, context: Any) -> AsyncIterator[list[Case]]: + del context + batch_size = self.batch_size or len(self.cases) or 1 + for start in range(0, len(self.cases), batch_size): + yield list(self.cases[start : start + batch_size]) + + +@dataclass(slots=True) +class TrialCaseLoader: + """Expand every base case into N trial cases in each emitted batch.""" + + base_loader: CaseLoader + trial_count: int + trial_input_key: str = "trial" + trial_count_input_key: str = "trial_count" + original_case_name_input_key: str = "original_case_name" + trial_name_template: str = "{case_name}_t{trial_index}" + trial_task_signature_template: str = "{task_signature}:trial:{trial_index}" + extra_input: dict[str, Any] = field(default_factory=dict) + extra_metadata: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.trial_count <= 0: + raise ValueError("trial_count must be > 0") + + async def batches(self, context: Any) -> AsyncIterator[list[Case]]: + async for cases in self.base_loader.batches(context): + expanded: list[Case] = [] + for trial_index in range(self.trial_count): + expanded.extend(self._trial_case(case, trial_index) for case in cases) + yield expanded + + async def split_exists(self) -> bool: + split_exists = getattr(self.base_loader, "split_exists", None) + if split_exists is None: + return True + return bool(await split_exists()) + + def _trial_case(self, case: Case, trial_index: int) -> Case: + trial_values = { + self.trial_input_key: trial_index, + self.trial_count_input_key: self.trial_count, + self.original_case_name_input_key: case.name, + } + format_values = { + "case_name": case.name, + "task_signature": case.task_signature, + "trial_index": trial_index, + "trial_count": self.trial_count, + } + return Case( + name=self.trial_name_template.format(**format_values), + task_signature=self.trial_task_signature_template.format(**format_values), + input={ + **dict(case.input), + **trial_values, + **dict(self.extra_input), + }, + rubric=case.rubric, + metadata={ + **dict(case.metadata), + **trial_values, + **dict(self.extra_metadata), + }, + ) + + +def make_trial_case_loader( + base_loader: CaseLoader, + trial_count: int, + *, + trial_input_key: str = "trial", + trial_count_input_key: str | None = None, + original_case_name_input_key: str = "original_case_name", +) -> TrialCaseLoader: + return TrialCaseLoader( + base_loader=base_loader, + trial_count=trial_count, + trial_input_key=trial_input_key, + trial_count_input_key=trial_count_input_key or f"{trial_input_key}_count", + original_case_name_input_key=original_case_name_input_key, + ) diff --git a/openviking/session/train/components/dataset_service.py b/openviking/session/train/components/dataset_service.py new file mode 100644 index 0000000000..0016587ddb --- /dev/null +++ b/openviking/session/train/components/dataset_service.py @@ -0,0 +1,703 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Generic HTTP service host for remote benchmark datasets.""" + +from __future__ import annotations + +import asyncio +import atexit +import json +import logging +import re +import shutil +import tempfile +import threading +import time +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from typing import Any +from uuid import uuid4 + +from fastapi import FastAPI, HTTPException +from pydantic import BaseModel, Field + +from openviking.session.train.context import ExecutionContext +from openviking.session.train.domain import ( + Case, + CriterionResult, + Experience, + ExperienceSet, + Rollout, + Rubric, + RubricCriterion, + RubricEvaluation, +) + +CaseLoaderFactory = Callable[[str, str, str, dict[str, Any]], Any] +RolloutExecutorFactory = Callable[[dict[str, Any]], Any] +logger = logging.getLogger(__name__) +_rollout_worker_state = threading.local() + +SENSITIVE_FIELD_NAMES = frozenset( + { + "api_key", + "openviking_api_key", + "root_api_key", + "authorization", + "x-api-key", + "token", + "secret", + } +) +_SENSITIVE_TEXT_RE = re.compile( + r"(?i)(openviking_api_key|root_api_key|api_key|x-api-key|authorization|token|secret)" + r"(\s*[:=]\s*)" + r"(['\"]?)([^'\"\s,}]+)(\3)" +) + + +def _redact_sensitive_text(text: str) -> str: + def _replace(match: re.Match[str]) -> str: + return ( + f"{match.group(1)}{match.group(2)}" + f"{match.group(3)}{match.group(5)}" + ) + + return _SENSITIVE_TEXT_RE.sub(_replace, text) + + +def redact_sensitive(value: Any) -> Any: + """Return a copy with known credential fields redacted for logs/errors.""" + if isinstance(value, str): + return _redact_sensitive_text(value) + if isinstance(value, dict): + redacted: dict[Any, Any] = {} + for key, item in value.items(): + key_text = str(key).lower() + if key_text in SENSITIVE_FIELD_NAMES or key_text.endswith("_api_key"): + redacted[key] = "" + else: + redacted[key] = redact_sensitive(item) + return redacted + if isinstance(value, list): + return [redact_sensitive(item) for item in value] + if isinstance(value, tuple): + return tuple(redact_sensitive(item) for item in value) + return value + + +class CasesQueryRequest(BaseModel): + dataset: str + domain: str + split: str + cursor: str | None = None + limit: int = Field(default=100, gt=0) + filters: dict[str, Any] = Field(default_factory=dict) + + +class RolloutExecuteRequest(BaseModel): + case: dict[str, Any] + policy_set: dict[str, Any] + execution_context: dict[str, Any] + options: dict[str, Any] = Field(default_factory=dict) + + +@dataclass(slots=True) +class RolloutExecution: + execution_id: str + status: str + created_at: float + updated_at: float + case_name: str + finished_at: float | None = None + # ``rollout`` is only populated on copies returned by ``get()`` for + # completed executions. Filesystem-backed records never keep Rollout + # objects alive across requests, so large payloads (messages, tool outputs, + # prompts, reasoning) are reloaded from disk on each poll and released as + # soon as the caller returns. + rollout: Rollout | None = None + error: str | None = None + + +class RolloutExecutionStore: + """Filesystem-backed execution store with zero in-memory state. + + Execution state is represented entirely by files under a root spool + directory. The root directory is created under ``tempfile.gettempdir()`` by + default so it is automatically cleared on service restart / host reboot; + there is no TTL or eviction because no payloads are kept in memory beyond + the scope of a single request. + + Directory layout:: + + / + running/.json # {status, created_at, updated_at, case_name} + completed/.json # {status, created_at, updated_at, finished_at, case_name, rollout: {...}} + failed/.json # {status, created_at, updated_at, finished_at, case_name, error} + """ + + _RUNNING = "running" + _COMPLETED = "completed" + _FAILED = "failed" + + def __init__(self, *, spool_dir: Path | None = None) -> None: + if spool_dir is None: + # Per-process temporary directory. Each service process gets its own + # spool root; the directory is removed on process exit via atexit so + # disk usage does not grow across runs. Concurrent service instances + # use separate directories (different mkdtemp suffixes) and do not + # collide. + spool_dir = Path(tempfile.mkdtemp(prefix="ov_rollout_spool_")) + self._owns_spool_dir = True + else: + spool_dir = spool_dir.expanduser().resolve() + self._owns_spool_dir = False + spool_dir = spool_dir.expanduser().resolve() + self._spool_dir = spool_dir + for sub in (self._RUNNING, self._COMPLETED, self._FAILED): + (spool_dir / sub).mkdir(parents=True, exist_ok=True) + if self._owns_spool_dir: + atexit.register(_cleanup_spool_dir, self._spool_dir) + + def _path(self, status: str, execution_id: str) -> Path: + return self._spool_dir / status / f"{execution_id}.json" + + @staticmethod + def _read_meta(path: Path) -> dict[str, Any] | None: + if not path.exists(): + return None + try: + with path.open("r", encoding="utf-8") as fh: + return json.load(fh) + except (OSError, json.JSONDecodeError): + return None + + async def create(self, *, case_name: str) -> RolloutExecution: + # File creation is the single source of truth. Use a lock-free + # atomic-write pattern; POSIX rename is atomic for same-directory files. + now = time.time() + execution_id = f"rollout_exec_{uuid4().hex}" + meta = { + "execution_id": execution_id, + "status": self._RUNNING, + "created_at": now, + "updated_at": now, + "case_name": case_name, + } + target = self._path(self._RUNNING, execution_id) + _atomic_write_json(target, meta) + return RolloutExecution( + execution_id=execution_id, + status=self._RUNNING, + created_at=now, + updated_at=now, + case_name=case_name, + ) + + async def get(self, execution_id: str) -> RolloutExecution | None: + for status in (self._COMPLETED, self._FAILED, self._RUNNING): + data = self._read_meta(self._path(status, execution_id)) + if data is None: + continue + rollout: Rollout | None = None + if status == self._COMPLETED: + rollout_data = data.get("rollout") + if isinstance(rollout_data, dict): + rollout = rollout_from_dict(rollout_data) + return RolloutExecution( + execution_id=data["execution_id"], + status=status, + created_at=float(data.get("created_at", 0.0)), + updated_at=float(data.get("updated_at", 0.0)), + case_name=data.get("case_name", ""), + finished_at=( + float(data["finished_at"]) if data.get("finished_at") is not None else None + ), + rollout=rollout, + error=data.get("error"), + ) + return None + + async def count_by_status(self) -> dict[str, int]: + counts: dict[str, int] = {} + for status in (self._RUNNING, self._COMPLETED, self._FAILED): + counts[status] = sum( + 1 for p in (self._spool_dir / status).iterdir() if p.suffix == ".json" + ) + return counts + + async def mark_completed(self, execution_id: str, rollout: Rollout) -> None: + now = time.time() + running_path = self._path(self._RUNNING, execution_id) + meta = self._read_meta(running_path) + # If the running record is gone (e.g. the service restarted while a task + # was in flight, or a prior failure already wrote a terminal record), + # fall back to a synthetic metadata skeleton so the completed payload + # is still persisted instead of raising a 500. + if meta is None: + completed_existing = self._read_meta(self._path(self._COMPLETED, execution_id)) + failed_existing = self._read_meta(self._path(self._FAILED, execution_id)) + existing = completed_existing or failed_existing or {} + meta = { + "execution_id": execution_id, + "status": self._RUNNING, + "created_at": existing.get("created_at", now), + "updated_at": existing.get("updated_at", now), + "case_name": existing.get("case_name", ""), + } + payload = { + **meta, + "status": self._COMPLETED, + "updated_at": now, + "finished_at": now, + "rollout": rollout_to_dict(rollout), + } + completed_path = self._path(self._COMPLETED, execution_id) + _atomic_write_json(completed_path, payload) + try: + running_path.unlink() + except FileNotFoundError: + pass + + async def mark_failed(self, execution_id: str, error: str) -> None: + now = time.time() + running_path = self._path(self._RUNNING, execution_id) + meta = self._read_meta(running_path) + if meta is None: + completed_existing = self._read_meta(self._path(self._COMPLETED, execution_id)) + failed_existing = self._read_meta(self._path(self._FAILED, execution_id)) + existing = completed_existing or failed_existing or {} + # Already in a terminal state: nothing to do. + if completed_existing is not None or failed_existing is not None: + return + meta = { + "execution_id": execution_id, + "status": self._RUNNING, + "created_at": existing.get("created_at", now), + "updated_at": existing.get("updated_at", now), + "case_name": existing.get("case_name", ""), + } + payload = { + **meta, + "status": self._FAILED, + "updated_at": now, + "finished_at": now, + "error": str(redact_sensitive({"error": error})["error"]), + } + failed_path = self._path(self._FAILED, execution_id) + _atomic_write_json(failed_path, payload) + try: + running_path.unlink() + except FileNotFoundError: + pass + + @property + def spool_dir(self) -> Path: + return self._spool_dir + + +def _cleanup_spool_dir(path: Path) -> None: + """Best-effort recursive removal of the per-process spool directory on exit.""" + try: + shutil.rmtree(path, ignore_errors=True) + except Exception: + pass + + +def _atomic_write_json(path: Path, data: dict[str, Any]) -> None: + """Write JSON atomically via temp-file + rename so readers never see partial writes.""" + tmp_path = path.with_suffix(path.suffix + ".tmp") + with tmp_path.open("w", encoding="utf-8") as fh: + json.dump(data, fh, ensure_ascii=False, default=str) + tmp_path.replace(path) + + +def create_dataset_service_app( + *, + service_name: str, + make_case_loader: CaseLoaderFactory, + make_rollout_executor: RolloutExecutorFactory, + max_rollout_concurrency: int | None = None, + rollout_thread_workers: int | None = None, +) -> FastAPI: + """Create a generic remote dataset service from train framework components.""" + + if max_rollout_concurrency is not None and max_rollout_concurrency <= 0: + raise ValueError("max_rollout_concurrency must be > 0") + if rollout_thread_workers is not None and rollout_thread_workers <= 0: + raise ValueError("rollout_thread_workers must be > 0") + + app = FastAPI(title=f"OpenViking {service_name} Dataset Service") + app.state.service_name = service_name + app.state.make_case_loader = make_case_loader + app.state.make_rollout_executor = make_rollout_executor + app.state.rollout_executions = RolloutExecutionStore() + app.state.max_rollout_concurrency = max_rollout_concurrency + app.state.rollout_semaphore = ( + asyncio.Semaphore(max_rollout_concurrency) if max_rollout_concurrency is not None else None + ) + app.state.rollout_thread_workers = rollout_thread_workers + app.state.rollout_thread_pool = ( + ThreadPoolExecutor( + max_workers=rollout_thread_workers, + thread_name_prefix=f"{service_name}-rollout", + ) + if rollout_thread_workers is not None + else None + ) + + @app.on_event("shutdown") + async def shutdown_rollout_thread_pool() -> None: + pool = app.state.rollout_thread_pool + if pool is not None: + pool.shutdown(wait=False, cancel_futures=True) + + @app.get("/health") + async def health() -> dict[str, Any]: + return { + "status": "ok", + "service": app.state.service_name, + "max_rollout_concurrency": app.state.max_rollout_concurrency, + "rollout_thread_workers": app.state.rollout_thread_workers, + "rollout_executions": await app.state.rollout_executions.count_by_status(), + } + + @app.post("/v1/cases/query") + async def query_cases(request: CasesQueryRequest) -> dict[str, Any]: + loader = app.state.make_case_loader( + request.dataset, + request.domain, + request.split, + dict(request.filters or {}), + ) + cases = await _load_case_page( + loader, + cursor=request.cursor, + limit=request.limit, + ) + next_offset = int(request.cursor or "0") + len(cases) + next_cursor = str(next_offset) if len(cases) >= request.limit else None + return { + "cases": [case_to_dict(case) for case in cases], + "next_cursor": next_cursor, + } + + @app.post("/v1/rollouts/execute") + async def execute_rollout(request: RolloutExecuteRequest) -> dict[str, Any]: + case = case_from_dict(request.case) + execution = await app.state.rollout_executions.create(case_name=case.name) + asyncio.create_task(_run_rollout_execution(app, execution.execution_id, request)) + return execution_to_dict(execution) + + @app.get("/v1/rollouts/executions/{execution_id}") + async def get_rollout_execution(execution_id: str) -> dict[str, Any]: + execution = await app.state.rollout_executions.get(execution_id) + if execution is None: + raise HTTPException( + status_code=404, + detail=f"Rollout execution not found: {execution_id}", + ) + return execution_to_dict(execution) + + return app + + +async def _run_rollout_execution( + app: FastAPI, + execution_id: str, + request: RolloutExecuteRequest, +) -> None: + case = case_from_dict(request.case) + try: + semaphore = app.state.rollout_semaphore + if semaphore is None: + rollout = await _execute_rollout_request_hosted(app, request, case) + else: + async with semaphore: + rollout = await _execute_rollout_request_hosted(app, request, case) + await app.state.rollout_executions.mark_completed(execution_id, rollout) + except Exception as exc: + logger.exception( + "rollout execution failed execution_id=%s case=%s", + execution_id, + case.name, + ) + try: + await app.state.rollout_executions.mark_failed( + execution_id, + str(redact_sensitive({"error": str(exc)})["error"]), + ) + except Exception: + # Last-ditch: do not let failures in the failure-recording path + # raise into the event loop (the task was already fire-and-forget). + logger.exception( + "rollout execution mark_failed itself failed execution_id=%s", + execution_id, + ) + + +async def _execute_rollout_request_hosted( + app: FastAPI, + request: RolloutExecuteRequest, + case: Case, +) -> Rollout: + """Execute one rollout either on the ASGI loop or in the service thread pool. + + Some benchmark backends (notably full agent-loop implementations) include + blocking synchronous sections that can starve the uvicorn event loop when + many rollouts are hosted by the same process. When a rollout thread pool is + configured, the whole rollout coroutine is driven by a fresh event loop in a + worker thread so request/polling endpoints remain responsive. + """ + + pool = getattr(app.state, "rollout_thread_pool", None) + if pool is None: + return await _execute_rollout_request(app, request, case) + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + pool, + _execute_rollout_request_in_thread, + app, + request, + case, + ) + + +def _execute_rollout_request_in_thread( + app: FastAPI, + request: RolloutExecuteRequest, + case: Case, +) -> Rollout: + loop = getattr(_rollout_worker_state, "loop", None) + if loop is None or loop.is_closed(): + loop = asyncio.new_event_loop() + _rollout_worker_state.loop = loop + asyncio.set_event_loop(loop) + # Keep the worker-thread event loop alive across rollout executions. Some + # provider stacks schedule async client cleanup tasks (e.g. httpx.aclose) + # late in the request lifecycle; using asyncio.run() here would close the + # loop immediately after each rollout and those cleanup tasks can then raise + # "RuntimeError: Event loop is closed". + try: + result = loop.run_until_complete(_execute_rollout_request(app, request, case)) + loop.run_until_complete(asyncio.sleep(0)) + return result + finally: + asyncio.set_event_loop(loop) + + +async def _execute_rollout_request( + app: FastAPI, + request: RolloutExecuteRequest, + case: Case, +) -> Rollout: + options = dict(request.options or {}) + executor = app.state.make_rollout_executor(options) + rollouts = await executor.execute( + [case], + policy_set_from_dict(request.policy_set), + ExecutionContext( + policy_snapshot_id=str(request.execution_context["policy_snapshot_id"]), + metadata=dict(request.execution_context.get("metadata") or {}), + ), + ) + return rollouts[0] + + +async def _load_case_page(loader: Any, *, cursor: str | None, limit: int) -> list[Case]: + offset = int(cursor or "0") + selected: list[Case] = [] + seen = 0 + async for batch in loader.batches(None): + for case in batch: + if seen < offset: + seen += 1 + continue + if len(selected) >= limit: + return selected + selected.append(case) + seen += 1 + return selected + + +def execution_to_dict(execution: RolloutExecution) -> dict[str, Any]: + data: dict[str, Any] = { + "execution_id": execution.execution_id, + "status": execution.status, + "case_name": execution.case_name, + "created_at": execution.created_at, + "updated_at": execution.updated_at, + "error": execution.error, + } + if execution.finished_at is not None: + data["duration_ms"] = round((execution.finished_at - execution.created_at) * 1000.0, 2) + if execution.rollout is not None: + data["rollout"] = rollout_to_dict(execution.rollout) + return data + + +def case_to_dict(case: Case) -> dict[str, Any]: + return { + "name": case.name, + "task_signature": case.task_signature, + "input": jsonable(case.input), + "rubric": { + "name": case.rubric.name, + "description": case.rubric.description, + "criteria": [ + { + "name": criterion.name, + "description": criterion.description, + "required": criterion.required, + "weight": criterion.weight, + "metadata": jsonable(criterion.metadata), + } + for criterion in case.rubric.criteria + ], + "metadata": jsonable(case.rubric.metadata), + }, + "metadata": jsonable(case.metadata), + } + + +def case_from_dict(data: dict[str, Any]) -> Case: + rubric = data["rubric"] + return Case( + name=data["name"], + task_signature=data["task_signature"], + input=dict(data.get("input") or {}), + rubric=Rubric( + name=rubric["name"], + description=rubric.get("description", ""), + criteria=[ + RubricCriterion( + name=item["name"], + description=item.get("description", ""), + required=bool(item.get("required", True)), + weight=float(item.get("weight", 1.0)), + metadata=dict(item.get("metadata") or {}), + ) + for item in rubric.get("criteria", []) + ], + metadata=dict(rubric.get("metadata") or {}), + ), + metadata=dict(data.get("metadata") or {}), + ) + + +def policy_set_from_dict(data: dict[str, Any]) -> ExperienceSet: + return ExperienceSet( + root_uri=data["root_uri"], + policies=[ + Experience( + name=item["name"], + uri=item["uri"], + version=int(item["version"]), + status=item["status"], + content=item["content"], + metadata=dict(item.get("metadata") or {}), + ) + for item in data.get("policies", []) + ], + metadata=dict(data.get("metadata") or {}), + ) + + +def policy_set_to_dict(policy_set: ExperienceSet) -> dict[str, Any]: + return { + "root_uri": policy_set.root_uri, + "policies": [ + { + "name": item.name, + "uri": item.uri, + "version": item.version, + "status": item.status, + "content": item.content, + "metadata": jsonable(item.metadata), + } + for item in policy_set.policies + ], + "metadata": jsonable(policy_set.metadata), + } + + +def rollout_to_dict(rollout: Rollout) -> dict[str, Any]: + return { + "case": case_to_dict(rollout.case), + "messages": [message.to_dict() for message in rollout.messages], + "policy_snapshot_id": rollout.policy_snapshot_id, + "evaluation": jsonable(evaluation_to_dict(rollout.evaluation)), + "metadata": jsonable(rollout.metadata), + } + + +def rollout_from_dict(data: dict[str, Any]) -> Rollout: + """Inverse of ``rollout_to_dict``. Used when reloading spooled Rollout payloads.""" + from openviking.message import Message + + return Rollout( + case=case_from_dict(data["case"]), + messages=[Message.from_dict(item) for item in data.get("messages", [])], + policy_snapshot_id=data["policy_snapshot_id"], + evaluation=evaluation_from_dict(data.get("evaluation")), + metadata=dict(data.get("metadata") or {}), + ) + + +def evaluation_from_dict(data: dict[str, Any] | None) -> RubricEvaluation | None: + if data is None: + return None + return RubricEvaluation( + passed=bool(data.get("passed")), + score=float(data.get("score") or 0.0), + criterion_results=[ + CriterionResult( + criterion_name=item.get("criterion_name", "unknown"), + passed=bool(item.get("passed")), + score=float(item.get("score") or 0.0), + feedback=[str(value) for value in item.get("feedback", [])], + evidence=[str(value) for value in item.get("evidence", [])], + metadata=dict(item.get("metadata") or {}), + ) + for item in data.get("criterion_results", []) + ], + feedback=[str(value) for value in data.get("feedback", [])], + metadata=dict(data.get("metadata") or {}), + ) + + +def jsonable(value: Any) -> Any: + if hasattr(value, "model_dump"): + return jsonable(value.model_dump(mode="json")) + if isinstance(value, Enum): + return value.value + if isinstance(value, dict): + return {str(jsonable(key)): jsonable(item) for key, item in value.items()} + if isinstance(value, list | tuple): + return [jsonable(item) for item in value] + return value + + +def evaluation_to_dict(evaluation: RubricEvaluation | None) -> dict[str, Any] | None: + if evaluation is None: + return None + return { + "passed": evaluation.passed, + "score": evaluation.score, + "criterion_results": [ + { + "criterion_name": result.criterion_name, + "passed": result.passed, + "score": result.score, + "feedback": result.feedback, + "evidence": result.evidence, + "metadata": jsonable(result.metadata), + } + for result in evaluation.criterion_results + ], + "feedback": evaluation.feedback, + "metadata": jsonable(evaluation.metadata), + } diff --git a/openviking/session/train/components/event_recorder.py b/openviking/session/train/components/event_recorder.py new file mode 100644 index 0000000000..9e6cc31fb5 --- /dev/null +++ b/openviking/session/train/components/event_recorder.py @@ -0,0 +1,179 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Streaming JSONL event recording for session train/eval pipelines.""" + +from __future__ import annotations + +import asyncio +import inspect +import json +from dataclasses import dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +from openviking.session.train.components.dataset_service import jsonable +from openviking.session.train.components.reporter import NoopPipelineLifecycleHook + + +@dataclass(slots=True) +class JsonlEventRecorder: + """Append train/eval events as one flushed JSON object per line.""" + + path: Path + default_fields: dict[str, Any] = field(default_factory=dict) + _lock: asyncio.Lock = field(init=False, repr=False) + _sequence: int = field(default=0, init=False) + + def __post_init__(self) -> None: + self.path = self.path.expanduser().resolve() + self._lock = asyncio.Lock() + + async def record(self, event: str, **fields: Any) -> None: + async with self._lock: + self._sequence += 1 + payload = { + "time": datetime.now(timezone.utc).isoformat(), + "sequence": self._sequence, + "event": event, + **self.default_fields, + **fields, + } + self.path.parent.mkdir(parents=True, exist_ok=True) + with self.path.open("a", encoding="utf-8") as file: + file.write(json.dumps(jsonable(payload), ensure_ascii=False, sort_keys=True)) + file.write("\n") + file.flush() + + +@dataclass(slots=True) +class CompositeEventRecorder: + """Fan out event records to multiple recorder implementations.""" + + recorders: tuple[Any, ...] + + async def record(self, event: str, **fields: Any) -> None: + for recorder in self.recorders: + record = getattr(recorder, "record", None) + if record is None: + continue + result = record(event, **fields) + if inspect.isawaitable(result): + await result + + +@dataclass(slots=True) +class JsonlPipelineEventHook(NoopPipelineLifecycleHook): + """Lifecycle hook that streams high-level pipeline reports to JSONL.""" + + recorder: JsonlEventRecorder + + async def on_epoch_start(self, *, epoch: int, context: Any) -> None: + await self.recorder.record( + "epoch_start", + stage="epoch_start", + **_merge_fields(_context_fields(context), {"epoch": epoch}), + ) + + async def on_train_rollout_report( + self, + *, + report: dict[str, Any], + context: Any, + ) -> None: + await self.recorder.record( + "train_rollout", + stage="train_rollout", + **_merge_fields(_context_fields(context), _report_fields(report)), + ) + + async def on_train_report( + self, + *, + report: dict[str, Any], + context: Any, + ) -> None: + await self.recorder.record( + "train_result", + stage="train", + **_merge_fields(_context_fields(context), _report_fields(report)), + ) + + async def on_eval_report( + self, + *, + label: str, + report: dict[str, Any], + context: Any, + ) -> None: + stage = str(report.get("rollout_stage") or label) + await self.recorder.record( + stage, + stage=stage, + **_merge_fields(_context_fields(context), _report_fields(report)), + ) + + async def on_run_summary( + self, + *, + title: str, + fields: dict[str, Any], + baseline_eval: dict[str, Any] | None = None, + final_eval: dict[str, Any] | None = None, + accuracy_delta: float | None = None, + output_path: str | None = None, + rollouts_root: str | None = None, + rollouts_index_path: str | None = None, + latest_failed_rollout: str | None = None, + ) -> None: + await self.recorder.record( + "run_summary", + stage="run_summary", + title=title, + fields=dict(fields), + baseline_eval=baseline_eval, + final_eval=final_eval, + accuracy_delta=accuracy_delta, + output_path=output_path, + rollouts_root=rollouts_root, + rollouts_index_path=rollouts_index_path, + latest_failed_rollout=latest_failed_rollout, + ) + + +def _merge_fields(*items: dict[str, Any]) -> dict[str, Any]: + merged: dict[str, Any] = {} + for item in items: + merged.update(item) + return merged + + +def _context_fields(context: Any) -> dict[str, Any]: + metadata = dict(getattr(context, "execution_metadata", {}) or {}) + fields: dict[str, Any] = {} + for key in ("epoch", "training", "rollout_stage", "eval_split"): + if key in metadata: + fields[key] = metadata[key] + return fields + + +def _report_fields(report: dict[str, Any]) -> dict[str, Any]: + excluded_keys = {"commit_results"} + fields = {key: value for key, value in report.items() if key not in excluded_keys} + commit_results = report.get("commit_results") + if isinstance(commit_results, list): + fields["commit_trace_ids"] = _commit_field_values(commit_results, "trace_id") + fields["commit_task_ids"] = _commit_field_values(commit_results, "task_id") + fields["commit_telemetry_ids"] = _commit_field_values(commit_results, "telemetry_id") + return fields + + +def _commit_field_values(commit_results: list[Any], key: str) -> list[str]: + values: list[str] = [] + for item in commit_results: + if not isinstance(item, dict): + continue + value = str(item.get(key) or "").strip() + if value: + values.append(value) + return values diff --git a/openviking/session/train/components/gradient_estimator.py b/openviking/session/train/components/gradient_estimator.py new file mode 100644 index 0000000000..80de733ffd --- /dev/null +++ b/openviking/session/train/components/gradient_estimator.py @@ -0,0 +1,290 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""ExtractLoop-backed GradientEstimator component.""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from typing import Any + +from openviking.message import Message +from openviking.server.identity import RequestContext +from openviking.session.memory.agent_experience_context_provider import ( + AgentExperienceContextProvider, +) +from openviking.session.memory.dataclass import MemoryFile, StoredLink +from openviking.session.memory.extract_loop import ExtractLoop +from openviking.session.memory.memory_isolation_handler import MemoryIsolationHandler +from openviking.session.train.domain import ExperienceSet, RolloutAnalysis, Trajectory +from openviking.session.train.gradients import PatchSemanticGradient +from openviking.session.train.utils import first_uri, safe_int +from openviking.storage.viking_fs import get_viking_fs +from openviking.telemetry import tracer +from openviking_cli.utils import get_logger +from openviking_cli.utils.config import get_openviking_config + +logger = get_logger(__name__) + + +@dataclass(slots=True) +class ExperienceGradientContext: + """Context for ExperienceGradientEstimator.""" + + request_context: RequestContext + messages: list[Message] + strict_extract_errors: bool = False + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class ExperienceGradientEstimator: + """Estimate PatchSemanticGradients via experience ExtractLoop. + + This component reuses AgentExperienceContextProvider and ExtractLoop but stops + before MemoryUpdater.apply_operations. The resolved operations are converted + into PatchSemanticGradient instances. + """ + + viking_fs: Any = None + vlm: Any = None + + @tracer( + "train.gradient_estimator.experience.estimate", + ignore_result=True, + ignore_args=True, + ) + async def estimate( + self, + analysis: RolloutAnalysis, + experience_set: ExperienceSet, + context: ExperienceGradientContext, + ) -> list[PatchSemanticGradient]: + if context is None or context.request_context is None: + raise ValueError("ExperienceGradientContext.request_context is required") + + extract_context = _context_with_analysis_messages(context, analysis) + + async def estimate_one(trajectory: Trajectory) -> list[PatchSemanticGradient]: + try: + operations = await self._run_extract_loop(trajectory, extract_context) + except Exception: + logger.exception("Experience gradient estimation failed") + if context.strict_extract_errors: + raise + return [] + if operations is None: + return [] + return _operations_to_gradients( + operations=operations, + trajectory=trajectory, + analysis=analysis, + experience_set=experience_set, + ) + + gradient_batches = await asyncio.gather( + *(estimate_one(trajectory) for trajectory in analysis.trajectories) + ) + return [gradient for batch in gradient_batches for gradient in batch] + + @tracer( + "train.gradient_estimator.experience.extract_loop", + ignore_result=True, + ignore_args=True, + ) + async def _run_extract_loop( + self, + trajectory: Trajectory, + context: ExperienceGradientContext, + ): + config = get_openviking_config() + vlm = self.vlm or config.vlm.get_vlm_instance() + viking_fs = self.viking_fs or get_viking_fs() + if viking_fs is None: + raise RuntimeError("VikingFS is required for experience gradient estimation") + + provider = AgentExperienceContextProvider( + messages=context.messages, + trajectory_summary=trajectory.content, + trajectory_uri=trajectory.uri, + ) + if hasattr(provider, "get_extract_context"): + extract_context = provider.get_extract_context() + else: + extract_context = context + isolation_handler = MemoryIsolationHandler( + context.request_context, + extract_context, + allowed_memory_types={"experiences"}, + ) + isolation_handler.prepare_messages() + + provider._isolation_handler = isolation_handler + provider._ctx = context.request_context + provider._viking_fs = viking_fs + + orchestrator = ExtractLoop( + vlm=vlm, + viking_fs=viking_fs, + ctx=context.request_context, + context_provider=provider, + isolation_handler=isolation_handler, + thinking=True, + ) + operations, _ = await orchestrator.run() + return operations + + +def _context_with_analysis_messages( + context: ExperienceGradientContext, + analysis: RolloutAnalysis, +) -> ExperienceGradientContext: + messages = analysis.metadata.get("rollout_messages") + if not messages: + return context + return ExperienceGradientContext( + request_context=context.request_context, + messages=list(messages), + strict_extract_errors=context.strict_extract_errors, + metadata=dict(context.metadata), + ) + + +def _operations_to_gradients( + *, + operations: Any, + trajectory: Trajectory, + analysis: RolloutAnalysis, + experience_set: ExperienceSet, +) -> list[PatchSemanticGradient]: + gradients: list[PatchSemanticGradient] = [] + for op in getattr(operations, "upsert_operations", []) or []: + if getattr(op, "memory_type", None) != "experiences": + continue + fields = dict(getattr(op, "memory_fields", {}) or {}) + after_content = str(fields.get("content") or "") + if not after_content.strip(): + continue + + old_file = getattr(op, "old_memory_file_content", None) + target_name = str(fields.get("experience_name") or _fallback_experience_name(op)) + target_uri = first_uri(getattr(op, "uris", []) or []) + base_version = _base_version(old_file, target_uri, experience_set) + after_file = _operation_after_file( + fields=fields, + target_name=target_name, + target_uri=target_uri, + old_file=old_file, + ) + + gradients.append( + PatchSemanticGradient( + before_file=old_file, + after_file=after_file, + base_version=base_version, + rationale=( + "ExtractLoop proposed an experience content update " + f"from trajectory {trajectory.uri}." + ), + links=[ + StoredLink( + from_uri=target_uri or "", + to_uri=trajectory.uri, + link_type="derived_from", + weight=1.0, + match_text=None, + description="", + ) + ], + confidence=_confidence(trajectory, analysis), + metadata={ + "memory_fields": fields, + "uris": list(getattr(op, "uris", []) or []), + "trajectory_outcome": trajectory.outcome, + "rubric_passed": analysis.evaluation.passed, + "supersedes": fields.get("supersedes"), + "training_category": _trajectory_training_category(trajectory, analysis), + }, + ) + ) + return gradients + + +def _trajectory_training_category( + trajectory: Trajectory, + analysis: RolloutAnalysis, +) -> str: + trajectory_metadata = dict(getattr(trajectory, "metadata", {}) or {}) + for key in ("training_category", "category"): + value = trajectory_metadata.get(key) + if value: + return str(value) + + analysis_metadata = dict(getattr(analysis, "metadata", {}) or {}) + for key in ("training_category", "category", "case_task_signature", "task_signature"): + value = analysis_metadata.get(key) + if value: + return str(value) + + if trajectory.retrieval_anchor: + return str(trajectory.retrieval_anchor) + return str(trajectory.name) + + +def _operation_after_file( + *, + fields: dict[str, Any], + target_name: str, + target_uri: str | None, + old_file: MemoryFile | None, +) -> MemoryFile: + extra_fields = dict(getattr(old_file, "extra_fields", {}) or {}) + for key, value in fields.items(): + if key != "content": + extra_fields[key] = value + extra_fields["memory_type"] = "experiences" + extra_fields["experience_name"] = target_name + return MemoryFile( + uri=target_uri, + content=str(fields.get("content") or ""), + links=list(getattr(old_file, "links", []) or []), + backlinks=list(getattr(old_file, "backlinks", []) or []), + memory_type="experiences", + extra_fields=extra_fields, + ) + + +def _fallback_experience_name(op: Any) -> str: + uri = first_uri(getattr(op, "uris", []) or []) + if uri: + return uri.rstrip("/").split("/")[-1].removesuffix(".md") + return "unknown_experience" + + +def _base_version( + old_file: Any, target_uri: str | None, experience_set: ExperienceSet +) -> int | None: + if old_file is not None: + fields = getattr(old_file, "extra_fields", {}) or {} + version = safe_int(fields.get("version")) + if version is not None: + return version + if target_uri: + for policy in experience_set.policies: + if policy.uri == target_uri: + return policy.version + return None + + +def _confidence(trajectory: Trajectory, analysis: RolloutAnalysis) -> float: + confidence = 0.5 + if analysis.evaluation.passed: + confidence += 0.2 + outcome = str(trajectory.outcome).lower() + if outcome == "success": + confidence += 0.2 + elif outcome in {"failure", "partial"}: + confidence -= 0.2 + elif outcome == "unfinished": + confidence -= 0.1 + return max(0.0, min(1.0, confidence)) diff --git a/openviking/session/train/components/memory_store.py b/openviking/session/train/components/memory_store.py new file mode 100644 index 0000000000..78850a9208 --- /dev/null +++ b/openviking/session/train/components/memory_store.py @@ -0,0 +1,199 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Load train-domain policy set objects from existing memory files.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from openviking.core.skill_loader import SkillLoader +from openviking.server.identity import RequestContext +from openviking.session.memory.utils.memory_file_utils import MemoryFileUtils +from openviking.session.train.domain import Policy, PolicySet, PolicyStatus +from openviking.storage.viking_fs import get_viking_fs +from openviking.telemetry import tracer + +_HIDDEN_MEMORY_FILES = {".overview.md", ".abstract.md"} +_ALLOWED_STATUSES = {"draft", "staging", "production", "deprecated", "archived"} + + +@dataclass(slots=True) +class ExperienceSetLoader: + """Build an ExperienceSet by reading an experiences directory.""" + + viking_fs: Any = None + + @tracer("train.experience_set_loader.load", ignore_result=True, ignore_args=True) + async def load(self, root_uri: str, ctx: RequestContext | None = None) -> PolicySet: + if ctx is None: + raise ValueError("ExperienceSetLoader.load requires request_context ctx") + viking_fs = self.viking_fs or get_viking_fs() + if viking_fs is None: + raise RuntimeError("VikingFS is required to load an ExperienceSet") + + try: + entries = await viking_fs.ls(root_uri, output="original", ctx=ctx) + except Exception: + entries = [] + + policies: list[Policy] = [] + for entry in entries or []: + if not isinstance(entry, dict): + continue + if entry.get("isDir") or entry.get("is_dir"): + continue + name = str(entry.get("name") or "") + uri = str(entry.get("uri") or "") + if not uri.endswith(".md") or name in _HIDDEN_MEMORY_FILES: + continue + if uri.endswith("/.overview.md") or uri.endswith("/.abstract.md"): + continue + + raw = await viking_fs.read_file(uri, ctx=ctx) or "" + mf = MemoryFileUtils.read(raw, uri=uri) + fields = dict(mf.extra_fields or {}) + policy_name = str(fields.get("experience_name") or name.removesuffix(".md")) + version = _safe_int(fields.get("version"), default=1) + status = _safe_status(fields.get("status")) + metadata = dict(fields) + metadata.setdefault("memory_type", mf.memory_type or fields.get("memory_type")) + policies.append( + Policy( + name=policy_name, + uri=uri, + version=version, + status=status, + content=mf.plain_content(), + metadata=metadata, + links=list(mf.links or []), + backlinks=list(mf.backlinks or []), + ) + ) + + policies.sort(key=lambda p: p.uri) + return PolicySet( + root_uri=root_uri, + policies=policies, + metadata={"source": "memory_store"}, + viking_fs=viking_fs, + request_context=ctx, + ) + + +def _safe_int(value: Any, *, default: int) -> int: + try: + parsed = int(value) + except (TypeError, ValueError): + return default + return parsed if parsed > 0 else default + + +def _safe_status(value: Any) -> PolicyStatus: + status = str(value or "production") + if status not in _ALLOWED_STATUSES: + return "production" + return status # type: ignore[return-value] + + +@dataclass(slots=True) +class SkillSetLoader: + """Build a PolicySet by reading a skills directory. + + Each skill is represented as a subdirectory containing a ``SKILL.md`` file + with YAML frontmatter. Skill-specific fields (description, allowed_tools, + tags, …) are stored in the policy ``metadata`` dict. + """ + + viking_fs: Any = None + + @tracer("train.skill_set_loader.load", ignore_result=True, ignore_args=True) + async def load(self, root_uri: str, ctx: RequestContext | None = None) -> PolicySet: + if ctx is None: + raise ValueError("SkillSetLoader.load requires request_context ctx") + viking_fs = self.viking_fs or get_viking_fs() + if viking_fs is None: + raise RuntimeError("VikingFS is required to load a SkillSet") + + try: + entries = await viking_fs.ls(root_uri, output="original", ctx=ctx) + except Exception: + entries = [] + + policies: list[Policy] = [] + for entry in entries or []: + if not isinstance(entry, dict): + continue + if not (entry.get("isDir") or entry.get("is_dir")): + continue + dir_name = str(entry.get("name") or "") + dir_uri = str(entry.get("uri") or "") + if not dir_name or not dir_uri: + continue + if dir_name.startswith("."): + continue + + skill_md_uri = f"{dir_uri.rstrip('/')}/SKILL.md" + try: + raw = await viking_fs.read_file(skill_md_uri, ctx=ctx) + except Exception: + continue + if not raw: + continue + + try: + skill = SkillLoader.parse(raw, source_path=skill_md_uri) + except Exception: + # Fall back to generic memory-file parsing if SKILL.md + # doesn't have valid frontmatter. + mf = MemoryFileUtils.read(raw, uri=skill_md_uri) + fields = dict(mf.extra_fields or {}) + skill_name = str(fields.get("skill_name") or dir_name) + version = _safe_int(fields.get("version"), default=1) + status = _safe_status(fields.get("status")) + metadata = dict(fields) + metadata.setdefault("memory_type", "skills") + metadata["description"] = fields.get("description", "") + policies.append( + Policy( + name=skill_name, + uri=skill_md_uri, + version=version, + status=status, + content=mf.plain_content(), + metadata=metadata, + links=list(mf.links or []), + backlinks=list(mf.backlinks or []), + ) + ) + continue + + version = 1 + status: PolicyStatus = "production" + metadata: dict[str, Any] = { + "memory_type": "skills", + "description": skill.get("description", ""), + "allowed_tools": list(skill.get("allowed_tools") or []), + "tags": list(skill.get("tags") or []), + } + policies.append( + Policy( + name=str(skill.get("name") or dir_name), + uri=skill_md_uri, + version=version, + status=status, + content=str(skill.get("content") or ""), + metadata=metadata, + links=[], + backlinks=[], + ) + ) + + policies.sort(key=lambda p: p.uri) + return PolicySet( + root_uri=root_uri, + policies=policies, + metadata={"source": "skill_store"}, + viking_fs=viking_fs, + request_context=ctx, + ) diff --git a/openviking/session/train/components/policy_optimizer.py b/openviking/session/train/components/policy_optimizer.py new file mode 100644 index 0000000000..f7226589f3 --- /dev/null +++ b/openviking/session/train/components/policy_optimizer.py @@ -0,0 +1,855 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Policy optimizer implementations.""" + +from __future__ import annotations + +from collections import defaultdict +from dataclasses import dataclass, field +from typing import Any + +from openviking.message import Message +from openviking.server.identity import RequestContext +from openviking.session.memory.dataclass import MemoryFile, StoredLink +from openviking.session.memory.extract_loop import ExtractLoop +from openviking.session.memory.memory_isolation_handler import MemoryIsolationHandler +from openviking.session.memory.memory_updater import ExtractContext +from openviking.session.memory.patch_merge_context_provider import ( + PatchMergeContextProvider, + PatchMergePatch, +) +from openviking.session.train.domain import ( + Policy, + PolicyPlanItem, + PolicySet, + PolicyUpdatePlan, +) +from openviking.session.train.interfaces import SemanticGradient +from openviking.session.train.utils import first_uri, safe_int +from openviking.telemetry import tracer +from openviking_cli.utils import get_logger +from openviking_cli.utils.config import get_openviking_config + +logger = get_logger(__name__) + + +@dataclass(slots=True) +class PatchMergePolicyOptimizerContext: + """Context for PatchMergePolicyOptimizer.""" + + request_context: RequestContext + messages: list[Message] = field(default_factory=list) + + +@dataclass(slots=True) +class PatchMergePolicyOptimizer: + """Merge patch gradients with ExtractLoop before producing update plan items.""" + + viking_fs: Any = None + vlm: Any = None + memory_type: str = "experiences" + + @tracer( + "train.policy_optimizer.patch_merge.plan", + ignore_result=True, + ignore_args=True, + ) + async def plan( + self, + gradients: list[SemanticGradient], + policy_set: PolicySet, + context: PatchMergePolicyOptimizerContext | None = None, + ) -> PolicyUpdatePlan: + if context is None: + raise ValueError("PatchMergePolicyOptimizerContext.request_context is required") + + patch_gradients = list(gradients) + if not patch_gradients: + return PolicyUpdatePlan( + items=[], + metadata={ + "optimizer": "patch_merge", + "memory_type": self.memory_type, + "gradient_count": len(gradients), + "patch_gradient_count": 0, + }, + ) + + operations = await self._run_merge_extract_loop( + gradients=patch_gradients, + policy_set=policy_set, + context=context, + ) + items = _operations_to_plan_items( + operations=operations, + gradients=patch_gradients, + policy_set=policy_set, + memory_type=self.memory_type, + ) + _log_merge_output( + target="all", + operations=operations, + plan_items=items, + console=False, + ) + + return PolicyUpdatePlan( + items=items, + metadata={ + "optimizer": "patch_merge", + "memory_type": self.memory_type, + "gradient_count": len(gradients), + "patch_gradient_count": len(patch_gradients), + "gradients": [ + _gradient_to_dict(idx, gradient) + for idx, gradient in enumerate(patch_gradients) + ], + }, + ) + + @tracer( + "train.policy_optimizer.patch_merge.extract_loop", + ignore_result=True, + ignore_args=True, + ) + async def _run_merge_extract_loop( + self, + *, + gradients: list[SemanticGradient], + policy_set: PolicySet, + context: PatchMergePolicyOptimizerContext, + ): + config = get_openviking_config() + vlm = self.vlm or config.vlm.get_vlm_instance() + viking_fs = self.viking_fs or policy_set.viking_fs + if viking_fs is None: + raise RuntimeError("VikingFS is required for patch-merge policy optimization") + + extract_context = ExtractContext(list(context.messages or [])) + provider = PatchMergeContextProvider( + memory_type=self.memory_type, + required_file_uris=_required_file_uris(gradients, policy_set), + patches=[_gradient_to_merge_patch(gradient) for gradient in gradients], + ) + provider._ctx = context.request_context + provider._viking_fs = viking_fs + provider._extract_context = extract_context + + isolation_handler = MemoryIsolationHandler( + context.request_context, + extract_context, + allowed_memory_types={self.memory_type}, + ) + isolation_handler.prepare_messages() + provider._isolation_handler = isolation_handler + + _seed_read_file_contents(provider, gradients, policy_set) + prefetch_messages = await provider.prefetch() + provider.prefetch = _constant_prefetch(prefetch_messages) + _log_merge_input( + target="all", + provider=provider, + gradients=gradients, + prefetch_messages=prefetch_messages, + console=False, + ) + + orchestrator = ExtractLoop( + vlm=vlm, + viking_fs=viking_fs, + ctx=context.request_context, + context_provider=provider, + isolation_handler=isolation_handler, + max_iterations=1, + thinking=self.memory_type in {"trajectories", "experiences"}, + ) + operations, _ = await orchestrator.run() + return operations + +def _constant_prefetch(messages: list[dict[str, Any]]): + async def prefetch() -> list[dict[str, Any]]: + return list(messages) + + return prefetch + +def _log_merge_input( + *, + target: str, + provider: PatchMergeContextProvider, + gradients: list[SemanticGradient], + prefetch_messages: list[dict[str, Any]], + console: bool, +) -> None: + lines = [ + "\n========== PatchMergePolicyOptimizer Input =========", + f"target: {target}", + f"memory_type: {provider.memory_type}", + f"required_file_uris: {provider.required_file_uris}", + f"gradient_count: {len(gradients)}", + ] + for idx, gradient in enumerate(gradients): + before_file = gradient.before_file + after_file = gradient.after_file + lines.extend( + [ + "", + f"[Gradient {idx}]", + f"target_name: {gradient.target_name}", + f"target_uri: {gradient.target_uri}", + f"base_version: {gradient.base_version}", + f"confidence: {gradient.confidence}", + f"links: {_links_to_dicts(gradient.links)}", + f"rationale: {gradient.rationale}", + ] + ) + if after_file is not None: + lines.extend( + [ + "before_file:", + _memory_file_summary(before_file), + "after_file:", + _memory_file_summary(after_file), + ] + ) + lines.extend(["", "[Prefetch Messages]"]) + for idx, message in enumerate(prefetch_messages): + lines.extend( + [f"--- message {idx} role={message.get('role')} ---", str(message.get("content"))] + ) + lines.append("===================================================\n") + tracer.info("\n".join(lines), console=console) + +def _log_merge_output( + *, + target: str, + operations: Any, + plan_items: list[PolicyPlanItem], + console: bool, +) -> None: + lines = [ + "\n========== PatchMergePolicyOptimizer Output =========", + f"target: {target}", + "[Resolved Operations]", + _dump_model_or_value(operations), + "", + "[Policy Plan Items]", + ] + for idx, item in enumerate(plan_items): + lines.extend( + [ + f"--- item {idx} ---", + f"kind: {item.kind}", + f"memory_type: {item.memory_type}", + f"target_name: {item.target_name}", + f"target_uri: {item.target_uri}", + f"base_version: {item.base_version}", + f"confidence: {item.confidence}", + f"links: {_links_to_dicts(item.links)}", + "before_content:", + str(item.before_content), + "after_content:", + str(item.after_content), + f"metadata: {item.metadata}", + ] + ) + lines.append("====================================================\n") + tracer.info("\n".join(lines), console=console) + +def _dump_model_or_value(value: Any) -> str: + dumper = getattr(value, "model_dump_json", None) + if dumper is not None: + try: + return str(dumper(indent=2)) + except TypeError: + return str(dumper()) + return str(value) + +def _memory_file_summary(file: MemoryFile | None) -> str: + if file is None: + return "None" + return _dump_model_or_value( + { + "uri": file.uri, + "memory_type": file.memory_type, + "content": file.content, + "links": file.links, + "backlinks": file.backlinks, + "extra_fields": file.extra_fields, + } + ) + +def _gradient_to_dict(index: int, gradient: SemanticGradient) -> dict[str, Any]: + result = { + "index": index, + "target_name": gradient.target_name, + "target_uri": gradient.target_uri, + "base_version": gradient.base_version, + "rationale": gradient.rationale, + "links": _links_to_dicts(gradient.links), + "confidence": gradient.confidence, + "metadata": dict(gradient.metadata), + } + before_file = gradient.before_file + after_file = gradient.after_file + if before_file is not None: + result["before_file"] = _memory_file_to_dict(before_file) + if after_file is not None: + result["after_file"] = _memory_file_to_dict(after_file) + return result + +def _memory_file_to_dict(file: MemoryFile) -> dict[str, Any]: + return { + "uri": file.uri, + "memory_type": file.memory_type, + "content": file.content, + "links": list(file.links or []), + "backlinks": list(file.backlinks or []), + "extra_fields": dict(file.extra_fields or {}), + } + + +def _links_to_dicts(links: list[StoredLink] | None) -> list[dict[str, Any]]: + return [link.model_dump() for link in links or []] + +def _gradient_to_merge_patch(gradient: SemanticGradient) -> PatchMergePatch: + return PatchMergePatch( + before_file=gradient.before_file, + after_file=gradient.after_file, + metadata={ + "base_version": gradient.base_version, + "rationale": gradient.rationale, + "links": _links_to_dicts(gradient.links), + "confidence": gradient.confidence, + "gradient_metadata": _compact_gradient_metadata(gradient.metadata), + }, + ) + + +def _compact_gradient_metadata(metadata: dict[str, Any]) -> dict[str, Any]: + compact = dict(metadata) + memory_fields = compact.get("memory_fields") + if isinstance(memory_fields, dict) and "content" in memory_fields: + compact["memory_fields"] = { + key: value for key, value in memory_fields.items() if key != "content" + } + return compact + + +def _required_file_uris( + gradients: list[SemanticGradient], + policy_set: PolicySet, +) -> list[str]: + uris: list[str] = [] + for gradient in gradients: + uri = gradient.target_uri + if not uri: + superseded = _find_superseded_policy(_gradient_supersedes(gradient), policy_set) + uri = superseded.uri if superseded is not None else None + if uri and uri not in uris: + uris.append(uri) + return uris + +def _seed_read_file_contents( + provider: PatchMergeContextProvider, + gradients: list[SemanticGradient], + policy_set: PolicySet, +) -> None: + for policy in policy_set.policies: + if policy.uri in provider.required_file_uris: + provider.read_file_contents[policy.uri] = _policy_to_memory_file( + policy, memory_type=provider.memory_type + ) + for gradient in gradients: + before_file = gradient.before_file + target_uri = gradient.target_uri + if before_file is None or target_uri in provider.read_file_contents: + continue + if target_uri: + provider.read_file_contents[target_uri] = before_file + + +def _policy_to_memory_file(policy: Policy, *, memory_type: str = "experiences") -> MemoryFile: + name_field = _name_field_for_memory_type(memory_type) + extra_fields = dict(policy.metadata) + extra_fields["memory_type"] = memory_type + extra_fields[name_field] = policy.name + extra_fields.setdefault("version", policy.version) + extra_fields.setdefault("status", policy.status) + return MemoryFile( + uri=policy.uri, + content=policy.content, + links=list(policy.links or []), + backlinks=list(policy.backlinks or []), + memory_type=memory_type, + extra_fields=extra_fields, + ) + +def _operations_to_plan_items( + *, + operations: Any, + gradients: list[SemanticGradient], + policy_set: PolicySet, + memory_type: str, +) -> list[PolicyPlanItem]: + items: list[PolicyPlanItem] = [] + source_links_by_target = _source_trajectory_links_by_target(gradients, policy_set) + superseded_policies = _superseded_policies_for_gradients(gradients, policy_set) + confidence_values = [float(gradient.confidence) for gradient in gradients] + confidence = max(confidence_values) if confidence_values else None + name_field = _name_field_for_memory_type(memory_type) + + upsert_output_count = _upsert_output_count(operations, memory_type=memory_type) + replacement_source_uris_by_target = _replacement_source_uris_by_target(operations) + upsert_target_uris: set[str] = set() + for op in getattr(operations, "upsert_operations", []) or []: + if getattr(op, "memory_type", None) != memory_type: + continue + fields = dict(getattr(op, "memory_fields", {}) or {}) + after_content = str(fields.get("content") or "") + if not after_content.strip(): + continue + target_name = str( + fields.get(name_field) + or fields.get("name") + or _fallback_policy_name(op, memory_type=memory_type) + ) + target_uri = first_uri(getattr(op, "uris", []) or []) + old_file = getattr(op, "old_memory_file_content", None) + before_content = old_file.plain_content() if old_file is not None else None + if before_content is None and target_uri: + policy = _find_policy_by_uri(policy_set, target_uri) + before_content = policy.content if policy is not None else None + items.append( + PolicyPlanItem( + kind="upsert", + memory_type=memory_type, + target_name=target_name, + target_uri=target_uri, + before_content=before_content, + after_content=after_content, + base_version=_base_version_from_old_file_or_policy( + old_file, + target_uri, + policy_set, + ), + confidence=confidence, + links=_source_trajectory_links_for_plan_item( + target_uri=target_uri or "", + target_name=target_name, + before_content=before_content, + after_content=after_content, + source_links_by_target=source_links_by_target, + replacement_source_uris=replacement_source_uris_by_target.get( + target_uri or "", + [], + ), + include_all_sources=upsert_output_count == 1, + ), + metadata={ + "rationale": "PatchMergeContextProvider merged semantic gradients via ExtractLoop.", + "merge_gradient_count": len(gradients), + "merge_memory_fields": fields, + "superseded_experience_uris": [ + policy.uri for policy in superseded_policies + ], + }, + ) + ) + if target_uri: + upsert_target_uris.add(target_uri) + + delete_uris: set[str] = set() + for old_file in getattr(operations, "delete_file_contents", []) or []: + target_uri = old_file.uri + target_name = str( + (old_file.extra_fields or {}).get(name_field) + or (target_uri.rstrip("/").split("/")[-1].removesuffix(".md") if target_uri else "") + ) + if not target_uri or target_uri in upsert_target_uris: + continue + items.append( + PolicyPlanItem( + kind="delete", + memory_type=memory_type, + target_name=target_name, + target_uri=target_uri, + before_content=old_file.plain_content(), + after_content=None, + confidence=confidence, + links=_source_trajectory_links_for_plan_item( + target_uri=target_uri, + target_name=target_name, + before_content=old_file.plain_content(), + after_content=None, + source_links_by_target=source_links_by_target, + replacement_source_uris=[], + ), + metadata={ + "rationale": "PatchMergeContextProvider merge requested memory deletion.", + "merge_gradient_count": len(gradients), + }, + ) + ) + delete_uris.add(target_uri) + + for policy in superseded_policies: + if policy.uri in upsert_target_uris or policy.uri in delete_uris: + continue + items.append( + PolicyPlanItem( + kind="delete", + memory_type=memory_type, + target_name=policy.name, + target_uri=policy.uri, + before_content=policy.content, + after_content=None, + base_version=policy.version, + confidence=confidence, + links=_source_trajectory_links_from_experience(policy), + metadata={ + "rationale": "Superseded by broader experience from semantic gradient.", + "merge_gradient_count": len(gradients), + "superseded_by": [ + item.target_uri or item.target_name + for item in items + if item.kind == "upsert" + ], + }, + ) + ) + delete_uris.add(policy.uri) + return items + + +def _name_field_for_memory_type(memory_type: str) -> str: + """Return the extra_fields key for the policy name in a given memory type.""" + if memory_type == "experiences": + return "experience_name" + if memory_type == "skills": + return "skill_name" + if memory_type.endswith("s"): + return f"{memory_type[:-1]}_name" + return f"{memory_type}_name" + + +def _fallback_policy_name(op: Any, *, memory_type: str) -> str: + uri = first_uri(getattr(op, "uris", []) or []) + if uri: + # For skills: path/to/skills/my_skill/SKILL.md → my_skill + if memory_type == "skills" and uri.endswith("/SKILL.md"): + parts = uri.rstrip("/").split("/") + if len(parts) >= 2: + return parts[-2] + return uri.rstrip("/").split("/")[-1].removesuffix(".md") + return f"unknown_{memory_type.rstrip('s')}" + +def _source_trajectory_links_from_experience(policy: Policy | None) -> list[StoredLink]: + if policy is None: + return [] + links: list[StoredLink] = [] + for link in list(getattr(policy, "links", []) or []): + try: + stored_link = link if isinstance(link, StoredLink) else StoredLink(**dict(link)) + except Exception: + continue + if _is_source_trajectory_link(stored_link): + links.append(stored_link) + return links + + +def _superseded_policies_for_gradients( + gradients: list[SemanticGradient], + policy_set: PolicySet, +) -> list[Policy]: + policies: list[Policy] = [] + seen: set[str] = set() + for gradient in gradients: + policy = _find_superseded_policy(_gradient_supersedes(gradient), policy_set) + if policy is None or policy.uri in seen: + continue + seen.add(policy.uri) + policies.append(policy) + return policies + + +def _merge_source_trajectory_links(links: list[StoredLink]) -> list[StoredLink]: + merged: dict[tuple[str, str, str | None], StoredLink] = {} + for link in links: + if not _is_source_trajectory_link(link): + continue + key = (link.from_uri, link.to_uri, link.match_text) + if key in merged: + existing = merged[key] + update: dict[str, Any] = {"weight": max(existing.weight, link.weight)} + if link.description: + update["description"] = link.description + if link.created_at and not existing.created_at: + update["created_at"] = link.created_at + merged[key] = existing.model_copy(update=update) + else: + merged[key] = link + return list(merged.values()) + + +def _remap_source_trajectory_links( + links: list[StoredLink], + *, + target_uri: str, +) -> list[StoredLink]: + return [ + link.model_copy(update={"from_uri": target_uri}) + for link in links + if target_uri and _is_source_trajectory_link(link) and link.to_uri + ] + + +def _source_trajectory_links_for_plan_item( + *, + target_uri: str, + target_name: str, + before_content: str | None, + after_content: str | None, + source_links_by_target: dict[tuple[str, str], list[StoredLink]], + replacement_source_uris: list[str] | None = None, + include_all_sources: bool = False, +) -> list[StoredLink]: + """Return only source trajectory links whose patch target maps to this plan item. + + Patch merge can reconcile several independent patch proposals into one or + more final policy files. Source links belong to the patch proposal that + produced them, not to the whole merge batch. Therefore link propagation must + follow proposal-target/replacement provenance instead of broadcasting all + gradient links to every upsert. + """ + + links: list[StoredLink] = [] + seen_source_keys: set[tuple[str, str]] = set() + candidate_keys = _plan_item_source_keys( + target_uri=target_uri, + target_name=target_name, + before_content=before_content, + after_content=after_content, + source_links_by_target=source_links_by_target, + replacement_source_uris=replacement_source_uris or [], + include_all_sources=include_all_sources, + ) + for key in candidate_keys: + if key in seen_source_keys: + continue + seen_source_keys.add(key) + links.extend(source_links_by_target.get(key, [])) + return _merge_source_trajectory_links( + _remap_source_trajectory_links(links, target_uri=target_uri) + ) + + +def _plan_item_source_keys( + *, + target_uri: str, + target_name: str, + before_content: str | None, + after_content: str | None, + source_links_by_target: dict[tuple[str, str], list[StoredLink]], + replacement_source_uris: list[str] | None = None, + include_all_sources: bool = False, +) -> list[tuple[str, str]]: + keys: list[tuple[str, str]] = [] + all_keys = list(source_links_by_target.keys()) + + def add(key: tuple[str, str]) -> None: + if key in source_links_by_target and key not in keys: + keys.append(key) + + uri = str(target_uri or "") + name = str(target_name or "") + if uri: + add(("uri", uri)) + if name: + add(("name", name)) + for source_uri in replacement_source_uris or []: + add(("uri", source_uri)) + + # Existing-file updates and replacement deletes should inherit links from + # the previous canonical file that the merge output is editing/replacing. + for key in all_keys: + kind, value = key + if kind == "uri" and uri and value == uri: + add(key) + elif kind == "name" and name and value == name: + add(key) + + # New proposals that keep their target URI/name may not have old content. + # If there is exactly one source candidate with the same rendered content, + # treat it as this plan item's source. This handles URI/name normalization + # without turning duplicate-content batches into a broadcast. + content = str(after_content or "").strip() + if content: + matches = [ + key + for key in all_keys + if key[0] == "content" and key[1].strip() == content + ] + if len(matches) == 1: + add(matches[0]) + + source_identities = _source_identity_keys(source_links_by_target) + if include_all_sources: + for key in source_identities: + add(key) + + # Single-patch merge: if the final URI/name was normalized, there is still + # only one possible source, so carry its provenance forward. + if not keys and len(source_identities) == 1: + keys.extend(source_identities) + + return keys + + +def _source_trajectory_links_by_target( + gradients: list[SemanticGradient], + policy_set: PolicySet, +) -> dict[tuple[str, str], list[StoredLink]]: + result: dict[tuple[str, str], list[StoredLink]] = defaultdict(list) + seen: set[tuple[tuple[str, str], str, str | None]] = set() + for gradient in gradients: + links = _merge_source_trajectory_links( + [ + *list(getattr(gradient, "links", []) or []), + *_superseded_source_trajectory_links(gradient, policy_set), + ] + ) + if not links: + continue + for key in _gradient_source_keys(gradient, policy_set): + for link in links: + dedupe_key = (key, link.to_uri, link.match_text) + if dedupe_key in seen: + continue + seen.add(dedupe_key) + result[key].append(link) + return dict(result) + + +def _gradient_source_keys( + gradient: SemanticGradient, + policy_set: PolicySet, +) -> list[tuple[str, str]]: + keys: list[tuple[str, str]] = [] + + def add(kind: str, value: Any) -> None: + text = str(value or "").strip() + if text and (kind, text) not in keys: + keys.append((kind, text)) + + add("uri", gradient.target_uri) + add("name", gradient.target_name) + after_file = getattr(gradient, "after_file", None) + if after_file is not None: + add("content", getattr(after_file, "content", "")) + before_file = getattr(gradient, "before_file", None) + if before_file is not None: + add("uri", getattr(before_file, "uri", None)) + fields = getattr(before_file, "extra_fields", {}) or {} + add("name", fields.get("experience_name") or fields.get("name")) + add("content", getattr(before_file, "content", "")) + + superseded_policy = _find_superseded_policy(_gradient_supersedes(gradient), policy_set) + if superseded_policy is not None: + add("uri", superseded_policy.uri) + add("name", superseded_policy.name) + add("content", superseded_policy.content) + + return keys + + +def _superseded_source_trajectory_links( + gradient: SemanticGradient, + policy_set: PolicySet, +) -> list[StoredLink]: + superseded_policy = _find_superseded_policy(_gradient_supersedes(gradient), policy_set) + return _source_trajectory_links_from_experience(superseded_policy) + + +def _upsert_output_count(operations: Any, *, memory_type: str) -> int: + count = 0 + for op in getattr(operations, "upsert_operations", []) or []: + if getattr(op, "memory_type", None) != memory_type: + continue + fields = dict(getattr(op, "memory_fields", {}) or {}) + if str(fields.get("content") or "").strip(): + count += 1 + return count + + +def _replacement_source_uris_by_target(operations: Any) -> dict[str, list[str]]: + replacements = getattr(operations, "delete_replacements", {}) or {} + if not isinstance(replacements, dict): + return {} + result: dict[str, list[str]] = defaultdict(list) + for source_uri, target_uri in replacements.items(): + source = str(source_uri or "").strip() + target = str(target_uri or "").strip() + if not source or not target or source == target: + continue + if source not in result[target]: + result[target].append(source) + return dict(result) + + +def _source_identity_keys( + source_links_by_target: dict[tuple[str, str], list[StoredLink]], +) -> list[tuple[str, str]]: + result: list[tuple[str, str]] = [] + for key in source_links_by_target: + if key[0] in {"uri", "name"} and key not in result: + result.append(key) + return result + + +def _is_source_trajectory_link(link: StoredLink) -> bool: + return ( + link.link_type == "derived_from" + and bool(link.to_uri) + and "/memories/trajectories/" in link.to_uri + ) + + +def _find_policy_by_uri(policy_set: PolicySet, uri: str) -> Policy | None: + for policy in policy_set.policies: + if policy.uri == uri: + return policy + return None + +def _base_version_from_old_file_or_policy( + old_file: Any, target_uri: str | None, policy_set: PolicySet +) -> int | None: + if old_file is not None: + version = safe_int((getattr(old_file, "extra_fields", {}) or {}).get("version")) + if version is not None: + return version + if target_uri: + policy = _find_policy_by_uri(policy_set, target_uri) + return policy.version if policy is not None else None + return None + + +def _gradient_supersedes(gradient: SemanticGradient) -> Any: + metadata = dict(getattr(gradient, "metadata", {}) or {}) + if metadata.get("supersedes") is not None: + return metadata.get("supersedes") + return (gradient.after_file.extra_fields or {}).get("supersedes") + + +def _find_superseded_policy(supersedes: Any, policy_set: PolicySet) -> Policy | None: + names: list[str] + if isinstance(supersedes, str): + names = [supersedes] + elif isinstance(supersedes, list): + names = [str(item) for item in supersedes] + else: + names = [] + for name in names: + for policy in policy_set.policies: + if policy.name == name: + return policy + return None diff --git a/openviking/session/train/components/policy_trainer.py b/openviking/session/train/components/policy_trainer.py new file mode 100644 index 0000000000..06c33e7b75 --- /dev/null +++ b/openviking/session/train/components/policy_trainer.py @@ -0,0 +1,719 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Batch and streaming trainers for session policy optimization. + +The trainers expose rollout-driven training primitives shared by offline and +realtime collection paths. They intentionally reuse the same downstream stages +as ``OfflinePolicyOptimizationPipeline`` while separating trainer concerns from +case rollout execution. +""" + +from __future__ import annotations + +import asyncio +import threading +from dataclasses import dataclass, field +from typing import Any, Hashable + +from openviking.session.memory.utils.streaming_batcher import ( + StreamingBatcher, + StreamingBatcherConfig, +) +from openviking.session.train.context import PipelineContext +from openviking.session.train.domain import ( + ExperienceSet, + PolicyApplyResult, + PolicyUpdatePlan, + Rollout, + RolloutAnalysis, + RolloutTrainingResult, + ScopedRolloutTrainingResult, +) +from openviking.session.train.engine import PolicyTrainingEngine +from openviking.session.train.interfaces import ( + GradientEstimator, + PolicyOptimizer, + PolicyUpdater, + RolloutAnalyzer, + SemanticGradient, +) +from openviking.session.train.utils import average_score, validate_rollouts_have_cases +from openviking.telemetry import tracer +from openviking_cli.utils import get_logger + +logger = get_logger(__name__) + + +@dataclass(slots=True) +class BatchPolicyTrainer: + """Train a policy from an explicit batch of rollout records.""" + + rollout_analyzer: RolloutAnalyzer + gradient_estimator: GradientEstimator + policy_optimizer: PolicyOptimizer + policy_updater: PolicyUpdater + _engine: PolicyTrainingEngine = field(init=False, repr=False) + + def __post_init__(self) -> None: + self._engine = PolicyTrainingEngine( + rollout_analyzer=self.rollout_analyzer, + gradient_estimator=self.gradient_estimator, + policy_optimizer=self.policy_optimizer, + policy_updater=self.policy_updater, + ) + + @tracer("train.batch_policy_trainer.train_rollouts", ignore_result=True, ignore_args=True) + async def train_rollouts( + self, + rollouts: list[Rollout], + policy_set: ExperienceSet, + context: PipelineContext | Any = None, + analyses: list[RolloutAnalysis] | None = None, + ) -> RolloutTrainingResult | ScopedRolloutTrainingResult: + ctx = _coerce_pipeline_context(context) + rollout_list = list(rollouts) + validate_rollouts_have_cases(rollout_list) + if analyses is None: + analyses = await self._engine.analyze_rollouts(rollout_list, ctx) + else: + analyses = list(analyses) + gradients = await self._engine.estimate_gradients(analyses, policy_set, ctx) + plan, apply_result = await self._engine.plan_and_apply( + gradients=gradients, + policy_set=policy_set, + ctx=ctx, + ) + return RolloutTrainingResult( + analyses=analyses, + gradients=gradients, + plan=plan, + apply_result=apply_result, + metadata={ + "policy_set_root_uri": apply_result.updated_policy_set.root_uri, + "rollout_count": len(rollout_list), + "analysis_count": len(analyses), + "gradient_count": len(gradients), + "score": average_score(analyses), + "source": "batch_rollouts", + }, + ) + + +@dataclass(slots=True) +class StreamingPolicyTrainerConfig: + """Configuration for automatic streaming rollout training.""" + + max_gradients_per_update: int = 16 + max_wait_seconds: float = 10.0 + timer_check_interval_seconds: float = 1.0 + trace_console: bool = False + + def __post_init__(self) -> None: + if self.max_gradients_per_update <= 0: + raise ValueError("max_gradients_per_update must be > 0") + if self.max_wait_seconds <= 0: + raise ValueError("max_wait_seconds must be > 0") + if self.timer_check_interval_seconds <= 0: + raise ValueError("timer_check_interval_seconds must be > 0") + + +@dataclass(frozen=True, slots=True) +class StreamingPolicyTrainerKey: + """Process-local registry key for one shared streaming trainer.""" + + account_id: str + user_id: str + policy_root_uri: str + + +@dataclass(slots=True) +class StreamingPolicyTrainer: + """Long-lived rollout trainer that batches concurrent semantic gradients. + + ``submit_rollout`` analyzes a rollout and estimates gradients immediately, + then appends those gradients to an in-memory buffer. A policy update is + automatically triggered either when the buffer reaches + ``max_gradients_per_update`` or when the oldest buffered gradient waits at + least ``max_wait_seconds``. If the submitting call triggers a count-based + flush, it waits for optimizer/apply completion before returning. + """ + + policy_set: ExperienceSet + rollout_analyzer: RolloutAnalyzer + gradient_estimator: GradientEstimator + policy_optimizer: PolicyOptimizer + policy_updater: PolicyUpdater + context: PipelineContext | Any = None + config: StreamingPolicyTrainerConfig = field(default_factory=StreamingPolicyTrainerConfig) + _core: PolicyTrainingEngine = field(init=False, repr=False) + _batcher: StreamingBatcher[_BufferedRolloutTraining, RolloutTrainingResult] = field( + init=False, repr=False + ) + _last_apply_result: PolicyApplyResult | None = field(init=False, default=None, repr=False) + _closed: bool = field(init=False, default=False, repr=False) + + def __post_init__(self) -> None: + self.context = _coerce_pipeline_context(self.context) + self._core = PolicyTrainingEngine( + rollout_analyzer=self.rollout_analyzer, + gradient_estimator=self.gradient_estimator, + policy_optimizer=self.policy_optimizer, + policy_updater=self.policy_updater, + ) + self._batcher = StreamingBatcher( + name="openviking-streaming-policy-trainer", + process_batch=self._process_batch, + config=StreamingBatcherConfig( + max_items_per_batch=self.config.max_gradients_per_update, + max_wait_seconds=self.config.max_wait_seconds, + timer_check_interval_seconds=self.config.timer_check_interval_seconds, + ), + item_size=lambda item: len(item.gradients), + result_metadata=lambda result: result.metadata, + ) + self._last_apply_result: PolicyApplyResult | None = None + self._closed = False + + @property + def last_apply_result(self) -> PolicyApplyResult | None: + return self._last_apply_result + + async def get_buffered_gradient_count(self) -> int: + """Return the current buffered gradient count under the buffer lock.""" + + return await self._batcher.get_buffered_size() + + @property + def closed(self) -> bool: + return self._closed + + async def close(self) -> RolloutTrainingResult | None: + """Stop the timer task and flush any buffered gradients once. + + ``close`` is idempotent. The first call cancels the background timer + and applies any remaining buffered gradients with ``flush_reason="close"``. + Subsequent calls are no-ops and return ``None``. + """ + + if self._closed: + return None + self._closed = True + return await self._batcher.close() + + @tracer("train.streaming_policy_trainer.submit_rollout", ignore_result=True, ignore_args=True) + async def submit_rollout( + self, + rollout: Rollout, + ) -> RolloutTrainingResult | ScopedRolloutTrainingResult: + """Submit one realtime rollout and wait for its batch update result. + + The rollout is analyzed and converted to gradients immediately, then + buffered by the shared count/time window. This method returns only + after the batch containing this rollout has been optimized and applied. + """ + + if self._closed: + raise RuntimeError("StreamingPolicyTrainer is closed") + validate_rollouts_have_cases([rollout]) + analysis = await self.rollout_analyzer.analyze(rollout, self.context.analysis_context) + gradients = await self.gradient_estimator.estimate( + analysis, + self.policy_set, + self.context.gradient_context, + ) + tracer.info( + "StreamingPolicyTrainer buffered rollout " + f"rollout_case={rollout.case.name} " + f"new_gradients={len(gradients)}", + console=self.config.trace_console, + ) + buffered = _BufferedRolloutTraining( + gradients=list(gradients), + analysis=analysis, + rollout=rollout, + ) + result = await self._batcher.submit(buffered) + self._last_apply_result = result.apply_result + scoped_result = _scope_training_result_to_submitter(result, buffered) + tracer.info( + "StreamingPolicyTrainer submit finished " + f"batch_id={result.metadata.get('batch_id')} " + f"batch_trace_id={result.metadata.get('batch_trace_id')} " + f"flush_reason={result.metadata.get('flush_reason')} " + f"rollout_count={result.metadata.get('rollout_count')} " + f"gradient_count={result.metadata.get('gradient_count')} " + f"written_uris={result.apply_result.written_uris} " + f"errors={result.apply_result.errors}", + console=self.config.trace_console, + ) + return scoped_result + + @tracer("train.streaming_policy_trainer.submit_gradients", ignore_result=True, ignore_args=True) + async def submit_gradients( + self, + gradients: list[SemanticGradient], + *, + analysis: RolloutAnalysis | None = None, + rollout: Rollout | None = None, + ) -> RolloutTrainingResult | ScopedRolloutTrainingResult: + """Submit pre-computed gradients directly to the streaming trainer. + + Unlike ``submit_rollout``, this method skips analysis and gradient + estimation. It is useful for memory types whose gradients are + produced during an earlier stage (e.g. session skills co-extracted + during trajectory analysis). + """ + if self._closed: + raise RuntimeError("StreamingPolicyTrainer is closed") + if not gradients: + # No gradients to submit — return a no-op result immediately. + return RolloutTrainingResult( + analyses=[analysis] if analysis is not None else [], + gradients=[], + plan=PolicyUpdatePlan(items=[], metadata={"no_op": True}), + apply_result=PolicyApplyResult( + updated_policy_set=self.policy_set, + written_uris=[], + errors=[], + metadata={"no_op": True}, + ), + metadata={"no_op": True, "gradient_count": 0}, + ) + tracer.info( + "StreamingPolicyTrainer buffered gradients " + f"new_gradients={len(gradients)}", + console=self.config.trace_console, + ) + buffered = _BufferedRolloutTraining( + gradients=list(gradients), + analysis=analysis, + rollout=rollout, + ) + result = await self._batcher.submit(buffered) + self._last_apply_result = result.apply_result + return _scope_training_result_to_submitter(result, buffered) + + + @tracer("train.streaming_policy_trainer.train_rollouts", ignore_result=True, ignore_args=True) + async def train_rollouts( + self, + rollouts: list[Rollout], + policy_set: ExperienceSet, + context: PipelineContext | Any = None, + analyses: list[RolloutAnalysis] | None = None, + ) -> RolloutTrainingResult: + del policy_set, context, analyses + results = await asyncio.gather(*[self.submit_rollout(rollout) for rollout in rollouts]) + return _combine_training_results(results, source="streaming_rollouts") + + async def _process_batch( + self, + items: list["_BufferedRolloutTraining"], + reason: str, + ) -> RolloutTrainingResult: + chunks = _chunks_buffered_items_by_gradient_count( + items, + self.config.max_gradients_per_update, + ) + gradients = [gradient for chunk in chunks for gradient in chunk.gradients] + analyses = _unique_by_identity( + [item.analysis for item in items if item.analysis is not None] + ) + rollouts = _unique_by_identity( + [item.rollout for item in items if item.rollout is not None] + ) + tracer.info( + "StreamingPolicyTrainer flush started " + f"reason={reason} " + f"rollout_count={len(rollouts)} " + f"analysis_count={len(analyses)} " + f"gradient_count={len(gradients)}", + console=self.config.trace_console, + ) + + plans: list[PolicyUpdatePlan] = [] + apply_results: list[PolicyApplyResult] = [] + for chunk_index, chunk in enumerate(chunks): + gradient_chunk = chunk.gradients + tracer.info( + "StreamingPolicyTrainer flush chunk started " + f"reason={reason} " + f"chunk_index={chunk_index} " + f"gradient_count={len(gradient_chunk)}", + console=self.config.trace_console, + ) + plan, apply_result = await self._core.plan_and_apply( + gradients=gradient_chunk, + policy_set=self.policy_set, + ctx=self.context, + ) + self.policy_set = apply_result.updated_policy_set + plans.append(plan) + apply_results.append(apply_result) + tracer.info( + "StreamingPolicyTrainer flush chunk finished " + f"reason={reason} " + f"chunk_index={chunk_index} " + f"written_uris={apply_result.written_uris} " + f"errors={apply_result.errors}", + console=self.config.trace_console, + ) + + plan = _combine_update_plans(plans) + apply_result = _combine_apply_results(apply_results, fallback_policy_set=self.policy_set) + self.policy_set = apply_result.updated_policy_set + self._last_apply_result = apply_result + chunk_gradient_counts = [len(chunk.gradients) for chunk in chunks] + result = RolloutTrainingResult( + analyses=analyses, + gradients=gradients, + plan=plan, + apply_result=apply_result, + metadata={ + "policy_set_root_uri": apply_result.updated_policy_set.root_uri, + "rollout_count": len(rollouts), + "analysis_count": len(analyses), + "gradient_count": len(gradients), + "chunk_count": len(chunk_gradient_counts), + "chunk_gradient_counts": chunk_gradient_counts, + "score": average_score(analyses), + "source": "streaming_rollouts", + "flush_reason": reason, + }, + ) + tracer.info( + "StreamingPolicyTrainer flush finished " + f"reason={reason} " + f"chunk_count={len(chunk_gradient_counts)} " + f"written_uris={apply_result.written_uris} " + f"errors={apply_result.errors}", + console=self.config.trace_console, + ) + return result + + +def _chunks_buffered_items_by_gradient_count( + items: list["_BufferedRolloutTraining"], + size: int, +) -> list["_BufferedRolloutTrainingChunk"]: + if size <= 0: + raise ValueError("chunk size must be > 0") + + all_gradients: list[SemanticGradient] = [] + for item in items: + all_gradients.extend(item.gradients) + + if not all_gradients: + return [_BufferedRolloutTrainingChunk(gradients=[])] + + chunks: list[_BufferedRolloutTrainingChunk] = [] + for start in range(0, len(all_gradients), size): + chunk_gradients = all_gradients[start : start + size] + chunks.append(_BufferedRolloutTrainingChunk(gradients=chunk_gradients)) + return chunks + + +def _combine_update_plans(plans: list[PolicyUpdatePlan]) -> PolicyUpdatePlan: + if not plans: + return PolicyUpdatePlan(items=[], metadata={"chunk_count": 0}) + items = [item for plan in plans for item in plan.items] + return PolicyUpdatePlan( + items=items, + metadata={ + "chunk_count": len(plans), + "chunk_item_counts": [len(plan.items) for plan in plans], + "chunks": [dict(plan.metadata or {}) for plan in plans], + }, + ) + + +def _combine_apply_results( + results: list[PolicyApplyResult], + *, + fallback_policy_set: ExperienceSet, +) -> PolicyApplyResult: + if not results: + return PolicyApplyResult(updated_policy_set=fallback_policy_set) + return PolicyApplyResult( + updated_policy_set=results[-1].updated_policy_set, + written_uris=[uri for result in results for uri in result.written_uris], + deleted_uris=[uri for result in results for uri in result.deleted_uris], + errors=[error for result in results for error in result.errors], + metadata={ + "chunk_count": len(results), + "chunk_metadata": [dict(result.metadata or {}) for result in results], + }, + ) + + +def _scope_training_result_to_submitter( + result: RolloutTrainingResult, + submitter: "_BufferedRolloutTraining", +) -> RolloutTrainingResult | ScopedRolloutTrainingResult: + """Return the submitting rollout's view of a shared streaming flush. + + StreamingBatcher intentionally gives every waiter the same batch result. + For per-commit consumers (memory_diff/case links), exposing all batch plan + items would make one trace appear to add every other concurrently flushed + experience. Keep the full batch result available via ``batch_result`` but + scope the top-level fields to the submitter's analyses and source + trajectories. + """ + + analysis = submitter.analysis + if analysis is None: + return result + + scoped_plan = _scope_plan_to_analysis( + result.plan, + analysis=analysis, + apply_result=result.apply_result, + ) + scoped_apply_result = _scope_apply_result_to_plan( + result.apply_result, + scoped_plan, + ) + metadata = dict(result.metadata or {}) + metadata.update( + { + "batch_rollout_count": metadata.get("rollout_count"), + "batch_analysis_count": metadata.get("analysis_count"), + "batch_gradient_count": metadata.get("gradient_count"), + "rollout_count": 1 if submitter.rollout is not None else 0, + "analysis_count": 1, + "gradient_count": len(submitter.gradients), + "source": "streaming_rollouts_scoped", + "scoped_to_submitter": True, + } + ) + return ScopedRolloutTrainingResult( + analyses=[analysis], + gradients=list(submitter.gradients), + plan=scoped_plan, + apply_result=scoped_apply_result, + batch_result=result, + metadata=metadata, + ) + + +def _scope_plan_to_analysis( + plan: PolicyUpdatePlan, + *, + analysis: RolloutAnalysis, + apply_result: PolicyApplyResult, +) -> PolicyUpdatePlan: + trajectory_uris = _analysis_trajectory_uris(analysis) + scoped_items = [ + item + for item in list(getattr(plan, "items", []) or []) + if _plan_item_belongs_to_trajectories( + item, + trajectory_uris=trajectory_uris, + ) + ] + metadata = dict(getattr(plan, "metadata", {}) or {}) + metadata.update( + { + "scoped_to_trajectory_uris": sorted(trajectory_uris), + "unscoped_item_count": len(getattr(plan, "items", []) or []), + } + ) + return PolicyUpdatePlan(items=scoped_items, metadata=metadata) + + +def _plan_item_belongs_to_trajectories( + item: Any, + *, + trajectory_uris: set[str], +) -> bool: + if not trajectory_uris: + return False + for link in getattr(item, "links", []) or []: + try: + if hasattr(link, "to_uri"): + to_uri = str(getattr(link, "to_uri", "") or "") + link_type = str(getattr(link, "link_type", "") or "") + elif isinstance(link, dict): + to_uri = str(link.get("to_uri") or "") + link_type = str(link.get("link_type") or "") + else: + continue + except Exception: + continue + if link_type == "derived_from" and to_uri in trajectory_uris: + return True + # Deletes may not carry fresh links when a merged replacement owns the + # source trajectory links. Keep only upserts in submitter-scoped views. + return False + + +def _scope_apply_result_to_plan( + apply_result: PolicyApplyResult, + plan: PolicyUpdatePlan, +) -> PolicyApplyResult: + plan_uris = { + _plan_item_uri(item, getattr(apply_result.updated_policy_set, "root_uri", "")) + for item in getattr(plan, "items", []) or [] + } + metadata = dict(getattr(apply_result, "metadata", {}) or {}) + metadata.update( + { + "unscoped_written_uris": list(getattr(apply_result, "written_uris", []) or []), + "unscoped_deleted_uris": list(getattr(apply_result, "deleted_uris", []) or []), + } + ) + return PolicyApplyResult( + updated_policy_set=apply_result.updated_policy_set, + written_uris=[uri for uri in getattr(apply_result, "written_uris", []) or [] if uri in plan_uris], + deleted_uris=[uri for uri in getattr(apply_result, "deleted_uris", []) or [] if uri in plan_uris], + errors=list(getattr(apply_result, "errors", []) or []), + metadata=metadata, + ) + + +def _analysis_trajectory_uris(analysis: RolloutAnalysis) -> set[str]: + return { + str(getattr(trajectory, "uri", "") or "") + for trajectory in getattr(analysis, "trajectories", []) or [] + if str(getattr(trajectory, "uri", "") or "") + } + + +def _plan_item_uri(item: Any, root_uri: str) -> str: + uri = str(getattr(item, "target_uri", "") or "") + if uri: + return uri + name = str(getattr(item, "target_name", "") or "new_experience") + return f"{root_uri.rstrip('/')}/{_safe_policy_filename(name)}.md" + + +def _safe_policy_filename(name: str) -> str: + import re + + filename = re.sub(r"[^a-zA-Z0-9_.-]+", "_", name.strip()).strip("._-") + return filename or "new_experience" + + +@dataclass(slots=True) +class _BufferedRolloutTraining: + gradients: list[SemanticGradient] + analysis: RolloutAnalysis | None = None + rollout: Rollout | None = None + + +@dataclass(slots=True) +class _BufferedRolloutTrainingChunk: + gradients: list[SemanticGradient] + + +_streaming_policy_trainer_registry: dict[Hashable, StreamingPolicyTrainer] = {} +_streaming_policy_trainer_registry_lock = threading.RLock() + + +async def get_streaming_policy_trainer( + *, + key: StreamingPolicyTrainerKey | Hashable, + policy_set: ExperienceSet, + rollout_analyzer: RolloutAnalyzer, + gradient_estimator: GradientEstimator, + policy_optimizer: PolicyOptimizer, + policy_updater: PolicyUpdater, + context: PipelineContext | Any = None, + config: StreamingPolicyTrainerConfig | None = None, +) -> StreamingPolicyTrainer: + """Get or create the process-global streaming trainer for one policy key.""" + + with _streaming_policy_trainer_registry_lock: + existing = _streaming_policy_trainer_registry.get(key) + if existing is not None: + return existing + trainer = StreamingPolicyTrainer( + policy_set=policy_set, + rollout_analyzer=rollout_analyzer, + gradient_estimator=gradient_estimator, + policy_optimizer=policy_optimizer, + policy_updater=policy_updater, + context=context, + config=config or StreamingPolicyTrainerConfig(), + ) + _streaming_policy_trainer_registry[key] = trainer + return trainer + + +def make_streaming_policy_trainer_key( + *, + policy_root_uri: str, + request_context: Any, +) -> StreamingPolicyTrainerKey: + """Build the default registry key from policy root and request context.""" + + user = getattr(request_context, "user", None) + account_id = ( + getattr(request_context, "account_id", None) + or getattr(user, "account_id", None) + or "default" + ) + user_id = getattr(request_context, "user_id", None) or getattr(user, "user_id", None) or "" + return StreamingPolicyTrainerKey( + account_id=str(account_id), + user_id=str(user_id), + policy_root_uri=policy_root_uri, + ) + + +def _coerce_pipeline_context(context: PipelineContext | Any = None) -> PipelineContext: + return context if isinstance(context, PipelineContext) else PipelineContext() + + +def _unique_by_identity(items: list[Any]) -> list[Any]: + seen: set[int] = set() + unique = [] + for item in items: + item_id = id(item) + if item_id in seen: + continue + seen.add(item_id) + unique.append(item) + return unique + + +def _combine_training_results( + results: list[RolloutTrainingResult], + *, + source: str, +) -> RolloutTrainingResult: + if not results: + return RolloutTrainingResult( + analyses=[], + gradients=[], + plan=PolicyUpdatePlan(metadata={"empty": True}), + apply_result=PolicyApplyResult(updated_policy_set=ExperienceSet(root_uri="", policies=[])), + metadata={ + "source": source, + "rollout_count": 0, + "analysis_count": 0, + "gradient_count": 0, + }, + ) + + last = results[-1] + last_unscoped = getattr(last, "batch_result", last) + analyses = _unique_by_identity([analysis for result in results for analysis in result.analyses]) + gradients = [gradient for result in results for gradient in result.gradients] + metadata = dict(last_unscoped.metadata) + metadata.update( + { + "source": source, + "rollout_count": len(analyses), + "analysis_count": len(analyses), + "gradient_count": len(gradients), + "score": average_score(analyses), + } + ) + return RolloutTrainingResult( + analyses=analyses, + gradients=gradients, + plan=last_unscoped.plan, + apply_result=last_unscoped.apply_result, + metadata=metadata, + ) diff --git a/openviking/session/train/components/policy_updater.py b/openviking/session/train/components/policy_updater.py new file mode 100644 index 0000000000..a43f3ac0e1 --- /dev/null +++ b/openviking/session/train/components/policy_updater.py @@ -0,0 +1,371 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""PolicyUpdater component implementations.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Any + +from openviking.session.memory.dataclass import ( + MemoryFile, + ResolvedOperation, + ResolvedOperations, + StoredLink, +) +from openviking.session.memory.memory_type_registry import create_default_registry +from openviking.session.memory.memory_updater import MemoryUpdater +from openviking.session.train.domain import ( + Policy, + PolicyApplyResult, + PolicyPlanItem, + PolicySet, + PolicyUpdatePlan, +) +from openviking.storage.viking_fs import get_viking_fs +from openviking.telemetry import tracer + +_EXPERIENCE_NAME_RE = re.compile(r"[^a-zA-Z0-9_.-]+") + + +@dataclass(slots=True) +class DryRunPolicyUpdater: + """PolicyUpdater that records a plan without writing files. + + Unlike a pure no-op, this updater simulates executable plan items into an + updated ExperienceSet snapshot, which makes tests and offline review useful + before enabling a writing updater. + """ + + simulate: bool = True + + @tracer("train.policy_updater.dry_run.apply", ignore_result=True, ignore_args=True) + async def apply( + self, + plan: PolicyUpdatePlan, + policy_set: PolicySet, + context: Any = None, + *, + transaction_handle: Any = None, + ) -> PolicyApplyResult: + del transaction_handle + del context + updated_policy_set = ( + _apply_items_to_snapshot(plan.items, policy_set) + if self.simulate and plan.items + else policy_set + ) + return PolicyApplyResult( + updated_policy_set=updated_policy_set, + written_uris=[], + metadata={ + "dry_run": True, + "simulated": self.simulate, + "plan": plan.metadata, + "item_count": len(plan.items), + }, + ) + + +@dataclass(slots=True) +class MemoryFilePolicyUpdater: + """PolicyUpdater that writes policy files via VikingFS. + + It consumes executable ``upsert`` and ``delete`` plan items. The updater + performs a lightweight base-content guard when ``before_content`` is + available to avoid blindly overwriting or deleting a diverged policy set + snapshot. + """ + + viking_fs: Any = None + vikingdb: Any = None + + @tracer("train.policy_updater.memory_file.apply", ignore_result=True, ignore_args=True) + async def apply( + self, + plan: PolicyUpdatePlan, + policy_set: PolicySet, + context: Any = None, + *, + transaction_handle: Any = None, + ) -> PolicyApplyResult: + viking_fs = self.viking_fs or get_viking_fs() + if viking_fs is None: + raise RuntimeError("VikingFS is required to apply policy update plans") + + updated_policy_set = _apply_items_to_snapshot(plan.items, policy_set) + operations, preflight_errors = _plan_to_resolved_operations( + plan=plan, + policy_set=policy_set, + updated_policy_set=updated_policy_set, + ) + updater = MemoryUpdater( + registry=create_default_registry(), + vikingdb=self.vikingdb, + transaction_handle=transaction_handle, + ) + updater._viking_fs = viking_fs + + apply_result = await updater.apply_operations( + operations, + context, + extract_context=None, + isolation_handler=None, + ) + errors = [*preflight_errors, *[f"{uri}: {exc}" for uri, exc in apply_result.errors]] + + return PolicyApplyResult( + updated_policy_set=updated_policy_set if not errors else policy_set, + written_uris=list(apply_result.written_uris + apply_result.edited_uris), + deleted_uris=list(apply_result.deleted_uris), + errors=errors, + metadata={ + "dry_run": False, + "item_count": len(plan.items), + "operation_upsert_count": len(operations.upsert_operations), + "operation_delete_count": len(operations.delete_file_contents), + }, + ) + + +def _apply_items_to_snapshot( + items: list[PolicyPlanItem], policy_set: PolicySet +) -> PolicySet: + policies_by_uri = {policy.uri: policy for policy in policy_set.policies} + result = list(policy_set.policies) + + for item in items: + uri = _target_uri(item, policy_set.root_uri) + + if item.kind == "delete": + existing = policies_by_uri.get(uri) or _find_policy( + PolicySet( + policy_set.root_uri, + result, + metadata=dict(policy_set.metadata), + viking_fs=policy_set.viking_fs, + request_context=policy_set.request_context, + ), + uri=None, + name=item.target_name, + ) + remove_uri = existing.uri if existing is not None else uri + result = [ + policy + for policy in result + if policy.uri != remove_uri and policy.name != item.target_name + ] + policies_by_uri.pop(remove_uri, None) + policies_by_uri.pop(uri, None) + continue + + if item.kind != "upsert" or item.after_content is None: + continue + existing = policies_by_uri.get(uri) or _find_policy( + PolicySet( + policy_set.root_uri, + result, + metadata=dict(policy_set.metadata), + viking_fs=policy_set.viking_fs, + request_context=policy_set.request_context, + ), + uri=None, + name=item.target_name, + ) + metadata = dict(existing.metadata) if existing is not None else {} + metadata.update(item.metadata.get("patch_metadata", {})) + metadata.setdefault("memory_type", item.memory_type or "experiences") + metadata["experience_name"] = item.target_name + version = (existing.version + 1) if existing is not None else 1 + updated = Policy( + name=item.target_name, + uri=uri, + version=version, + status=(existing.status if existing is not None else "draft"), + content=item.after_content, + metadata=metadata, + links=list(existing.links or []) if existing is not None else [], + backlinks=list(existing.backlinks or []) if existing is not None else [], + ) + if existing is None: + result.append(updated) + else: + result = [updated if policy.uri == existing.uri else policy for policy in result] + policies_by_uri[uri] = updated + + result.sort(key=lambda policy: policy.uri) + return PolicySet( + root_uri=policy_set.root_uri, + policies=result, + metadata=dict(policy_set.metadata), + viking_fs=policy_set.viking_fs, + request_context=policy_set.request_context, + ) + + +def _find_policy( + policy_set: PolicySet, + *, + uri: str | None, + name: str, +) -> Policy | None: + for policy in policy_set.policies: + if uri and policy.uri == uri: + return policy + if not uri and policy.name == name: + return policy + return None + + +def _target_uri(item: PolicyPlanItem, root_uri: str) -> str: + if item.target_uri: + return item.target_uri + return f"{root_uri.rstrip('/')}/{_safe_experience_filename(item.target_name)}.md" + + + +def _plan_to_resolved_operations( + *, + plan: PolicyUpdatePlan, + policy_set: PolicySet, + updated_policy_set: PolicySet, +) -> tuple[ResolvedOperations, list[str]]: + upserts: list[ResolvedOperation] = [] + deletes: list[MemoryFile] = [] + links: list[StoredLink] = [] + errors: list[str] = [] + + for item in plan.items: + uri = _target_uri(item, policy_set.root_uri) + current = _find_policy(policy_set, uri=uri, name=item.target_name) + if ( + current is not None + and item.before_content is not None + and _normalize_guard_content(current.content) + != _normalize_guard_content(item.before_content) + ): + errors.append( + "base content mismatch for " + f"{item.target_name}: expected gradient before_content" + ) + continue + + if item.kind == "delete": + deletes.append(_policy_or_plan_item_memory_file(item, uri=uri, current=current)) + continue + + if item.kind != "upsert": + continue + if item.after_content is None: + errors.append(f"missing after_content for {item.target_name}") + continue + + updated = _find_policy(updated_policy_set, uri=uri, name=item.target_name) + if updated is None: + errors.append( + f"planned policy not found after simulation: {item.target_name}" + ) + continue + + upserts.append( + ResolvedOperation( + old_memory_file_content=_policy_to_memory_file(current) + if current is not None + else None, + memory_fields={ + **dict(updated.metadata), + "memory_type": item.memory_type or "experiences", + "experience_name": updated.name, + "content": updated.content, + "status": updated.status, + }, + memory_type=item.memory_type or "experiences", + uris=[uri], + ) + ) + links.extend(_source_trajectory_links(exp_uri=uri, links=item.links)) + + return ( + ResolvedOperations( + upsert_operations=upserts, + delete_file_contents=deletes, + errors=[], + resolved_links=links, + ), + errors, + ) + + +def _policy_or_plan_item_memory_file( + item: PolicyPlanItem, + *, + uri: str, + current: Policy | None, +) -> MemoryFile: + if current is not None: + return _policy_to_memory_file(current) + return MemoryFile( + uri=uri, + content=item.before_content or "", + memory_type=item.memory_type or "experiences", + extra_fields={ + "memory_type": item.memory_type or "experiences", + "experience_name": item.target_name, + **({"version": item.base_version} if item.base_version is not None else {}), + }, + ) + + +def _policy_to_memory_file(policy: Policy | None) -> MemoryFile | None: + if policy is None: + return None + return MemoryFile( + uri=policy.uri, + content=policy.content, + links=list(policy.links or []), + backlinks=list(policy.backlinks or []), + memory_type="experiences", + extra_fields={ + **dict(policy.metadata), + "memory_type": "experiences", + "experience_name": policy.name, + "version": policy.version, + "status": policy.status, + }, + ) + + +def _source_trajectory_links( + *, + exp_uri: str, + links: list[StoredLink], +) -> list[StoredLink]: + result: list[StoredLink] = [] + seen: set[tuple[str, str | None]] = set() + for link in links or []: + if ( + link.link_type != "derived_from" + or not link.to_uri + or "/memories/trajectories/" not in link.to_uri + ): + continue + key = (link.to_uri, link.match_text) + if key in seen: + continue + seen.add(key) + update = {"from_uri": exp_uri, "match_text": None, "description": ""} + if not link.created_at: + update["created_at"] = datetime.now(timezone.utc).isoformat() + result.append(link.model_copy(update=update)) + return result + + +def _safe_experience_filename(name: str) -> str: + filename = _EXPERIENCE_NAME_RE.sub("_", name.strip()).strip("._-") + return filename or "new_experience" + + +def _normalize_guard_content(content: str) -> str: + return content.strip() diff --git a/openviking/session/train/components/progress.py b/openviking/session/train/components/progress.py new file mode 100644 index 0000000000..fafb98fa93 --- /dev/null +++ b/openviking/session/train/components/progress.py @@ -0,0 +1,379 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Terminal progress helpers for train components.""" + +from __future__ import annotations + +import sys +import time +from dataclasses import dataclass, field +from typing import Any, Callable + +try: # pragma: no cover - exercised through integration/TTY usage + from rich.console import Console + from rich.progress import Progress, ProgressColumn, Task, TextColumn + from rich.table import Column + from rich.text import Text +except Exception: # pragma: no cover - fallback for minimal environments + Console = None + Progress = None + ProgressColumn = object # type: ignore[assignment,misc] + Task = Any # type: ignore[misc,assignment] + TextColumn = None + Column = None + Text = None + + +class ThreeStateBarColumn(ProgressColumn): + """Render successful/failed/running/pending work as a rich progress bar.""" + + def __init__(self, bar_width: int = 42) -> None: + table_column = Column(ratio=1, no_wrap=True) if Column is not None else None + super().__init__(table_column=table_column) + self.bar_width = bar_width + + def render(self, task: Task) -> Text: + bar_width = self.bar_width or 40 + total = max(int(task.total or 0), 0) + processed = max(int(task.completed or 0), 0) + failed = max(int(task.fields.get("failed", 0) or 0), 0) + succeeded = max(int(task.fields.get("succeeded", processed - failed) or 0), 0) + running = max(int(task.fields.get("running", 0) or 0), 0) + + if total <= 0: + return Text("─" * bar_width, style="bar.back") + + succeeded = min(succeeded, total) + failed = min(failed, total - succeeded) + running = min(running, total - succeeded - failed) + succeeded_width = int(bar_width * succeeded / total) + failed_width = int(bar_width * (succeeded + failed) / total) - succeeded_width + running_width = ( + int(bar_width * (succeeded + failed + running) / total) - succeeded_width - failed_width + ) + pending_width = bar_width - succeeded_width - failed_width - running_width + + bar = Text() + if succeeded_width > 0: + bar.append("█" * succeeded_width, style="green") + if failed_width > 0: + bar.append("█" * failed_width, style="red") + if running_width > 0: + bar.append("▓" * running_width, style="yellow") + if pending_width > 0: + bar.append("░" * pending_width, style="dim") + return bar + + +class ProgressSummaryColumn(ProgressColumn): + """Render processed count plus non-zero active/failure counts.""" + + def render(self, task: Task) -> Text: + failed = max(int(task.fields.get("failed", 0) or 0), 0) + running = max(int(task.fields.get("running", 0) or 0), 0) + + summary = Text("(") + summary.append(f"{int(task.completed or 0)}/{int(task.total or 0)}") + if failed > 0: + summary.append(", ") + summary.append(f"{failed} failed", style="bold red") + if running > 0: + summary.append(", ") + summary.append(f"{running} running", style="bold yellow") + summary.append(")") + return summary + + +@dataclass(slots=True) +class ProgressPrinter: + """Render a four-state progress indicator for train components. + + The public API intentionally matches the previous lightweight printer so + callers only need ``render()``, ``start_one()``, ``complete_one()`` and + ``finish()``. ``fail_one()`` records failed work in red. Rich rendering is + used for interactive terminals; a compact text fallback is kept for minimal + environments. + """ + + total: int + label: str + enabled: bool + description: str = "" + pending: int = 0 + running: int = 0 + completed: int = 0 + failed: int = 0 + _finished: bool = False + _use_rich: bool = field(init=False, default=False) + _progress: Any = field(init=False, default=None) + _task_id: Any = field(init=False, default=None) + _started: bool = field(init=False, default=False) + _description_printed: bool = field(init=False, default=False) + _started_at: float | None = field(init=False, default=None) + + def __post_init__(self) -> None: + if self.pending == 0 and self.running == 0 and self.completed == 0 and self.failed == 0: + self.pending = max(0, self.total) + self._use_rich = bool( + self.enabled + and self.total > 0 + and Progress is not None + and Console is not None + and TextColumn is not None + and sys.stderr.isatty() + ) + + def render(self) -> None: + if not self.enabled or self.total <= 0: + return + self._mark_started() + self._print_description_once() + if self._use_rich: + self._ensure_rich_started() + self._update_rich() + return + self._write_text() + + def start_one(self) -> None: + if not self.enabled or self.total <= 0 or self._finished: + return + if self.pending > 0: + self.pending -= 1 + self.running += 1 + self._write() + + def complete_one(self) -> None: + if not self.enabled or self.total <= 0 or self._finished: + return + if self.running > 0: + self.running -= 1 + elif self.pending > 0: + self.pending -= 1 + self.completed = min(self.total, self.completed + 1) + self._write() + + def fail_one(self) -> None: + if not self.enabled or self.total <= 0 or self._finished: + return + if self.running > 0: + self.running -= 1 + elif self.pending > 0: + self.pending -= 1 + self.failed = min(self.total - self.completed, self.failed + 1) + self._write() + + def advance(self) -> None: + """Compatibility helper for callers that only track completion.""" + self.complete_one() + + def finish(self) -> None: + if not self.enabled or self.total <= 0 or self._finished: + return + self._finished = True + if self._use_rich: + self._update_rich() + if self._progress is not None and self._started: + self._progress.stop() + return + self._write_text(newline=True) + + def _write(self) -> None: + self._mark_started() + self._print_description_once() + if self._use_rich: + self._ensure_rich_started() + self._update_rich() + return + self._write_text() + + def _mark_started(self) -> None: + if self._started_at is None: + self._started_at = time.monotonic() + + def _print_description_once(self) -> None: + if self._description_printed or not self.description: + return + self._description_printed = True + if self._use_rich: + line = Text() + line.append(format_label(self.label), style=label_style(self.label)) + line.append(f" {self.description}") + Console(stderr=True, soft_wrap=False).print(line) + return + sys.stdout.write(f"[{self.label}] {self.description}\n") + sys.stdout.flush() + + def _elapsed_seconds(self) -> float: + if self._started_at is None: + return 0.0 + return max(0.0, time.monotonic() - self._started_at) + + def _ensure_rich_started(self) -> None: + if self._progress is not None: + return + self._progress = Progress( + ThreeStateBarColumn(), + TextColumn("[progress.percentage]{task.percentage:>3.0f}%"), + ProgressSummaryColumn(), + console=Console(stderr=True, soft_wrap=False), + transient=False, + ) + self._task_id = self._progress.add_task( + format_label(self.label), + total=self.total, + running=0, + succeeded=0, + failed=0, + ) + self._progress.start() + self._started = True + + def _update_rich(self) -> None: + if self._progress is None: + return + self._progress.update( + self._task_id, + total=self.total, + completed=self.completed + self.failed, + succeeded=self.completed, + failed=self.failed, + running=self.running, + ) + + def _write_text(self, *, newline: bool = False) -> None: + width = 24 + succeeded_width, failed_width, running_width, pending_width = _state_widths( + pending=self.pending, + running=self.running, + completed=self.completed, + failed=self.failed, + total=self.total, + width=width, + ) + bar = "C" * succeeded_width + "F" * failed_width + "R" * running_width + "P" * pending_width + processed = self.completed + self.failed + percent = (processed / self.total) * 100.0 + suffix = "\n" if newline else "" + sys.stdout.write( + f"\r[{self.label}] [{bar}] {percent:6.2f}% " + f"{_progress_summary(processed=processed, total=self.total, failed=self.failed, running=self.running)}" + f"{suffix}" + ) + sys.stdout.flush() + + +def _state_widths( + *, + pending: int, + running: int, + completed: int, + failed: int, + total: int, + width: int, +) -> tuple[int, int, int, int]: + values = [completed, failed, running, pending] + if total <= 0 or width <= 0: + return 0, 0, 0 + + exact = [(value / total) * width if value > 0 else 0.0 for value in values] + widths = [int(value) for value in exact] + + for idx, value in enumerate(values): + if value > 0 and widths[idx] == 0: + widths[idx] = 1 + + while sum(widths) > width: + candidates = [ + idx for idx, value in enumerate(values) if widths[idx] > (1 if value > 0 else 0) + ] + if not candidates: + break + idx = max(candidates, key=lambda item: widths[item]) + widths[idx] -= 1 + + while sum(widths) < width: + idx = max(range(len(values)), key=lambda item: exact[item] - int(exact[item])) + widths[idx] += 1 + + return widths[0], widths[1], widths[2], widths[3] + + +def _progress_summary(*, processed: int, total: int, failed: int, running: int) -> str: + extras: list[str] = [] + if failed > 0: + extras.append(f"{failed} failed") + if running > 0: + extras.append(f"{running} running") + if extras: + return f"({processed}/{total}, {', '.join(extras)})" + return f"({processed}/{total})" + + +def format_duration(seconds: float) -> str: + total_seconds = max(0, int(round(seconds))) + hours, remainder = divmod(total_seconds, 3600) + minutes, secs = divmod(remainder, 60) + if hours: + return f"{hours}h{minutes}m{secs}s" + if minutes: + return f"{minutes}m{secs}s" + return f"{secs}s" + + +def format_label(label: str) -> str: + return f"[{label.upper()}]" + + +def label_style(label: str) -> str: + if label.endswith("_start"): + return "bold yellow" + if "final" in label or label in { + "train", + "train_rollout", + "test_rollout", + "baseline_test_rollout", + } or label.endswith("_rollout"): + return "bold green" + return "bold cyan" + + +async def run_with_progress( + items: list[Any], + *, + coroutine_factory: Callable[[Any, int], Any], + total: int | None = None, + label: str, + enabled: bool, + description: str = "", + concurrency: int, +) -> list[Any]: + """Run ``coroutine_factory(item, index)`` for each item with a ProgressPrinter. + + Creates a semaphore, starts/fails/completes progress rows per item, and + guarantees ``progress.finish()`` in a finally block. Returns the gathered + results in input order. + """ + import asyncio + + n = total if total is not None else len(items) + progress = ProgressPrinter(total=n, label=label, enabled=enabled, description=description) + progress.render() + semaphore = asyncio.Semaphore(concurrency) + + async def _run_one(item: Any, index: int) -> Any: + async with semaphore: + progress.start_one() + try: + result = await coroutine_factory(item, index) + except Exception: + progress.fail_one() + raise + progress.complete_one() + return result + + try: + return list( + await asyncio.gather(*(_run_one(item, idx) for idx, item in enumerate(items))) + ) + finally: + progress.finish() diff --git a/openviking/session/train/components/remote.py b/openviking/session/train/components/remote.py new file mode 100644 index 0000000000..bf2bba1172 --- /dev/null +++ b/openviking/session/train/components/remote.py @@ -0,0 +1,314 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""HTTP-backed CaseLoader and RolloutExecutor implementations.""" + +from __future__ import annotations + +import asyncio +import inspect +from collections.abc import AsyncIterator +from dataclasses import dataclass, field +from typing import Any + +import httpx + +from openviking.session.train.components.dataset_service import ( + case_from_dict, + case_to_dict, + policy_set_to_dict, + rollout_from_dict, +) +from openviking.session.train.components.progress import run_with_progress +from openviking.session.train.context import ExecutionContext +from openviking.session.train.domain import ( + Case, + ExperienceSet, + Rollout, +) + + +@dataclass(slots=True) +class RemoteCaseLoader: + """Load Case batches from a benchmark/environment HTTP service.""" + + service_url: str + dataset: str + domain: str + split: str + batch_size: int | None = None + limit: int | None = None + filters: dict[str, Any] = field(default_factory=dict) + timeout_seconds: float = 60.0 + + async def batches(self, context: Any = None) -> AsyncIterator[list[Case]]: + del context + if self.batch_size is not None and self.batch_size <= 0: + raise ValueError("batch_size must be > 0") + if self.limit is not None and self.limit <= 0: + raise ValueError("limit must be > 0") + + remaining = self.limit + cursor: str | None = None + async with httpx.AsyncClient(base_url=self.service_url.rstrip("/"), timeout=self.timeout_seconds) as client: + while True: + request_limit = self.batch_size or remaining + if request_limit is None: + request_limit = 100 + if remaining is not None: + request_limit = min(request_limit, remaining) + if request_limit <= 0: + return + response = await client.post( + "/v1/cases/query", + json={ + "dataset": self.dataset, + "domain": self.domain, + "split": self.split, + "cursor": cursor, + "limit": request_limit, + "filters": self.filters, + }, + ) + response.raise_for_status() + data = response.json() + cases = [case_from_dict(item) for item in data.get("cases", [])] + if not cases: + return + yield cases + if remaining is not None: + remaining -= len(cases) + if remaining <= 0: + return + cursor = data.get("next_cursor") + if not cursor: + return + + async def split_exists(self) -> bool: + async with httpx.AsyncClient(base_url=self.service_url.rstrip("/"), timeout=self.timeout_seconds) as client: + response = await client.post( + "/v1/cases/query", + json={ + "dataset": self.dataset, + "domain": self.domain, + "split": self.split, + "cursor": None, + "limit": 1, + "filters": self.filters, + }, + ) + response.raise_for_status() + return bool(response.json().get("cases")) + + +@dataclass(slots=True) +class RemoteRolloutExecutor: + """Execute rollouts through a benchmark/environment HTTP service.""" + + service_url: str + options: dict[str, Any] = field(default_factory=dict) + concurrency: int = 20 + request_timeout_seconds: float = 300.0 + poll_interval_seconds: float = 2.0 + execution_timeout_seconds: float = 3600.0 + missing_execution_grace_seconds: float = 60.0 + show_progress: bool = False + progress_label: str = "rollout" + on_rollout_complete: Any | None = None + + def __post_init__(self) -> None: + if self.concurrency <= 0: + raise ValueError("concurrency must be > 0") + if self.request_timeout_seconds <= 0: + raise ValueError("request_timeout_seconds must be > 0") + if self.poll_interval_seconds <= 0: + raise ValueError("poll_interval_seconds must be > 0") + if self.execution_timeout_seconds <= 0: + raise ValueError("execution_timeout_seconds must be > 0") + if self.missing_execution_grace_seconds <= 0: + raise ValueError("missing_execution_grace_seconds must be > 0") + + async def execute( + self, + cases: list[Case], + policy_set: ExperienceSet, + context: ExecutionContext, + ) -> list[Rollout]: + case_list = list(cases) + stage_label = _progress_stage_label( + context.metadata.get("stage"), + default=self.progress_label, + ) + timeout = httpx.Timeout(self.request_timeout_seconds) + async with httpx.AsyncClient(base_url=self.service_url.rstrip("/"), timeout=timeout) as client: + + async def _execute(case: Case, index: int) -> Rollout: + rollout = await self._execute_one(client, case, policy_set, context) + await self._emit_rollout_complete( + rollout=rollout, + index=index, + context=context, + ) + return rollout + + return await run_with_progress( + case_list, + coroutine_factory=_execute, + label=stage_label, + enabled=self.show_progress, + description=f"Running {len(case_list)} rollouts, concurrency={self.concurrency}", + concurrency=self.concurrency, + ) + + async def _emit_rollout_complete( + self, + *, + rollout: Rollout, + index: int, + context: ExecutionContext, + ) -> None: + if self.on_rollout_complete is None: + return + result = self.on_rollout_complete( + rollout=rollout, + index=index, + context=context, + ) + if inspect.isawaitable(result): + await result + + async def _execute_one( + self, + client: httpx.AsyncClient, + case: Case, + policy_set: ExperienceSet, + context: ExecutionContext, + ) -> Rollout: + response = await client.post( + "/v1/rollouts/execute", + json={ + "case": case_to_dict(case), + "policy_set": policy_set_to_dict(policy_set), + "execution_context": { + "policy_snapshot_id": context.policy_snapshot_id, + "metadata": context.metadata, + }, + "options": _remote_execution_options(self.options), + }, + ) + response.raise_for_status() + execution_id = _require_execution_id(response.json(), case=case) + return await self._poll_execution(client, execution_id, case=case) + + async def _poll_execution( + self, + client: httpx.AsyncClient, + execution_id: str, + *, + case: Case, + ) -> Rollout: + loop = asyncio.get_running_loop() + started_at = loop.time() + deadline = started_at + self.execution_timeout_seconds + missing_execution_deadline = started_at + self.missing_execution_grace_seconds + transient_errors = 0 + missing_execution_errors = 0 + last_transient_error: BaseException | None = None + while True: + try: + response = await client.get(f"/v1/rollouts/executions/{execution_id}") + transient_errors = 0 + except (httpx.ReadError, httpx.ConnectError, httpx.RemoteProtocolError, httpx.TimeoutException) as exc: + transient_errors += 1 + last_transient_error = exc + if asyncio.get_running_loop().time() >= deadline: + raise TimeoutError( + f"rollout execution {execution_id} polling timed out for case {case.name} " + f"after {self.execution_timeout_seconds}s; last polling error: " + f"{type(exc).__name__}: {exc}" + ) from exc + await asyncio.sleep(min(self.poll_interval_seconds * transient_errors, 10.0)) + continue + if response.status_code == 404: + missing_execution_errors += 1 + if ( + loop.time() >= deadline + or loop.time() >= missing_execution_deadline + ): + raise RuntimeError( + f"rollout execution {execution_id} was not found while polling " + f"case {case.name}; observed {missing_execution_errors} 404 response(s) " + f"over {loop.time() - started_at:.1f}s. Last response: " + f"{_response_text(response)}" + ) + await asyncio.sleep( + min(self.poll_interval_seconds * missing_execution_errors, 10.0) + ) + continue + response.raise_for_status() + data = response.json() + status = data.get("status") + if status == "completed": + rollout_data = data.get("rollout") + if not isinstance(rollout_data, dict): + raise RuntimeError(f"rollout execution {execution_id} completed without rollout") + return rollout_from_dict(rollout_data) + if status == "failed": + raise RuntimeError( + f"rollout execution {execution_id} failed for case {case.name}: " + f"{data.get('error') or 'unknown error'}" + ) + if asyncio.get_running_loop().time() >= deadline: + last_error_text = ( + f"; last polling error: {type(last_transient_error).__name__}: " + f"{last_transient_error}" + if last_transient_error is not None + else "" + ) + raise TimeoutError( + f"rollout execution {execution_id} timed out for case {case.name} " + f"after {self.execution_timeout_seconds}s{last_error_text}" + ) + await asyncio.sleep(self.poll_interval_seconds) + + +def _progress_stage_label(stage: Any, *, default: str) -> str: + stage_text = str(stage or "") + stage_parts = stage_text.split(maxsplit=1) + stage_name = stage_parts[0] if stage_parts else "" + if _is_progress_stage_name(stage_name): + return f"{stage_name}_start" + if stage_name.endswith("_start") and _is_progress_stage_name(stage_name[:-6]): + return stage_name + return default + + +def _is_progress_stage_name(stage_name: str) -> bool: + return ( + stage_name == "train_rollout" + or stage_name == "test_rollout" + or stage_name.endswith("_rollout") + ) + + +def _remote_execution_options(options: dict[str, Any]) -> dict[str, Any]: + execution_options = dict(options) + execution_options.pop("concurrency", None) + return execution_options + + +def _require_execution_id(data: dict[str, Any], *, case: Case) -> str: + execution_id = data.get("execution_id") + if not isinstance(execution_id, str) or not execution_id: + raise RuntimeError(f"rollout service did not return execution_id for case {case.name}") + return execution_id + + +def _response_text(response: httpx.Response, *, max_chars: int = 500) -> str: + try: + text = response.text + except Exception: + return "" + text = text.replace("\n", "\\n") + if len(text) > max_chars: + return text[:max_chars] + "..." + return text diff --git a/openviking/session/train/components/report_builder.py b/openviking/session/train/components/report_builder.py new file mode 100644 index 0000000000..f2707015f8 --- /dev/null +++ b/openviking/session/train/components/report_builder.py @@ -0,0 +1,445 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Default report builders for session train/eval results.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any + +from openviking.session.train.components.reporter import NoopPipelineLifecycleHook +from openviking.session.train.domain import ( + CriterionResult, + ExperienceSet, + PipelineEpochResult, + PipelineEvaluationResult, + Rollout, + RolloutAnalysis, +) + + +@dataclass(slots=True) +class PipelineReportBuilder: + """Build serializable summary reports from pipeline domain objects.""" + + trial_index_key: str = "trial" + memory_tool_name_prefix: str = "openviking" + + def evaluation_report( + self, + result: PipelineEvaluationResult, + ) -> dict[str, Any]: + rewards = [float(analysis.evaluation.score) for analysis in result.analyses] + passed_count = sum(1 for analysis in result.analyses if analysis.evaluation.passed) + case_count = len(result.analyses) + return { + "epoch": result.epoch, + **_eval_metadata_fields(result.metadata), + "case_count": case_count, + "accuracy": _ratio(passed_count, case_count), + "passed_count": passed_count, + "average_reward": _average(rewards), + "rewards": rewards, + "snapshot_ids": list(result.policy_snapshot_ids), + "metadata": dict(result.metadata), + "memory_usage": self.memory_usage_from_analyses(result.analyses), + "cost_seconds": result.metadata.get("cost_seconds"), + } + + def trial_evaluation_report( + self, + result: PipelineEvaluationResult, + *, + trial_count: int | None = None, + ) -> dict[str, Any]: + if trial_count is None: + trial_count = self._trial_count_from_analyses(result.analyses) + analyses_by_trial: dict[int, list[RolloutAnalysis]] = { + trial_index: [] for trial_index in range(trial_count) + } + for analysis in result.analyses: + trial_index = self.analysis_trial_index(analysis) + analyses_by_trial.setdefault(trial_index, []).append(analysis) + + trials = [ + { + "trial": trial_index, + **self.evaluation_summary_from_analyses( + analyses_by_trial.get(trial_index, []) + ), + } + for trial_index in range(trial_count) + ] + accuracies = [ + float(item["accuracy"]) + for item in trials + if item.get("accuracy") is not None + ] + average_rewards = [ + float(item["average_reward"]) + for item in trials + if item.get("average_reward") is not None + ] + overall = self.evaluation_report(result) + case_counts = [int(item["case_count"]) for item in trials] + return { + **overall, + "trial_count": trial_count, + "case_count_per_trial": case_counts[0] if len(set(case_counts)) == 1 else None, + "case_counts_per_trial": case_counts, + "total_rollout_count": len(result.analyses), + "accuracy_mean": _average(accuracies), + "accuracy_std": _stddev(accuracies), + "average_reward_mean": _average(average_rewards), + "average_reward_std": _stddev(average_rewards), + "trials": trials, + # Keep callers simple: accuracy/average_reward denote the + # trial-level mean when trials are enabled. + "accuracy": _average(accuracies), + "average_reward": _average(average_rewards), + } + + def evaluation_summary_from_analyses( + self, + analyses: list[RolloutAnalysis], + ) -> dict[str, Any]: + rewards = [float(analysis.evaluation.score) for analysis in analyses] + passed_count = sum(1 for analysis in analyses if analysis.evaluation.passed) + case_count = len(analyses) + return { + "case_count": case_count, + "accuracy": _ratio(passed_count, case_count), + "passed_count": passed_count, + "average_reward": _average(rewards), + "rewards": rewards, + "memory_usage": self.memory_usage_from_analyses(analyses), + } + + def analysis_trial_index(self, analysis: RolloutAnalysis) -> int: + rollout = analysis.metadata.get("rollout") + if isinstance(rollout, Rollout): + value = rollout.case.input.get( + self.trial_index_key, + rollout.case.metadata.get(self.trial_index_key, 0), + ) + else: + value = 0 + try: + return int(value) + except (TypeError, ValueError): + return 0 + + def _trial_count_from_analyses(self, analyses: list[RolloutAnalysis]) -> int: + for analysis in analyses: + rollout = analysis.metadata.get("rollout") + if not isinstance(rollout, Rollout): + continue + value = rollout.case.input.get( + f"{self.trial_index_key}_count", + rollout.case.metadata.get(f"{self.trial_index_key}_count"), + ) + if value is not None: + return int(value) + return 1 + + def train_rollout_report( + self, + *, + epoch: int, + rollouts: list[Rollout], + snapshot_id: str, + ) -> dict[str, Any]: + analyses = _analyses_from_rollout_evaluations(rollouts) + train_eval = self.train_evaluation_report(analyses) + cache_hits = [ + rollout + for rollout in rollouts + if isinstance(rollout.metadata, dict) + and bool(rollout.metadata.get("train_rollout_cache_hit")) + ] + return { + "epoch": epoch, + "snapshot_id": snapshot_id, + "snapshot_ids": [snapshot_id], + **train_eval, + "cache_hit_count": len(cache_hits), + "cache_miss_count": max(len(rollouts) - len(cache_hits), 0), + "from_cache": bool(cache_hits) and len(cache_hits) == len(rollouts), + } + + def train_epoch_report( + self, + epoch_result: PipelineEpochResult, + *, + rollout_report: dict[str, Any] | None = None, + ) -> dict[str, Any]: + train_eval = self.train_evaluation_report(epoch_result.analyses) + commit_results = list(epoch_result.apply_result.metadata.get("commit_results", [])) + errors = [error for item in commit_results if (error := item.get("error"))] + snapshot_ids = list(epoch_result.policy_snapshot_ids) + display_eval = rollout_report or train_eval + return { + "epoch": epoch_result.epoch, + "case_count": display_eval["case_count"], + "accuracy": display_eval["accuracy"], + "passed_count": display_eval["passed_count"], + "average_reward": display_eval["average_reward"], + "train_rollout": rollout_report, + "train_eval": train_eval, + "batch_count": len(snapshot_ids), + "gradient_count": len(epoch_result.gradients), + "committed_rollout_count": len(commit_results), + "errors": errors, + "failed_commit_trace_ids": _failed_commit_trace_ids(commit_results), + "failed_commit_telemetry_ids": _failed_commit_telemetry_ids(commit_results), + "snapshot_ids": snapshot_ids, + "commit_results": commit_results, + "metadata": dict(epoch_result.metadata), + "memory_usage": self.memory_usage_from_analyses(epoch_result.analyses), + "cost_seconds": epoch_result.metadata.get("cost_seconds"), + } + + def train_evaluation_report( + self, + analyses: list[RolloutAnalysis], + ) -> dict[str, Any]: + rewards = [float(analysis.evaluation.score) for analysis in analyses] + return { + **self.evaluation_summary_from_analyses(analyses), + "reward_std": _stddev(rewards), + "case_results": [_train_case_evaluation_result(analysis) for analysis in analyses], + } + + def memory_usage_from_analyses( + self, + analyses: list[RolloutAnalysis], + ) -> dict[str, Any]: + rollouts = [ + rollout + for analysis in analyses + if isinstance((rollout := analysis.metadata.get("rollout")), Rollout) + ] + rollout_count = len(rollouts) + memory_context_count = 0 + memory_tool_call_rollout_count = 0 + memory_tool_call_total = 0 + for rollout in rollouts: + metadata = rollout.metadata or {} + if str(metadata.get("memory") or "").strip(): + memory_context_count += 1 + tool_call_count = self.memory_tool_call_count(metadata.get("tools_used")) + if tool_call_count: + memory_tool_call_rollout_count += 1 + memory_tool_call_total += tool_call_count + return { + "rollout_count": rollout_count, + "memory_context_count": memory_context_count, + "memory_context_ratio": _ratio(memory_context_count, rollout_count), + "memory_tool_call_rollout_count": memory_tool_call_rollout_count, + "memory_tool_call_rollout_ratio": _ratio( + memory_tool_call_rollout_count, + rollout_count, + ), + "memory_tool_call_total": memory_tool_call_total, + } + + def memory_tool_call_count(self, tools_used: Any) -> int: + if not isinstance(tools_used, list): + return 0 + count = 0 + for tool_info in tools_used: + if not isinstance(tool_info, dict): + continue + tool_name = str(tool_info.get("tool_name") or "") + if tool_name.startswith(self.memory_tool_name_prefix): + count += 1 + return count + + def accuracy_delta( + self, + baseline_eval: dict[str, Any] | None, + final_eval: dict[str, Any] | None, + ) -> float | None: + if not baseline_eval or not final_eval: + return None + baseline = baseline_eval.get("accuracy") + final = final_eval.get("accuracy") + if baseline is None or final is None: + return None + return float(final) - float(baseline) + + +@dataclass(slots=True) +class PipelineReportHook(NoopPipelineLifecycleHook): + """Lifecycle hook that builds default serializable pipeline reports.""" + + def on_train_rollout_end( + self, + *, + epoch: int, + rollouts: list[Any], + snapshot_id: str, + policy_set: ExperienceSet, + context: Any, + ) -> dict[str, Any] | None: + del policy_set + return _context_report_builder(context).train_rollout_report( + epoch=epoch, + rollouts=list(rollouts), + snapshot_id=snapshot_id, + ) + + def on_epoch_end( + self, + *, + epoch_result: PipelineEpochResult, + policy_set: ExperienceSet, + context: Any, + ) -> Any: + del policy_set + from openviking.session.train.context import PipelineHookDecision + + return PipelineHookDecision( + report=_context_report_builder(context).train_epoch_report( + epoch_result, + rollout_report=epoch_result.metadata.get("train_rollout_report"), + ) + ) + + def on_eval_end( + self, + *, + evaluation_result: PipelineEvaluationResult, + policy_set: ExperienceSet, + context: Any, + ) -> dict[str, Any] | None: + del policy_set + eval_trials = int(getattr(context, "eval_trials", 1) or 1) + builder = _context_report_builder(context) + if eval_trials > 1: + return builder.trial_evaluation_report( + evaluation_result, + trial_count=eval_trials, + ) + return builder.evaluation_report(evaluation_result) + + +def _context_report_builder(context: Any) -> PipelineReportBuilder: + report_builder = getattr(context, "report_builder", None) + if report_builder is not None: + return report_builder + return PipelineReportBuilder( + trial_index_key=str(getattr(context, "trial_index_key", "trial") or "trial") + ) + + +def _eval_metadata_fields(metadata: dict[str, Any]) -> dict[str, Any]: + fields: dict[str, Any] = {} + if metadata.get("rollout_stage") is not None: + fields["rollout_stage"] = metadata["rollout_stage"] + if metadata.get("eval_split") is not None: + fields["split"] = metadata["eval_split"] + return fields + + +def _analyses_from_rollout_evaluations(rollouts: list[Rollout]) -> list[RolloutAnalysis]: + analyses: list[RolloutAnalysis] = [] + for idx, rollout in enumerate(rollouts): + if rollout.evaluation is None: + raise ValueError( + "report builder requires rollout.evaluation; " + f"missing index={idx}, case={rollout.case.name}" + ) + analyses.append( + RolloutAnalysis( + evaluation=rollout.evaluation, + trajectories=[], + metadata={ + "rollout": rollout, + "rollout_messages": rollout.messages, + "policy_snapshot_id": rollout.policy_snapshot_id, + "evaluation_source": "rollout", + }, + ) + ) + return analyses + + +def _train_case_evaluation_result(analysis: RolloutAnalysis) -> dict[str, Any]: + evaluation = analysis.evaluation + rollout = analysis.metadata.get("rollout") + case = rollout.case if isinstance(rollout, Rollout) else None + rollout_metadata = dict(rollout.metadata) if isinstance(rollout, Rollout) else {} + case_input = dict(case.input) if case is not None else {} + return { + "case_name": case.name if case is not None else "", + "task_signature": case.task_signature if case is not None else "", + "data_split": case_input.get("data_split") or rollout_metadata.get("data_split"), + "task_no": case_input.get("task_no") or rollout_metadata.get("task_no"), + "task_id": case_input.get("task_id") or rollout_metadata.get("task_id"), + "passed": evaluation.passed, + "score": float(evaluation.score), + "reward": rollout_metadata.get("reward", evaluation.metadata.get("reward")), + "evaluation_result": rollout_metadata.get( + "evaluation_result", + evaluation.metadata.get("evaluation_result"), + ), + "feedback": list(evaluation.feedback), + "criterion_results": [ + _criterion_result_report(result) for result in evaluation.criterion_results + ], + "policy_snapshot_id": analysis.metadata.get("policy_snapshot_id"), + } + + +def _criterion_result_report(result: CriterionResult) -> dict[str, Any]: + return { + "criterion_name": result.criterion_name, + "passed": result.passed, + "score": float(result.score), + "feedback": list(result.feedback), + "evidence": list(result.evidence), + "metadata": dict(result.metadata), + } + + +def _failed_commit_trace_ids(commit_results: list[dict[str, Any]]) -> list[str]: + trace_ids: list[str] = [] + for item in commit_results: + if not item.get("error"): + continue + trace_id = str(item.get("trace_id") or "").strip() + if trace_id: + trace_ids.append(trace_id) + return trace_ids + + +def _failed_commit_telemetry_ids(commit_results: list[dict[str, Any]]) -> list[str]: + telemetry_ids: list[str] = [] + for item in commit_results: + if not item.get("error"): + continue + telemetry_id = str(item.get("telemetry_id") or "").strip() + if telemetry_id: + telemetry_ids.append(telemetry_id) + return telemetry_ids + + +def _average(values: list[float]) -> float | None: + if not values: + return None + return sum(values) / len(values) + + +def _stddev(values: list[float]) -> float | None: + if not values: + return None + mean = sum(values) / len(values) + return math.sqrt(sum((value - mean) ** 2 for value in values) / len(values)) + + +def _ratio(numerator: int, denominator: int) -> float | None: + if denominator <= 0: + return None + return numerator / denominator diff --git a/openviking/session/train/components/reporter.py b/openviking/session/train/components/reporter.py new file mode 100644 index 0000000000..1b5fe7ab9f --- /dev/null +++ b/openviking/session/train/components/reporter.py @@ -0,0 +1,656 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Console reporting helpers for session training pipelines.""" + +from __future__ import annotations + +import inspect +import sys +from collections.abc import Awaitable +from dataclasses import dataclass, field +from typing import Any, Protocol + +try: # pragma: no cover - cosmetic terminal rendering + from rich.console import Console + from rich.text import Text +except Exception: # pragma: no cover - rich is optional + Console = None + Text = None + +from openviking.session.train.components.progress import format_duration, format_label, label_style + +HookResult = Awaitable[None] | None +ReportHookResult = Awaitable[dict[str, Any] | None] | dict[str, Any] | None +DecisionHookResult = Awaitable[Any] | Any | None + + +class PipelineLifecycleHook(Protocol): + """Lifecycle hook extension point for train/eval pipelines.""" + + def on_epoch_start(self, *, epoch: int, context: Any) -> HookResult: ... + + def on_train_rollout_end( + self, + *, + epoch: int, + rollouts: list[Any], + snapshot_id: str, + policy_set: Any, + context: Any, + ) -> ReportHookResult: ... + + def on_epoch_end( + self, + *, + epoch_result: Any, + policy_set: Any, + context: Any, + ) -> DecisionHookResult: ... + + def on_eval_end( + self, + *, + evaluation_result: Any, + policy_set: Any, + context: Any, + ) -> ReportHookResult: ... + + def on_eval_report( + self, + *, + label: str, + report: dict[str, Any], + context: Any, + ) -> HookResult: ... + + def on_train_rollout_report( + self, + *, + report: dict[str, Any], + context: Any, + ) -> HookResult: ... + + def on_train_report( + self, + *, + report: dict[str, Any], + context: Any, + ) -> HookResult: ... + + def on_run_summary( + self, + *, + title: str, + fields: dict[str, Any], + baseline_eval: dict[str, Any] | None = None, + final_eval: dict[str, Any] | None = None, + accuracy_delta: float | None = None, + output_path: str | None = None, + rollouts_root: str | None = None, + rollouts_index_path: str | None = None, + latest_failed_rollout: str | None = None, + ) -> HookResult: ... + + +class NoopPipelineLifecycleHook: + """Base class for lifecycle hooks that only need to override some events.""" + + def on_epoch_start(self, *, epoch: int, context: Any) -> None: + del epoch, context + + def on_train_rollout_end( + self, + *, + epoch: int, + rollouts: list[Any], + snapshot_id: str, + policy_set: Any, + context: Any, + ) -> None: + del epoch, rollouts, snapshot_id, policy_set, context + + def on_epoch_end( + self, + *, + epoch_result: Any, + policy_set: Any, + context: Any, + ) -> None: + del epoch_result, policy_set, context + + def on_eval_end( + self, + *, + evaluation_result: Any, + policy_set: Any, + context: Any, + ) -> None: + del evaluation_result, policy_set, context + + def on_eval_report( + self, + *, + label: str, + report: dict[str, Any], + context: Any, + ) -> None: + del label, report, context + + def on_train_rollout_report( + self, + *, + report: dict[str, Any], + context: Any, + ) -> None: + del report, context + + def on_train_report( + self, + *, + report: dict[str, Any], + context: Any, + ) -> None: + del report, context + + def on_run_summary( + self, + *, + title: str, + fields: dict[str, Any], + baseline_eval: dict[str, Any] | None = None, + final_eval: dict[str, Any] | None = None, + accuracy_delta: float | None = None, + output_path: str | None = None, + rollouts_root: str | None = None, + rollouts_index_path: str | None = None, + latest_failed_rollout: str | None = None, + ) -> None: + del ( + title, + fields, + baseline_eval, + final_eval, + accuracy_delta, + output_path, + rollouts_root, + rollouts_index_path, + latest_failed_rollout, + ) + + +async def emit_run_summary( + context: Any, + *, + title: str, + fields: dict[str, Any], + baseline_eval: dict[str, Any] | None = None, + final_eval: dict[str, Any] | None = None, + accuracy_delta: float | None = None, + output_path: str | None = None, + rollouts_root: str | None = None, + rollouts_index_path: str | None = None, + latest_failed_rollout: str | None = None, +) -> None: + """Emit a run-level summary event to lifecycle hooks on a pipeline context.""" + + lifecycle_hooks = list(getattr(context, "lifecycle_hooks", []) or []) + for hook in lifecycle_hooks: + result = hook.on_run_summary( + title=title, + fields=fields, + baseline_eval=baseline_eval, + final_eval=final_eval, + accuracy_delta=accuracy_delta, + output_path=output_path, + rollouts_root=rollouts_root, + rollouts_index_path=rollouts_index_path, + latest_failed_rollout=latest_failed_rollout, + ) + if inspect.isawaitable(result): + await result + + +@dataclass(slots=True) +class ConsolePipelineReporter(NoopPipelineLifecycleHook): + """Default stdout lifecycle hook for batch train/eval runners.""" + + use_rich: bool | None = None + _epoch_summaries: dict[int, dict[str, Any]] = field(init=False, default_factory=dict) + _printed_epoch_summaries: set[int] = field(init=False, default_factory=set) + + def __post_init__(self) -> None: + if self.use_rich is None: + self.use_rich = Console is not None and Text is not None and sys.stdout.isatty() + + def on_eval_report( + self, + *, + label: str, + report: dict[str, Any], + context: Any, + ) -> None: + del context + label = str(report.get("rollout_stage") or label) + split = report.get("split") + trial_count = int(report.get("trial_count") or 1) + if trial_count > 1: + self._print_line( + label, + [ + ("epoch", report["epoch"]), + *_split_field(split), + ("trials", trial_count, "cyan"), + ("cases_per_trial", report.get("case_count_per_trial") or "varies"), + ( + "total_rollouts", + report.get("total_rollout_count", report["case_count"]), + "cyan", + ), + ( + "accuracy", + fmt_percent(report.get("accuracy_mean")), + _accuracy_style(report.get("accuracy_mean")), + ), + ("", f"± {fmt_percentage_point_abs(report.get('accuracy_std'))}", "yellow"), + *_cost_field(report), + ], + ) + self._remember_eval_report(label, report) + if _is_epoch_test_report(label, report): + self._print_epoch_summary(int(report["epoch"])) + self._print_stage_separator() + return + self._print_line( + label, + [ + ("epoch", report["epoch"]), + *_split_field(split), + ("cases", report["case_count"]), + ( + "accuracy", + fmt_percent(report["accuracy"]), + _accuracy_style(report.get("accuracy")), + ), + ( + "passed", + f"{report['passed_count']}/{report['case_count']}", + _passed_style(report), + ), + *_cost_field(report), + ], + ) + self._remember_eval_report(label, report) + if _is_epoch_test_report(label, report): + self._print_epoch_summary(int(report["epoch"])) + self._print_stage_separator() + + def on_epoch_start(self, *, epoch: int, context: Any) -> None: + del context + text = f" epoch {epoch} " + width = 44 + left = max((width - len(text)) // 2, 1) + right = max(width - len(text) - left, 1) + line = f"{'=' * left}{text}{'=' * right}" + if not self.use_rich: + print(line) + return + Console().print(line, style="bold cyan") + + def on_train_rollout_report( + self, + *, + report: dict[str, Any], + context: Any, + ) -> None: + del context + self._remember_train_rollout_report(report) + self._print_line( + "train_rollout", + [ + ("epoch", report["epoch"]), + ("cases", report["case_count"]), + *_cache_field(report), + ( + "accuracy", + fmt_percent(report["accuracy"]), + _accuracy_style(report.get("accuracy")), + ), + ( + "passed", + f"{report['passed_count']}/{report['case_count']}", + _passed_style(report), + ), + *_cost_field(report), + ], + ) + self._print_stage_separator() + + def on_train_report( + self, + *, + report: dict[str, Any], + context: Any, + ) -> None: + self._remember_train_report(report) + error_count = len(report["errors"]) + self._print_line( + "train", + [ + ("epoch", report["epoch"]), + ("commits", report["committed_rollout_count"], "cyan"), + ("errors", error_count, "green" if error_count == 0 else "red bold"), + *_cost_field(report), + ], + ) + if report.get("errors"): + trace_ids = report.get("failed_commit_trace_ids") or [] + telemetry_ids = report.get("failed_commit_telemetry_ids") or [] + if trace_ids: + print(f"[train] failed_commit_trace_ids={','.join(trace_ids)}") + else: + print("[train] failed_commit_trace_ids=") + if telemetry_ids: + print(f"[train] failed_commit_telemetry_ids={','.join(telemetry_ids)}") + if not _has_epoch_eval(context): + self._print_epoch_summary(int(report["epoch"])) + self._print_stage_separator() + + def on_run_summary( + self, + *, + title: str, + fields: dict[str, Any], + baseline_eval: dict[str, Any] | None = None, + final_eval: dict[str, Any] | None = None, + accuracy_delta: float | None = None, + output_path: str | None = None, + rollouts_root: str | None = None, + rollouts_index_path: str | None = None, + latest_failed_rollout: str | None = None, + ) -> None: + print(f"==== {title} ====") + for key, value in fields.items(): + if value is not None: + print(f"{key}: {value}") + if baseline_eval: + self._report_eval_line("baseline", baseline_eval) + if final_eval: + self._report_eval_line("final", final_eval) + if accuracy_delta is not None: + print(f"accuracy delta: {fmt_percentage_point(accuracy_delta)}") + if output_path: + print(f"report: {output_path}") + if rollouts_root: + print(f"rollouts: {rollouts_root}") + if rollouts_index_path: + print(f"rollouts_index: {rollouts_index_path}") + if latest_failed_rollout: + print(f"latest_failed_rollout: {latest_failed_rollout}") + + def _print_stage_separator(self) -> None: + print("-" * 60) + + def _print_line(self, label: str, fields: list[tuple[Any, ...]]) -> None: + if not self.use_rich: + print( + f"[{label}] " + + " ".join( + _plain_field(item) + for item in fields + ) + ) + return + console = Console() + line = Text() + line.append(format_label(label), style=label_style(label)) + for item in fields: + key = str(item[0]) + value = str(item[1]) + value_style = str(item[2]) if len(item) > 2 else "default" + line.append(" ") + if key: + line.append(f"{key}=", style="dim") + line.append(value, style=value_style) + console.print(line) + + def _remember_train_rollout_report(self, report: dict[str, Any]) -> None: + epoch = _report_epoch(report) + if epoch is None: + return + self._epoch_summaries.setdefault(epoch, {})["train_rollout"] = dict(report) + + def _remember_train_report(self, report: dict[str, Any]) -> None: + epoch = _report_epoch(report) + if epoch is None: + return + self._epoch_summaries.setdefault(epoch, {})["train"] = dict(report) + + def _remember_eval_report(self, label: str, report: dict[str, Any]) -> None: + epoch = _report_epoch(report) + if epoch is None: + return + self._epoch_summaries.setdefault(epoch, {})["test"] = { + **dict(report), + "label": label, + } + + def _print_epoch_summary(self, epoch: int) -> None: + if epoch in self._printed_epoch_summaries: + return + summary = self._epoch_summaries.get(epoch) or {} + train_data = _summary_train_report(summary) + test_data = summary.get("test") + if train_data is None and test_data is None: + return + self._printed_epoch_summaries.add(epoch) + + header = f" epoch {epoch} summary " + width = max(44, len(header) + 8) + left = max((width - len(header)) // 2, 1) + right = max(width - len(header) - left, 1) + border_top = f"{'=' * left}{header}{'=' * right}" + border_bottom = "=" * len(border_top) + self._print_summary_fragments([(border_top, "cyan bold")]) + if train_data is not None: + self._print_summary_fragments(_train_summary_fragments(train_data)) + if test_data is not None: + self._print_summary_fragments(_test_summary_fragments(test_data)) + self._print_summary_fragments([(border_bottom, "cyan bold")]) + + def _print_summary_fragments(self, fragments: list[tuple[str, str]]) -> None: + if not self.use_rich: + print("".join(_style_plain(text, style) for text, style in fragments)) + return + line = Text() + for text, style in fragments: + line.append(text, style=style or "default") + Console().print(line) + + def _report_eval_line(self, label: str, data: dict[str, Any]) -> None: + trial_count = int(data.get("trial_count") or 1) + if trial_count > 1: + print( + f"{label} accuracy: {fmt_percent(data.get('accuracy_mean'))} ± " + f"{fmt_percentage_point_abs(data.get('accuracy_std'))} " + f"(trials={trial_count}, " + f"cases_per_trial={data.get('case_count_per_trial') or 'varies'})" + ) + return + print( + f"{label} accuracy: " + f"{fmt_percent(data['accuracy'])} " + f"({data['passed_count']}/{data['case_count']})" + ) + + +def _cost_field(report: dict[str, Any]) -> list[tuple[str, str, str]]: + cost_seconds = report.get("cost_seconds") + if cost_seconds is None: + return [] + return [("cost", format_duration(float(cost_seconds)), "magenta bold")] + + +def _cache_field(report: dict[str, Any]) -> list[tuple[str, str, str]]: + hit_count = int(report.get("cache_hit_count") or 0) + if hit_count <= 0: + return [] + total = int(report.get("case_count") or 0) + if bool(report.get("from_cache")): + value = "all" if total <= 0 else f"{hit_count}/{total}" + else: + value = f"partial({hit_count}/{total})" if total > 0 else str(hit_count) + return [("from_cache", value, "cyan bold")] + + +def _split_field(split: Any) -> list[tuple[str, str, str]]: + if split is None: + return [] + return [("split", str(split), "cyan")] + + +def _plain_field(item: tuple[Any, ...]) -> str: + text = f"{item[0]}={item[1]}" if item[0] else str(item[1]) + if len(item) <= 2 or item[0] != "accuracy": + return text + return f"accuracy={_style_plain(str(item[1]), str(item[2]))}" + + +def _style_plain(text: str, style: str) -> str: + if not style: + return text + parts: list[str] = [] + style_tokens = set(style.split()) + if "bold" in style_tokens: + parts.append("1") + if "red" in style_tokens: + parts.append("31") + elif "green" in style_tokens: + parts.append("32") + elif "yellow" in style_tokens: + parts.append("33") + elif "cyan" in style_tokens: + parts.append("36") + elif "magenta" in style_tokens: + parts.append("35") + if not parts: + return text + return f"\033[{';'.join(parts)}m{text}\033[0m" + + +def _report_epoch(report: dict[str, Any]) -> int | None: + try: + return int(report["epoch"]) + except (KeyError, TypeError, ValueError): + return None + + +def _has_epoch_eval(context: Any) -> bool: + return getattr(context, "eval_each_epoch_case_loader", None) is not None + + +def _is_epoch_test_report(label: str, report: dict[str, Any]) -> bool: + label_text = str(label) + return ( + ( + label_text == "test_rollout" + or label_text.startswith("epoch_") + or label_text.startswith("eval_") + ) + and label_text.endswith("_rollout") + and report.get("epoch") is not None + and int(report.get("epoch") or 0) >= 0 + ) + + +def _summary_train_report(summary: dict[str, Any]) -> dict[str, Any] | None: + train_rollout = summary.get("train_rollout") + if isinstance(train_rollout, dict): + return train_rollout + train = summary.get("train") + if isinstance(train, dict): + nested = train.get("train_rollout") + if isinstance(nested, dict): + return nested + return None + + +def _train_summary_fragments(data: dict[str, Any]) -> list[tuple[str, str]]: + accuracy = data.get("accuracy") + passed = data.get("passed_count") + total = data.get("case_count") + fragments = [ + ("TRAIN accuracy: ", "bold"), + (fmt_percent(accuracy), _accuracy_style(accuracy)), + ] + if passed is not None and total is not None: + fragments.extend([(" passed=", "default"), (f"{passed}/{total}", _passed_style(data))]) + return fragments + + +def _test_summary_fragments(data: dict[str, Any]) -> list[tuple[str, str]]: + trial_count = int(data.get("trial_count") or 1) + if trial_count > 1: + accuracy = data.get("accuracy_mean") + fragments = [ + ("EVAL accuracy: ", "bold"), + (fmt_percent(accuracy), _accuracy_style(accuracy)), + (" ± ", "default"), + (fmt_percentage_point_abs(data.get("accuracy_std")), "yellow"), + (" trials=", "default"), + (str(trial_count), "cyan"), + ] + return fragments + + accuracy = data.get("accuracy") + fragments = [ + ("EVAL accuracy: ", "bold"), + (fmt_percent(accuracy), _accuracy_style(accuracy)), + (" passed=", "default"), + (f"{data.get('passed_count')}/{data.get('case_count')}", _passed_style(data)), + ] + return fragments + + +def fmt_score(value: Any) -> str: + if value is None: + return "n/a" + return f"{float(value):.6f}" + + +def fmt_percent(value: Any) -> str: + if value is None: + return "n/a" + return f"{float(value) * 100:.2f}%" + + +def fmt_percentage_point(value: Any) -> str: + if value is None: + return "n/a" + return f"{float(value) * 100:+.2f}pp" + + +def fmt_percentage_point_abs(value: Any) -> str: + if value is None: + return "n/a" + return f"{float(value) * 100:.2f}pp" + + +def _accuracy_style(value: Any) -> str: + if value is None: + return "dim" + score = float(value) + if score >= 0.8: + return "green bold" + if score >= 0.5: + return "yellow bold" + return "red bold" + + +def _passed_style(data: dict[str, Any]) -> str: + case_count = int(data.get("case_count") or 0) + passed_count = int(data.get("passed_count") or 0) + if case_count > 0 and passed_count == case_count: + return "green bold" + if passed_count == 0: + return "red bold" + return "yellow bold" diff --git a/openviking/session/train/components/rollout_artifact_recorder.py b/openviking/session/train/components/rollout_artifact_recorder.py new file mode 100644 index 0000000000..39f0d209e5 --- /dev/null +++ b/openviking/session/train/components/rollout_artifact_recorder.py @@ -0,0 +1,963 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Rollout artifact recording for batch train/eval pipelines.""" + +from __future__ import annotations + +import difflib +import json +import re +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from openviking.message import ToolPart +from openviking.session.train.components.dataset_service import ( + case_to_dict, + evaluation_to_dict, + jsonable, +) +from openviking.session.train.components.reporter import NoopPipelineLifecycleHook +from openviking.session.train.domain import ( + PipelineEpochResult, + PipelineEvaluationResult, + Rollout, + RolloutAnalysis, +) + + +@dataclass(slots=True) +class RolloutArtifactIndex: + """Serializable index of recorded rollout artifacts.""" + + run_dir: str + rollouts_root: str + case_groups: list[dict[str, Any]] = field(default_factory=list) + latest_failed_rollout: str | None = None + + def to_dict(self) -> dict[str, Any]: + return { + "run_dir": self.run_dir, + "rollouts_root": self.rollouts_root, + "latest_failed_rollout": self.latest_failed_rollout, + "case_groups": self.case_groups, + } + + +class RolloutArtifactRecorder(NoopPipelineLifecycleHook): + """Write per-case/per-rollout artifacts for all case groups. + + Each case group and all its rollouts are written to disk so success/failure + trials can be compared by an LLM or inspected manually. + + Inherits from NoopPipelineLifecycleHook so it can be registered as a + pipeline lifecycle hook; only on_epoch_end and on_eval_end are overridden. + """ + + def __init__( + self, + *, + run_dir: Path, + client: Any | None = None, + latest_pointer_path: Path | None = None, + ) -> None: + self.run_dir = run_dir.expanduser().resolve() + self.rollouts_root = self.run_dir / "rollouts" + self.client = client + self.latest_pointer_path = ( + latest_pointer_path.expanduser().resolve() if latest_pointer_path else None + ) + self._case_groups: dict[str, dict[str, Any]] = {} + self._latest_failed_rollout: Path | None = None + + def record_rollout_completion( + self, + *, + rollout: Rollout, + index: int, + context: Any, + ) -> None: + metadata = dict(getattr(context, "metadata", {}) or {}) + training = bool(metadata.get("training")) + epoch = int(metadata.get("epoch", 0) or 0) + stage = _stage_from_execution_metadata(metadata) + commit_index = index if training else None + records = [ + _RolloutRecord( + rollout=rollout, + evaluation=_rollout_evaluation_or_default(rollout), + stage=stage, + epoch=epoch, + commit_index=commit_index, + artifact_state="rollout_done" if training else "complete", + ) + ] + for group_id, group_records in self._group_records(records).items(): + self._write_group(group_id, group_records) + self._write_index() + + def record_eval( + self, + *, + label: str, + epoch: int, + analyses: list[RolloutAnalysis], + ) -> None: + grouped = self._group_records( + [ + _RolloutRecord( + rollout=rollout, + evaluation=analysis.evaluation, + stage=_stage_dir(label, epoch=epoch), + epoch=epoch, + ) + for analysis in analyses + if isinstance((rollout := analysis.metadata.get("rollout")), Rollout) + ] + ) + for group_id, records in grouped.items(): + self._write_group(group_id, records) + + def on_train_rollout_end( + self, + *, + epoch: int, + rollouts: list[Any], + snapshot_id: str, + policy_set: Any, + context: Any, + ) -> None: + del snapshot_id, policy_set, context + self.record_train_rollouts(epoch=epoch, rollouts=list(rollouts)) + self._write_index() + + def record_train_rollouts( + self, + *, + epoch: int, + rollouts: list[Rollout], + ) -> None: + records = [ + _RolloutRecord( + rollout=rollout, + evaluation=_rollout_evaluation_or_default(rollout), + stage=_stage_dir("train_rollout", epoch=epoch), + epoch=epoch, + commit_index=idx, + artifact_state="rollout_done", + ) + for idx, rollout in enumerate(rollouts) + ] + grouped = self._group_records(records) + for group_id, group_records in grouped.items(): + self._write_group(group_id, group_records) + + async def record_train_epoch( + self, + *, + epoch: int, + analyses: list[RolloutAnalysis], + commit_results: list[dict[str, Any]], + ) -> None: + commit_by_index = { + int(item["index"]): item + for item in commit_results + if isinstance(item, dict) and item.get("index") is not None + } + records: list[_RolloutRecord] = [] + for idx, analysis in enumerate(analyses): + rollout = analysis.metadata.get("rollout") + if not isinstance(rollout, Rollout): + continue + commit_result = commit_by_index.get(idx) + records.append( + _RolloutRecord( + rollout=rollout, + evaluation=analysis.evaluation, + stage=_stage_dir("train_rollout", epoch=epoch), + epoch=epoch, + commit_result=commit_result, + commit_index=idx, + artifact_state=_artifact_state_from_commit_result(commit_result), + ) + ) + grouped = self._group_records(records) + for group_id, group_records in grouped.items(): + self._rewrite_commit_artifact_group(group_id, group_records) + await self._write_train_commit_artifacts(group_records) + + def record_train_commit_result(self, event: str, **fields: Any) -> None: + if event not in {"train_commit_submitted", "train_commit_done", "train_commit_failed"}: + return + train_dir = self._train_rollout_dir_from_event_fields(fields) + commit_dir = self._commit_rollout_dir_from_event_fields(fields) + if train_dir is None or commit_dir is None: + return + commit_result = _commit_result_from_event(event, fields) + _write_json(commit_dir / "commit_result.json", commit_result) + status_path = train_dir / "status.json" + if status_path.exists(): + try: + status = json.loads(status_path.read_text(encoding="utf-8")) + except json.JSONDecodeError: + status = {} + status.update( + { + "artifact_state": _artifact_state_from_commit_event(event), + "commit_error": commit_result.get("error"), + "commit_task_status": commit_result.get("task_status"), + "archive_uri": commit_result.get("archive_uri"), + "commit_path": str(commit_dir), + "commit_result_path": str(commit_dir / "commit_result.json"), + } + ) + _write_json(status_path, status) + + if commit_result.get("error"): + self._latest_failed_rollout = train_dir + self._update_rollout_index_entry( + path=str(train_dir), + updates={ + "artifact_state": _artifact_state_from_commit_event(event), + "commit_error": commit_result.get("error"), + "archive_uri": commit_result.get("archive_uri"), + "commit_task_status": commit_result.get("task_status"), + "commit_path": str(commit_dir), + "commit_result_path": str(commit_dir / "commit_result.json"), + }, + ) + self._write_index() + + async def on_epoch_end( + self, + *, + epoch_result: PipelineEpochResult, + policy_set: Any, + context: Any, + ) -> None: + """Lifecycle hook: write rollout artifacts immediately after each training epoch. + + This ensures rollouts are persisted incrementally instead of waiting + for the full pipeline to finish, which is important for long runs and + crash recovery. + """ + commit_results = list( + epoch_result.apply_result.metadata.get("commit_results", []) or [] + ) + await self.record_train_epoch( + epoch=epoch_result.epoch, + analyses=epoch_result.analyses, + commit_results=commit_results, + ) + # Also update the index file incrementally so it stays current. + self._write_index() + + def on_eval_end( + self, + *, + evaluation_result: PipelineEvaluationResult, + policy_set: Any, + context: Any, + ) -> None: + """Lifecycle hook: write eval rollout artifacts immediately after each eval pass.""" + label = str( + evaluation_result.metadata.get("rollout_stage") or "test_rollout" + ) + self.record_eval( + label=label, + epoch=evaluation_result.epoch, + analyses=evaluation_result.analyses, + ) + self._write_index() + + def finalize(self) -> RolloutArtifactIndex: + return self._write_index() + + def _write_index(self) -> RolloutArtifactIndex: + """Write rollouts_index.json with current state (incremental update).""" + case_groups = sorted(self._case_groups.values(), key=lambda item: item["case_group_id"]) + index = RolloutArtifactIndex( + run_dir=str(self.run_dir), + rollouts_root=str(self.rollouts_root), + case_groups=case_groups, + latest_failed_rollout=str(self._latest_failed_rollout) if self._latest_failed_rollout else None, + ) + self.run_dir.mkdir(parents=True, exist_ok=True) + index_path = self.run_dir / "rollouts_index.json" + index_path.write_text( + json.dumps(index.to_dict(), ensure_ascii=False, indent=2), + encoding="utf-8", + ) + if case_groups: + self.rollouts_root.mkdir(parents=True, exist_ok=True) + if self.latest_pointer_path is not None: + self.latest_pointer_path.parent.mkdir(parents=True, exist_ok=True) + self.latest_pointer_path.write_text(str(self.rollouts_root) + "\n", encoding="utf-8") + return index + + def _group_records(self, records: list["_RolloutRecord"]) -> dict[str, list["_RolloutRecord"]]: + grouped: dict[str, list[_RolloutRecord]] = {} + for record in records: + grouped.setdefault(_case_group_id(record.rollout), []).append(record) + return grouped + + def _write_group(self, group_id: str, records: list["_RolloutRecord"]) -> None: + if not records: + return + group_dir = self.rollouts_root / group_id + group_dir.mkdir(parents=True, exist_ok=True) + case = records[0].rollout.case + _write_json(group_dir / "case.json", case_to_dict(case)) + + group_entry = self._case_groups.setdefault( + group_id, + { + "case_group_id": group_id, + "path": str(group_dir), + "case_name": _original_case_name(records[0].rollout), + "task_id": _task_id(records[0].rollout), + "task_no": _task_no(records[0].rollout), + "split": _split(records[0].rollout), + "rollouts": [], + }, + ) + + seen_paths = {item["path"] for item in group_entry["rollouts"]} + for record in records: + rollout_dir = group_dir / record.stage / _rollout_dir_name(record) + rollout_dir.mkdir(parents=True, exist_ok=True) + self._write_rollout_artifacts(rollout_dir, record) + rollout_index = _rollout_index(record, rollout_dir) + if rollout_index["path"] not in seen_paths: + group_entry["rollouts"].append(rollout_index) + seen_paths.add(rollout_index["path"]) + if not record.passed or _commit_failed(record.commit_result): + self._latest_failed_rollout = rollout_dir + self._write_group_readme(group_dir, group_entry) + + def _write_rollout_artifacts(self, rollout_dir: Path, record: "_RolloutRecord") -> None: + rollout = record.rollout + _write_json(rollout_dir / "status.json", _status_payload(record)) + _write_json(rollout_dir / "rollout.json", _rollout_payload(record)) + _write_json(rollout_dir / "messages.json", [message.to_dict() for message in rollout.messages]) + _write_json(rollout_dir / "tool_calls.json", _tool_calls(rollout)) + _write_json(rollout_dir / "evaluation.json", evaluation_to_dict(record.evaluation)) + (rollout_dir / "memory_context.md").write_text(_memory_context(rollout), encoding="utf-8") + task_case_skill = _task_case_experience_skill(rollout) + if task_case_skill: + (rollout_dir / "task_case_experience_skill.md").write_text( + task_case_skill, + encoding="utf-8", + ) + (rollout_dir / "prompt_for_llm.md").write_text(_prompt_for_llm(record), encoding="utf-8") + # Full commit messages (as sent to session.commit) + commit_msgs = _build_commit_messages(rollout) + _write_json(rollout_dir / "commit_messages.json", commit_msgs) + (rollout_dir / "commit_messages.md").write_text( + _format_commit_messages_markdown(commit_msgs), encoding="utf-8" + ) + + async def _write_train_commit_artifacts(self, records: list["_RolloutRecord"]) -> None: + if self.client is None: + return + for record in records: + if record.commit_result is None: + continue + archive_uri = str(record.commit_result.get("archive_uri") or "").strip() + if not archive_uri: + continue + train_dir = self._train_rollout_dir(record) + commit_dir = self._commit_rollout_dir(record) + try: + memory_diff = await self.client.read(f"{archive_uri}/memory_diff.json") + except Exception as exc: # best-effort artifact enrichment + _write_json( + commit_dir / "memory_diff_error.json", + {"archive_uri": archive_uri, "error": str(exc)}, + ) + self._update_rollout_status( + train_dir, + memory_diff_error=str(exc), + commit_path=str(commit_dir), + ) + self._update_rollout_index_entry( + path=str(train_dir), + updates={"memory_diff_error": str(exc), "commit_path": str(commit_dir)}, + ) + self._write_index() + continue + (commit_dir / "memory_diff.json").write_text(str(memory_diff), encoding="utf-8") + (commit_dir / "memory_diff.md").write_text( + _format_memory_diff_markdown(_parse_memory_diff(memory_diff)), + encoding="utf-8", + ) + self._update_rollout_status( + train_dir, + artifact_state="memory_diff_done", + memory_diff_path=str(commit_dir / "memory_diff.json"), + memory_diff_markdown_path=str(commit_dir / "memory_diff.md"), + commit_path=str(commit_dir), + ) + self._update_rollout_index_entry( + path=str(train_dir), + updates={ + "artifact_state": "memory_diff_done", + "memory_diff_path": str(commit_dir / "memory_diff.json"), + "memory_diff_markdown_path": str(commit_dir / "memory_diff.md"), + "commit_path": str(commit_dir), + }, + ) + self._write_index() + + def _update_rollout_status(self, rollout_dir: Path, **updates: Any) -> None: + status_path = rollout_dir / "status.json" + if not status_path.exists(): + return + try: + status = json.loads(status_path.read_text(encoding="utf-8")) + except json.JSONDecodeError: + status = {} + status.update(updates) + _write_json(status_path, status) + + def _rewrite_commit_artifact_group( + self, + group_id: str, + records: list["_RolloutRecord"], + ) -> None: + group_entry = self._case_groups.get(group_id) + if group_entry is None: + self._write_group(group_id, records) + group_entry = self._case_groups.get(group_id) + if group_entry is None: + return + for record in records: + train_dir = self._train_rollout_dir(record) + commit_dir = self._commit_rollout_dir(record) + if record.commit_result is not None: + _write_json(commit_dir / "commit_result.json", record.commit_result) + updates = { + "artifact_state": record.artifact_state, + "commit_error": ( + record.commit_result.get("error") if record.commit_result else None + ), + "archive_uri": ( + record.commit_result.get("archive_uri") if record.commit_result else None + ), + "commit_task_status": ( + record.commit_result.get("task_status") if record.commit_result else None + ), + "commit_path": str(commit_dir), + "commit_result_path": str(commit_dir / "commit_result.json") + if record.commit_result is not None + else None, + } + self._update_rollout_status(train_dir, **updates) + self._update_rollout_index_entry(path=str(train_dir), updates=updates) + if not record.passed or _commit_failed(record.commit_result): + self._latest_failed_rollout = train_dir + self._write_group_readme(self.rollouts_root / group_id, group_entry) + + def _update_rollout_index_entry(self, *, path: str, updates: dict[str, Any]) -> None: + for group_entry in self._case_groups.values(): + for item in group_entry.get("rollouts", []): + if item.get("path") == path: + item.update(updates) + return + + def _train_rollout_dir(self, record: "_RolloutRecord") -> Path: + return ( + self.rollouts_root + / _case_group_id(record.rollout) + / f"epoch_{record.epoch}" + / _stage_leaf("train_rollout") + / _rollout_dir_name(record) + ) + + def _commit_rollout_dir(self, record: "_RolloutRecord") -> Path: + return ( + self.rollouts_root + / _case_group_id(record.rollout) + / f"epoch_{record.epoch}" + / _stage_leaf("train") + / _rollout_dir_name(record) + ) + + def _train_rollout_dir_from_event_fields(self, fields: dict[str, Any]) -> Path | None: + return self._rollout_dir_from_event_fields(fields, phase=_stage_leaf("train_rollout")) + + def _commit_rollout_dir_from_event_fields(self, fields: dict[str, Any]) -> Path | None: + return self._rollout_dir_from_event_fields(fields, phase=_stage_leaf("train")) + + def _rollout_dir_from_event_fields(self, fields: dict[str, Any], *, phase: str) -> Path | None: + split = fields.get("split") + task_no = fields.get("task_no") + task_id = fields.get("case_task_id") or fields.get("case_name") + epoch = fields.get("epoch") + index = fields.get("index") + if split is None or task_no is None or task_id is None or epoch is None or index is None: + return None + group_id = ( + f"{_safe_fragment(split)}_task_" + f"{_safe_fragment(str(task_no))}_" + f"{_safe_fragment(task_id)}" + )[:120] + return self.rollouts_root / group_id / f"epoch_{epoch}" / phase / f"trial_{index}" + + def _write_group_readme(self, group_dir: Path, group_entry: dict[str, Any]) -> None: + failed = [item for item in group_entry["rollouts"] if not item.get("passed") or item.get("commit_error")] + lines = [ + f"# Rollout artifact group: {group_entry['case_group_id']}", + "", + f"- split: {group_entry.get('split')}", + f"- task_no: {group_entry.get('task_no')}", + f"- task_id: {group_entry.get('task_id')}", + f"- rollouts: {len(group_entry['rollouts'])}", + f"- failed_rollouts: {len(failed)}", + "", + "## Rollouts", + ] + for item in group_entry["rollouts"]: + status = "FAIL" if (not item.get("passed") or item.get("commit_error")) else "PASS" + lines.append( + f"- [{status}] {item.get('stage')} {item.get('rollout_name')} " + f"score={item.get('score')} path={item.get('path')}" + ) + lines.extend( + [ + "", + "## Suggested LLM prompt", + "", + "Read this directory recursively. Compare successful and failed rollouts for the same task. ", + "Focus on whether the injected memory_context.md was missing, misleading, ignored, or helpful.", + ] + ) + (group_dir / "README.md").write_text("\n".join(lines) + "\n", encoding="utf-8") + + +@dataclass(slots=True) +class _RolloutRecord: + rollout: Rollout + evaluation: Any + stage: str + epoch: int + commit_result: dict[str, Any] | None = None + commit_index: int | None = None + artifact_state: str = "complete" + + @property + def passed(self) -> bool: + return bool(getattr(self.evaluation, "passed", False)) + + @property + def score(self) -> float: + return float(getattr(self.evaluation, "score", 0.0) or 0.0) + + +def _write_json(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(jsonable(value), ensure_ascii=False, indent=2), encoding="utf-8") + + +@dataclass(slots=True) +class RolloutArtifactEventRecorder: + """Event recorder adapter that enriches rollout artifacts from commit events.""" + + recorder: RolloutArtifactRecorder + + def record(self, event: str, **fields: Any) -> None: + self.recorder.record_train_commit_result(event, **fields) + + +def _rollout_evaluation_or_default(rollout: Rollout) -> Any: + if rollout.evaluation is not None: + return rollout.evaluation + from openviking.session.train.components.session_commit import ( + _rollout_evaluation_or_default as default_evaluation, + ) + + return default_evaluation(rollout) + + +def _commit_result_from_event(event: str, fields: dict[str, Any]) -> dict[str, Any]: + return { + "index": fields.get("index"), + "session_id": fields.get("session_id"), + "stage": fields.get("stage"), + "task_id": fields.get("task_id"), + "archive_uri": fields.get("archive_uri"), + "trace_id": fields.get("trace_id"), + "telemetry_id": fields.get("telemetry_id"), + "task_status": fields.get("task_status"), + "score": fields.get("score"), + "error": fields.get("error"), + "event": event, + "artifact_state": _artifact_state_from_commit_event(event), + } + + +def _artifact_state_from_commit_event(event: str) -> str: + if event == "train_commit_submitted": + return "commit_submitted" + if event == "train_commit_done": + return "commit_done" + if event == "train_commit_failed": + return "commit_failed" + return "rollout_done" + + +def _artifact_state_from_commit_result(commit_result: dict[str, Any] | None) -> str: + if not commit_result: + return "rollout_done" + if commit_result.get("error"): + return "commit_failed" + return "commit_done" + + +def _stage_from_execution_metadata(metadata: dict[str, Any]) -> str: + stage = str(metadata.get("rollout_stage") or metadata.get("stage") or "") + epoch = int(metadata.get("epoch", 0) or 0) + if not stage: + stage = "train_rollout" if bool(metadata.get("training")) else "test_rollout" + return _stage_dir(stage.split(maxsplit=1)[0], epoch=epoch) + + +def _status_payload(record: _RolloutRecord) -> dict[str, Any]: + rollout = record.rollout + return { + "stage": record.stage, + "epoch": record.epoch, + "rollout_name": _rollout_name(record), + "case_group_id": _case_group_id(rollout), + "case_name": rollout.case.name, + "original_case_name": _original_case_name(rollout), + "split": _split(rollout), + "task_no": _task_no(rollout), + "task_id": _task_id(rollout), + "trial": _trial(rollout), + "passed": record.passed, + "score": record.score, + "policy_snapshot_id": rollout.policy_snapshot_id, + "has_memory_context": bool(_memory_context(rollout).strip()), + "has_task_case_experience_skill": bool(_task_case_experience_skill(rollout).strip()), + "task_case_experience_skill_path": "task_case_experience_skill.md" + if _task_case_experience_skill(rollout).strip() + else None, + "artifact_state": record.artifact_state, + "commit_error": record.commit_result.get("error") if record.commit_result else None, + "commit_task_status": ( + record.commit_result.get("task_status") if record.commit_result else None + ), + "archive_uri": record.commit_result.get("archive_uri") if record.commit_result else None, + } + + +def _rollout_payload(record: _RolloutRecord) -> dict[str, Any]: + rollout = record.rollout + return { + "case": case_to_dict(rollout.case), + "policy_snapshot_id": rollout.policy_snapshot_id, + "metadata": jsonable(rollout.metadata), + "evaluation": evaluation_to_dict(record.evaluation), + } + + +def _rollout_index(record: _RolloutRecord, rollout_dir: Path) -> dict[str, Any]: + return { + "rollout_name": _rollout_name(record), + "stage": record.stage, + "epoch": record.epoch, + "trial": _trial(record.rollout), + "passed": record.passed, + "score": record.score, + "artifact_state": record.artifact_state, + "path": str(rollout_dir), + "commit_error": record.commit_result.get("error") if record.commit_result else None, + "archive_uri": record.commit_result.get("archive_uri") if record.commit_result else None, + "commit_task_status": ( + record.commit_result.get("task_status") if record.commit_result else None + ), + } + + +def _tool_calls(rollout: Rollout) -> list[dict[str, Any]]: + calls: list[dict[str, Any]] = [] + for message_index, message in enumerate(rollout.messages): + for part in message.parts: + if isinstance(part, ToolPart): + calls.append( + { + "message_index": message_index, + "message_id": message.id, + "role": message.role, + "tool_id": part.tool_id, + "tool_name": part.tool_name, + "tool_status": part.tool_status, + "tool_input": jsonable(part.tool_input), + "tool_output": part.tool_output, + } + ) + return calls + + +def _prompt_for_llm(record: _RolloutRecord) -> str: + status = _status_payload(record) + return "\n".join( + [ + "# Analyze this rollout", + "", + "Read all files in this directory, especially:", + "- memory_context.md: memory injected into the agent prompt at rollout time", + "- messages.json and tool_calls.json: trajectory", + "- evaluation.json: failure signal", + "- memory_diff.json: training memory update result when present", + "", + "## Status", + "", + "```json", + json.dumps(jsonable(status), ensure_ascii=False, indent=2), + "```", + "", + "Please identify whether the failure is caused by missing memory, " + "wrong memory, ignored memory, bad tool use, or task ambiguity.", + ] + ) + "\n" + + +def _memory_context(rollout: Rollout) -> str: + metadata = rollout.metadata or {} + value = metadata.get("memory") + if value is None: + return "" + return str(value) + + +def _task_case_experience_skill(rollout: Rollout) -> str: + metadata = rollout.metadata or {} + value = metadata.get("task_case_experience_skill") + if value is None: + return "" + return str(value) + + +def _case_group_id(rollout: Rollout) -> str: + split = _safe_fragment(_split(rollout) or "split") + task_no = _safe_fragment( + str(_task_no(rollout) if _task_no(rollout) is not None else "x") + ) + task_id = _safe_fragment( + str(_task_id(rollout) or _original_case_name(rollout) or rollout.case.name) + ) + return f"{split}_task_{task_no}_{task_id}"[:120] + + +def _rollout_dir_name(record: _RolloutRecord) -> str: + return _safe_fragment(_rollout_name(record)) + + +def _rollout_name(record: _RolloutRecord) -> str: + trial = _trial(record.rollout) + if trial is not None: + return f"trial_{trial}" + if record.commit_index is not None: + return f"trial_{record.commit_index}" + return _safe_fragment(record.rollout.case.name) + + +def _stage_dir(label: str, *, epoch: int | None = None) -> str: + stage = _stage_leaf(label) + return stage if epoch is None else f"epoch_{epoch}/{stage}" + + +def _stage_leaf(label: str) -> str: + stage = _safe_fragment(label or "rollout") + order = { + "train_rollout": 1, + "train": 2, + "eval_train_rollout": 3, + "test_rollout": 4, + "baseline_test_rollout": 0, + "final_test_rollout": 5, + }.get(stage) + return f"{order}.{stage}" if order is not None else stage + + +def _original_case_name(rollout: Rollout) -> str: + return str( + rollout.case.input.get("original_case_name") + or rollout.case.metadata.get("original_case_name") + or rollout.metadata.get("original_case_name") + or rollout.case.name + ) + + +def _split(rollout: Rollout) -> str | None: + value = ( + rollout.case.input.get("data_split") + or rollout.metadata.get("data_split") + or rollout.case.input.get("split") + or rollout.metadata.get("split") + ) + return str(value) if value is not None else None + + +def _task_no(rollout: Rollout) -> Any: + return rollout.case.input.get("task_no", rollout.metadata.get("task_no")) + + +def _task_id(rollout: Rollout) -> Any: + return rollout.case.input.get("task_id", rollout.metadata.get("task_id")) + + +def _trial(rollout: Rollout) -> Any: + for key in ("eval_trial", "train_trial"): + if key in rollout.case.input: + return rollout.case.input.get(key) + if key in rollout.case.metadata: + return rollout.case.metadata.get(key) + if key in rollout.metadata: + return rollout.metadata.get(key) + return None + + +def _commit_failed(commit_result: dict[str, Any] | None) -> bool: + return bool(commit_result and commit_result.get("error")) + + +def _safe_fragment(value: Any) -> str: + text = str(value) + text = re.sub(r"[^A-Za-z0-9_.-]+", "_", text).strip("_") + return text or "unknown" + + +def _build_commit_messages(rollout: Rollout) -> list[dict[str, Any]]: + """Build the full message list as sent to session.commit. + + Matches the message assembly in session_commit._commit_one: + [case_spec] + rollout.messages + [evaluation] + """ + from openviking.session.train.components.session_commit import ( + _case_spec_message_to_request, + _evaluation_message_to_request, + _message_to_request, + ) + + messages: list[dict[str, Any]] = [_case_spec_message_to_request(rollout)] + for msg in rollout.messages: + messages.append(_message_to_request(msg)) + messages.append(_evaluation_message_to_request(rollout)) + return messages + + +def _format_commit_messages_markdown(messages: list[dict[str, Any]]) -> str: + """Format commit messages as a readable Markdown document.""" + lines: list[str] = ["# Commit Messages", ""] + for idx, msg in enumerate(messages): + role = msg.get("role", "unknown") + parts = msg.get("parts", []) + lines.append(f"## [{idx}] {role}") + lines.append("") + for part in parts: + part_type = part.get("type", "text") + if part_type == "text": + text = part.get("text", "") + # Indent to make it a blockquote / code block if needed + lines.append(text) + elif part_type == "tool": + status = str(part.get("tool_status") or "") + label = "Tool result" if status in {"completed", "error"} else "Tool call" + lines.append(f"**{label}:** `{part.get('tool_name', '?')}` status={status or '?'}") + if part.get("tool_input") is not None: + lines.append("") + lines.append("```json") + lines.append(json.dumps(part.get("tool_input"), ensure_ascii=False, indent=2)) + lines.append("```") + if part.get("tool_output"): + content = str(part.get("tool_output", "")) + lines.append("") + lines.append("```") + lines.append(content) + lines.append("```") + elif part_type == "tool_call": + lines.append(f"**Tool call:** `{part.get('tool_name', '?')}`") + lines.append("") + lines.append("```json") + lines.append(json.dumps(part.get("tool_input", {}), ensure_ascii=False, indent=2)) + lines.append("```") + elif part_type == "tool_result": + lines.append(f"**Tool result:** `{part.get('tool_name', '?')}`") + lines.append("") + content = str(part.get("text", part.get("tool_result", ""))) + lines.append("```") + lines.append(content) + lines.append("```") + else: + lines.append(f"*[{part_type} part]*") + lines.append("") + lines.append("---") + lines.append("") + return "\n".join(lines) + + +def _parse_memory_diff(memory_diff: Any) -> dict[str, Any]: + if isinstance(memory_diff, dict): + return memory_diff + if isinstance(memory_diff, str): + try: + parsed = json.loads(memory_diff) + except json.JSONDecodeError: + return {} + return parsed if isinstance(parsed, dict) else {} + return {} + + +def _format_memory_diff_markdown(memory_diff: dict[str, Any]) -> str: + lines = ["# Memory Diff", ""] + summary = memory_diff.get("summary") if isinstance(memory_diff, dict) else None + if isinstance(summary, dict): + lines.extend( + [ + "## Summary", + "", + f"- adds: {summary.get('total_adds', 0)}", + f"- updates: {summary.get('total_updates', 0)}", + f"- deletes: {summary.get('total_deletes', 0)}", + "", + ] + ) + operations = memory_diff.get("operations") if isinstance(memory_diff, dict) else None + operations = operations if isinstance(operations, dict) else {} + for item in operations.get("adds", []) or []: + if isinstance(item, dict): + lines.extend(_memory_diff_file_section("add", item.get("uri"), "", item.get("after", ""))) + for item in operations.get("updates", []) or []: + if isinstance(item, dict): + lines.extend( + _memory_diff_file_section( + "update", + item.get("uri"), + item.get("before", ""), + item.get("after", ""), + ) + ) + return "\n".join(lines).rstrip() + "\n" + + +def _memory_diff_file_section(kind: str, uri: Any, before: Any, after: Any) -> list[str]: + path = str(uri or "unknown") + old_path = "/dev/null" if kind == "add" else path + diff = difflib.unified_diff( + str(before or "").splitlines(), + str(after or "").splitlines(), + fromfile=old_path, + tofile=path, + lineterm="", + ) + return [ + f"## {kind}: `{path}`", + "", + "```diff", + *diff, + "```", + "", + ] diff --git a/openviking/session/train/components/rollout_executor.py b/openviking/session/train/components/rollout_executor.py new file mode 100644 index 0000000000..a2b5ce2a44 --- /dev/null +++ b/openviking/session/train/components/rollout_executor.py @@ -0,0 +1,125 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""RolloutExecutor component implementations.""" + +from __future__ import annotations + +import json +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any +from uuid import uuid4 + +from openviking.message import Message, TextPart +from openviking.session.train.context import ExecutionContext +from openviking.session.train.domain import Case, ExperienceSet, Rollout +from openviking.telemetry import tracer +from openviking_cli.utils.config import get_openviking_config + +PromptBuilder = Callable[[Case, ExperienceSet, ExecutionContext], str] + + +@dataclass(slots=True) +class SingleTurnLLMRolloutExecutor: + """Execute each Case with one plain LLM call. + + This is a minimal RolloutExecutor for offline training bootstrap. It does + not run tools or a full agent loop; future agent-loop components can implement + the same RolloutExecutor interface. + """ + + vlm: Any = None + prompt_builder: PromptBuilder | None = None + thinking: bool | None = None + + @tracer("train.rollout_executor.single_turn.execute", ignore_result=True, ignore_args=True) + async def execute( + self, + cases: list[Case], + policy_set: ExperienceSet, + context: ExecutionContext, + ) -> list[Rollout]: + vlm = self.vlm or get_openviking_config().vlm + rollouts: list[Rollout] = [] + for case in cases: + prompt = self._build_prompt(case, policy_set, context) + response = await vlm.get_completion_async(prompt=prompt, thinking=self.thinking) + assistant_text = _response_text(response) + rollouts.append( + Rollout( + case=case, + messages=[ + Message( + id=f"rollout-user-{uuid4().hex}", + role="user", + parts=[TextPart(text=prompt)], + ), + Message( + id=f"rollout-assistant-{uuid4().hex}", + role="assistant", + parts=[TextPart(text=assistant_text)], + ), + ], + policy_snapshot_id=context.policy_snapshot_id, + ) + ) + return rollouts + + def _build_prompt( + self, + case: Case, + policy_set: ExperienceSet, + context: ExecutionContext, + ) -> str: + if self.prompt_builder is not None: + return self.prompt_builder(case, policy_set, context) + return default_single_turn_prompt(case, policy_set, context) + + +def default_single_turn_prompt( + case: Case, + policy_set: ExperienceSet, + context: ExecutionContext, +) -> str: + """Build a simple prompt containing policy experiences and case input.""" + + experiences = "\n\n".join( + f"### {policy.name} v{policy.version} [{policy.status}]\n{policy.content}" + for policy in policy_set.policies + ) + if not experiences: + experiences = "(no experience policies available)" + + return "\n".join( + [ + "You are executing an offline training case for OpenViking.", + "Use the current experience policies when they are relevant.", + "Return the best final answer/action for the case.", + "", + f"Policy snapshot: {context.policy_snapshot_id}", + "", + "# Experience Policies", + experiences, + "", + "# Case", + f"Name: {case.name}", + f"Task signature: {case.task_signature}", + "Input:", + json.dumps(case.input, ensure_ascii=False, indent=2, sort_keys=True), + "", + "# Rubric", + f"{case.rubric.name}: {case.rubric.description}", + *[ + f"- {criterion.name} ({'required' if criterion.required else 'optional'}, " + f"weight={criterion.weight}): {criterion.description}" + for criterion in case.rubric.criteria + ], + ] + ) + + +def _response_text(response: Any) -> str: + content = getattr(response, "content", None) + if content is not None: + return str(content) + return str(response or "") diff --git a/openviking/session/train/components/session_commit.py b/openviking/session/train/components/session_commit.py new file mode 100644 index 0000000000..e927bb1c56 --- /dev/null +++ b/openviking/session/train/components/session_commit.py @@ -0,0 +1,550 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""PolicyTrainer implementation backed by OpenViking session.commit.""" + +from __future__ import annotations + +import asyncio +import json +import time +from dataclasses import dataclass +from typing import Any +from uuid import uuid4 + +from openviking.session.train.components.progress import run_with_progress +from openviking.session.train.context import PipelineContext +from openviking.session.train.domain import ( + CriterionResult, + ExperienceSet, + PolicyApplyResult, + PolicyUpdatePlan, + Rollout, + RolloutAnalysis, + RolloutTrainingResult, + RubricEvaluation, +) +from openviking.session.train.utils import average_score, validate_rollouts_have_cases +from openviking_cli.client.http import AsyncHTTPClient + +_TRAINING_COMMIT_MEMORY_TYPES = ("cases", "trajectories", "experiences") +_TRAINING_CASE_SPEC_PROTOCOL = "openviking.batch_train.case_spec.v1" +_TRAINING_CASE_SPEC_HEADER = "# OpenViking Batch Training CaseSpec v1" +_SESSION_BATCH_ADD_MESSAGE_LIMIT = 100 + + +@dataclass(slots=True) +class SessionCommitPolicyTrainer: + """Train remotely by writing rollout messages to sessions and committing them.""" + + client: AsyncHTTPClient + run_id: str = "" + keep_recent_count: int = 0 + poll_interval_seconds: float = 2.0 + timeout_seconds: float | None = None + commit_concurrency: int = 20 + show_progress: bool = False + progress_label: str = "session-commit" + event_recorder: Any | None = None + + def __post_init__(self) -> None: + if not self.run_id: + self.run_id = _new_run_id() + if self.poll_interval_seconds <= 0: + raise ValueError("poll_interval_seconds must be > 0") + if self.timeout_seconds is not None and self.timeout_seconds <= 0: + raise ValueError("timeout_seconds must be > 0") + if self.commit_concurrency <= 0: + raise ValueError("commit_concurrency must be > 0") + + async def train_rollouts( + self, + rollouts: list[Rollout], + policy_set: ExperienceSet, + context: PipelineContext | Any = None, + analyses: list[RolloutAnalysis] | None = None, + ) -> RolloutTrainingResult: + rollout_list = list(rollouts) + validate_rollouts_have_cases(rollout_list) + if analyses is not None and len(analyses) != len(rollout_list): + raise ValueError( + "SessionCommitPolicyTrainer analyses length must match rollouts length when provided" + ) + execution_metadata = dict(getattr(context, "execution_metadata", {}) or {}) + + async def _commit(rollout: Rollout, idx: int) -> dict[str, Any]: + return await self._commit_one( + rollout, + idx, + execution_metadata=execution_metadata, + ) + + commit_results = await run_with_progress( + rollout_list, + coroutine_factory=_commit, + label="train_start", + enabled=self.show_progress, + description=( + f"Processing {len(rollout_list)} rollouts, " + f"concurrency={self.commit_concurrency}" + ), + concurrency=self.commit_concurrency, + ) + analysis_list = [_analysis_from_rollout(rollout) for rollout in rollout_list] + errors = [item["error"] for item in commit_results if item.get("error")] + apply_result = PolicyApplyResult( + updated_policy_set=policy_set, + errors=errors, + metadata={ + "committed_rollout_count": len(commit_results), + "commit_results": commit_results, + "run_id": self.run_id, + }, + ) + return RolloutTrainingResult( + analyses=analysis_list, + gradients=[], + plan=PolicyUpdatePlan(metadata={"trainer": "session_commit", "run_id": self.run_id}), + apply_result=apply_result, + metadata={ + "policy_set_root_uri": policy_set.root_uri, + "rollout_count": len(rollout_list), + "analysis_count": len(analysis_list), + "gradient_count": 0, + "score": average_score(analysis_list), + "source": "session_commit_trainer", + "run_id": self.run_id, + }, + ) + + async def _commit_one( + self, + rollout: Rollout, + index: int, + *, + execution_metadata: dict[str, Any] | None = None, + ) -> dict[str, Any]: + session_id = _session_id_for_rollout(rollout, run_id=self.run_id) + stage = "prepare_messages" + try: + messages = ( + [_case_spec_message_to_request(rollout)] + + [_message_to_request(message) for message in rollout.messages] + + [_evaluation_message_to_request(rollout)] + ) + stage = "create_session" + await self.client.create_session( + session_id=session_id, + memory_policy=_training_commit_memory_policy(), + ) + stage = "batch_add_messages" + await self._batch_add_messages(session_id, messages) + stage = "commit_session" + commit_result = await self.client.commit_session( + session_id, + telemetry=True, + keep_recent_count=self.keep_recent_count, + ) + task_id = str(commit_result.get("task_id") or "") + archive_uri = str(commit_result.get("archive_uri") or "") + trace_id = _commit_trace_id(commit_result) + telemetry_id = _commit_telemetry_id(commit_result) + await self._record_event( + "train_commit_submitted", + rollout=rollout, + index=index, + session_id=session_id, + stage=stage, + execution_metadata=execution_metadata, + task_id=task_id, + archive_uri=archive_uri, + trace_id=trace_id, + telemetry_id=telemetry_id, + score=_rollout_score(rollout), + ) + stage = "wait_task" + task = await self._wait_task(task_id) if task_id else None + task_error = _task_error(task) + if task_error: + print( + f"[session_commit] failed stage={stage} session_id={session_id} " + f"task_id={task_id} trace_id={trace_id or ''} " + f"error={task_error}", + flush=True, + ) + await self._record_event( + "train_commit_failed" if task_error else "train_commit_done", + rollout=rollout, + index=index, + session_id=session_id, + stage=stage, + execution_metadata=execution_metadata, + task_id=task_id, + archive_uri=archive_uri, + trace_id=trace_id, + telemetry_id=telemetry_id, + task_status=task.get("status") if isinstance(task, dict) else None, + score=_rollout_score(rollout), + error=task_error, + ) + return { + "index": index, + "session_id": session_id, + "stage": stage, + "task_id": task_id, + "archive_uri": archive_uri, + "trace_id": trace_id, + "telemetry_id": telemetry_id, + "task_status": task.get("status") if isinstance(task, dict) else None, + "score": _rollout_score(rollout), + "error": task_error, + } + except Exception as exc: + print( + f"[session_commit] failed stage={stage} session_id={session_id} " + f"task_id= trace_id= error={exc}", + flush=True, + ) + await self._record_event( + "train_commit_failed", + rollout=rollout, + index=index, + session_id=session_id, + stage=stage, + execution_metadata=execution_metadata, + task_id="", + archive_uri="", + trace_id=None, + telemetry_id=None, + task_status="failed", + score=_rollout_score(rollout), + error=str(exc), + ) + return { + "index": index, + "session_id": session_id, + "stage": stage, + "task_id": "", + "archive_uri": "", + "trace_id": None, + "telemetry_id": None, + "task_status": "failed", + "score": _rollout_score(rollout), + "error": str(exc), + } + + + async def _record_event( + self, + event: str, + *, + rollout: Rollout, + index: int, + session_id: str, + stage: str, + execution_metadata: dict[str, Any] | None = None, + **fields: Any, + ) -> None: + if self.event_recorder is None: + return + record = getattr(self.event_recorder, "record", None) + if record is None: + return + payload = { + "index": index, + "stage": stage, + "session_id": session_id, + **_rollout_event_fields( + rollout, + execution_metadata=execution_metadata, + ), + **fields, + } + result = record(event, **payload) + if asyncio.iscoroutine(result): + await result + + async def _batch_add_messages(self, session_id: str, messages: list[dict[str, Any]]) -> None: + for start in range(0, len(messages), _SESSION_BATCH_ADD_MESSAGE_LIMIT): + await self.client.batch_add_messages( + session_id, messages[start : start + _SESSION_BATCH_ADD_MESSAGE_LIMIT] + ) + + async def _wait_task(self, task_id: str) -> dict[str, Any]: + deadline = ( + asyncio.get_running_loop().time() + self.timeout_seconds + if self.timeout_seconds is not None + else None + ) + while True: + task = await self.client.get_task(task_id) + if task and task.get("status") in {"completed", "failed"}: + return task + if deadline is not None and asyncio.get_running_loop().time() >= deadline: + return {"task_id": task_id, "status": "timeout", "error": "commit task timeout"} + await asyncio.sleep(self.poll_interval_seconds) + + +def _rollout_event_fields( + rollout: Rollout, + *, + execution_metadata: dict[str, Any] | None = None, +) -> dict[str, Any]: + case = rollout.case + metadata = rollout.metadata or {} + rollout_execution_metadata = metadata.get("execution_metadata", {}) + if not isinstance(rollout_execution_metadata, dict): + rollout_execution_metadata = {} + event_execution_metadata = dict(rollout_execution_metadata) + event_execution_metadata.update(execution_metadata or {}) + case_input = case.input or {} + return { + "epoch": event_execution_metadata.get("epoch"), + "training": event_execution_metadata.get("training"), + "rollout_stage": event_execution_metadata.get("rollout_stage") + or event_execution_metadata.get("stage"), + "case_name": case.name, + "task_signature": case.task_signature, + "split": ( + case_input.get("data_split") + or metadata.get("data_split") + or case_input.get("split") + or metadata.get("split") + ), + "task_no": ( + case_input.get("task_no") + if case_input.get("task_no") is not None + else metadata.get("task_no") + ), + "case_task_id": case_input.get("task_id") or metadata.get("task_id"), + "task_id": case_input.get("task_id") or metadata.get("task_id"), + "policy_snapshot_id": rollout.policy_snapshot_id, + "passed": bool(rollout.evaluation.passed) if rollout.evaluation is not None else None, + } + + +def _training_commit_memory_policy() -> dict[str, Any]: + return { + "memory_types": list(_TRAINING_COMMIT_MEMORY_TYPES), + "working_memory": {"enabled": False}, + } + + +def _analysis_from_rollout(rollout: Rollout) -> RolloutAnalysis: + return RolloutAnalysis( + evaluation=_rollout_evaluation_or_default(rollout), + trajectories=[], + metadata={ + "rollout": rollout, + "rollout_messages": rollout.messages, + "policy_snapshot_id": rollout.policy_snapshot_id, + "evaluation_source": "rollout" + if rollout.evaluation is not None + else "session_commit_default", + }, + ) + + +def _rollout_evaluation_or_default(rollout: Rollout) -> RubricEvaluation: + if rollout.evaluation is not None: + return rollout.evaluation + return RubricEvaluation( + passed=False, + score=0.0, + criterion_results=[ + CriterionResult( + criterion_name="rollout_evaluation_provided", + passed=False, + score=0.0, + feedback=["Rollout executor did not provide evaluation."], + evidence=[], + metadata={"source": "session_commit_default"}, + ) + ], + feedback=["Rollout executor did not provide evaluation."], + metadata={"source": "session_commit_default"}, + ) + + +def _rollout_score(rollout: Rollout) -> float: + if rollout.evaluation is None: + return 0.0 + return float(rollout.evaluation.score) + +def _task_error(task: dict[str, Any] | None) -> str | None: + if task is None: + return None + if task.get("status") == "failed": + return str(task.get("error") or "task failed") + if task.get("status") == "timeout": + return str(task.get("error") or "task timeout") + return None + + +def _commit_trace_id(commit_result: dict[str, Any]) -> str | None: + trace_id = commit_result.get("trace_id") + return str(trace_id) if trace_id else None + + +def _commit_telemetry_id(commit_result: dict[str, Any]) -> str | None: + telemetry = commit_result.get("telemetry") + if not isinstance(telemetry, dict): + return None + telemetry_id = telemetry.get("id") + return str(telemetry_id) if telemetry_id else None + + +def _session_id_for_rollout(rollout: Rollout, *, run_id: str) -> str: + safe_name = _safe_session_fragment(rollout.case.name) + metadata = rollout.metadata or {} + execution_metadata = metadata.get("execution_metadata", {}) + epoch = execution_metadata.get("epoch", "0") + task_no = metadata.get("task_no", "0") + split = metadata.get("data_split", "tau2") + return f"tau2_train_{run_id}_{split}_e{epoch}_t{task_no}_{safe_name}" + + +def _safe_session_fragment(value: str) -> str: + return "".join(ch if ch.isalnum() or ch in "_.-" else "_" for ch in value)[:80] or "case" + + +def _new_run_id() -> str: + return f"{int(time.time())}_{uuid4().hex[:8]}" + + +def _case_spec_message_to_request(rollout: Rollout) -> dict[str, Any]: + text = ( + f"{_TRAINING_CASE_SPEC_HEADER}\n\n" + "The following structured case and rubric describe the task that " + "produced this rollout. It is control-plane metadata for the " + "batch training pipeline.\n\n" + f"```json\n{_case_spec_payload_json(rollout)}\n```" + ) + return { + "role": "system", + "parts": [{"type": "text", "text": text}], + } + + +def _case_spec_payload_json(rollout: Rollout) -> str: + import json + + return json.dumps(_case_spec_payload(rollout), ensure_ascii=False, indent=2, sort_keys=True) + + +def _case_spec_payload(rollout: Rollout) -> dict[str, Any]: + case = rollout.case + return { + "protocol": _TRAINING_CASE_SPEC_PROTOCOL, + "case": { + "name": _stable_case_name(rollout), + "task_signature": _stable_task_signature(rollout), + "input": _case_input_payload(case.input), + "metadata": _stable_case_metadata(rollout), + "rubric": { + "name": case.rubric.name, + "description": case.rubric.description, + "criteria": [ + { + "name": criterion.name, + "description": criterion.description, + "required": criterion.required, + "weight": criterion.weight, + } + for criterion in case.rubric.criteria + ], + }, + }, + } + + +def _stable_case_name(rollout: Rollout) -> str: + case = rollout.case + return str( + case.input.get("original_case_name") + or case.metadata.get("original_case_name") + or rollout.metadata.get("original_case_name") + or case.name + ) + + +def _stable_task_signature(rollout: Rollout) -> str: + case = rollout.case + if case.input.get("original_case_name") or case.metadata.get("original_case_name"): + return str(case.task_signature).split(":trial:", 1)[0] + return case.task_signature + + +def _stable_case_metadata(rollout: Rollout) -> dict[str, Any]: + metadata = dict(rollout.case.metadata or {}) + metadata.setdefault("rollout_case_name", rollout.case.name) + metadata.setdefault("rollout_task_signature", rollout.case.task_signature) + return metadata + + +def _case_input_payload(case_input: dict[str, Any]) -> dict[str, Any]: + allowed_keys = ( + "domain", + "split", + "data_split", + "task_id", + "task_no", + "user_query", + ) + return {key: case_input[key] for key in allowed_keys if key in case_input} + + +def _evaluation_message_to_request(rollout: Rollout) -> dict[str, Any]: + text = ( + "# OpenViking OutcomeEvaluation\n\n" + "The following structured evaluation describes the outcome of the " + "preceding rollout. Use it as the training signal when extracting " + "training memories.\n\n" + f"```json\n{_evaluation_payload_json(rollout)}\n```" + ) + return { + "role": "user", + "parts": [{"type": "text", "text": text}], + } + + +def _evaluation_payload_json(rollout: Rollout) -> str: + return json.dumps( + {"evaluation": _evaluation_payload(rollout.evaluation)}, + ensure_ascii=False, + indent=2, + sort_keys=True, + ) + + +def _evaluation_payload(evaluation: RubricEvaluation | None) -> dict[str, Any] | None: + if evaluation is None: + return None + return { + "passed": evaluation.passed, + "score": evaluation.score, + "feedback": evaluation.feedback, + "criterion_results": [ + { + "criterion_name": result.criterion_name, + "passed": result.passed, + "score": result.score, + "feedback": result.feedback, + "evidence": result.evidence, + "metadata": result.metadata, + } + for result in evaluation.criterion_results + ], + "metadata": evaluation.metadata, + } + + +def _message_to_request(message: Any) -> dict[str, Any]: + data = message.to_dict() + request = { + "role": data["role"], + "parts": data.get("parts", []), + "created_at": data.get("created_at"), + } + if data.get("peer_id") is not None: + request["peer_id"] = data["peer_id"] + return request diff --git a/openviking/session/train/components/skill_policy_updater.py b/openviking/session/train/components/skill_policy_updater.py new file mode 100644 index 0000000000..2bea523905 --- /dev/null +++ b/openviking/session/train/components/skill_policy_updater.py @@ -0,0 +1,317 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""PolicyUpdater that writes skill files via SkillProcessor / SkillOperationUpdater.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from openviking.core.skill_loader import SkillLoader +from openviking.server.identity import RequestContext +from openviking.session.memory.dataclass import ( + MemoryFile, + ResolvedOperation, + ResolvedOperations, +) +from openviking.session.skill import SkillOperationUpdater +from openviking.session.skill.session_skill_context_provider import ( + SESSION_SKILL_MEMORY_TYPE, + load_skill_extract_registry, +) +from openviking.session.train.domain import ( + Policy, + PolicyApplyResult, + PolicyPlanItem, + PolicySet, + PolicyUpdatePlan, +) +from openviking.storage.viking_fs import get_viking_fs +from openviking.telemetry import tracer +from openviking.utils.skill_processor import SkillProcessor +from openviking_cli.utils import get_logger + +logger = get_logger(__name__) + + +@dataclass(slots=True) +class SkillPolicyUpdater: + """PolicyUpdater that writes skill files to a skills directory. + + For new skills (no existing file) the full ``SkillProcessor.process_skill`` + pipeline is used (validation, privacy, overview, index). For existing + skills, the merged content is serialized to SKILL.md and written via + ``ContentWriteCoordinator``. + + ``delete`` operations remove the entire skill subdirectory. + """ + + skill_processor: SkillProcessor | None = None + viking_fs: Any = None + vikingdb: Any = None + memory_type: str = "skills" + + @tracer("train.policy_updater.skill.apply", ignore_result=True, ignore_args=True) + async def apply( + self, + plan: PolicyUpdatePlan, + policy_set: PolicySet, + context: Any = None, + *, + transaction_handle: Any = None, + ) -> PolicyApplyResult: + ctx = _coerce_request_context(context) + if ctx is None: + raise ValueError("SkillPolicyUpdater.apply requires a request context") + viking_fs = self.viking_fs or get_viking_fs() + if viking_fs is None: + raise RuntimeError("VikingFS is required to apply skill policy updates") + + updated_policy_set = _apply_items_to_snapshot(plan.items, policy_set) + operations = _plan_to_resolved_operations( + plan=plan, + policy_set=policy_set, + updated_policy_set=updated_policy_set, + ) + if not operations.upsert_operations and not operations.delete_file_contents: + return PolicyApplyResult( + updated_policy_set=policy_set, + written_uris=[], + deleted_uris=[], + errors=[], + metadata={"dry_run": False, "item_count": 0, "memory_type": self.memory_type}, + ) + + registry = load_skill_extract_registry() + processor = self.skill_processor or SkillProcessor() + updater = SkillOperationUpdater( + registry=registry, + skill_processor=processor, + viking_fs=viking_fs, + ) + result = await updater.apply_operations(operations, ctx) + + errors = [f"{uri}: {exc}" for uri, exc in result.errors] + + # Handle deletes (SkillOperationUpdater doesn't support delete ops) + delete_errors: list[str] = [] + deleted_uris: list[str] = [] + for old_file in operations.delete_file_contents: + if not old_file.uri: + continue + try: + skill_root = _root_uri_from_skill_md(old_file.uri) + await viking_fs.rm(skill_root, ctx=ctx, lock_handle=transaction_handle) + deleted_uris.append(old_file.uri) + except Exception as exc: + delete_errors.append(f"{old_file.uri}: {exc}") + + all_errors = [*errors, *delete_errors] + return PolicyApplyResult( + updated_policy_set=updated_policy_set if not all_errors else policy_set, + written_uris=list(result.written_uris + result.edited_uris), + deleted_uris=deleted_uris, + errors=all_errors, + metadata={ + "dry_run": False, + "item_count": len(plan.items), + "memory_type": self.memory_type, + "operation_upsert_count": len(operations.upsert_operations), + "operation_delete_count": len(operations.delete_file_contents), + }, + ) + + +def _coerce_request_context(context: Any) -> RequestContext | None: + if context is None: + return None + if isinstance(context, dict): + return context.get("request_context") or context.get("ctx") + # Try duck-typing for common context wrappers + for attr in ("request_context", "ctx", "apply_context"): + value = getattr(context, attr, None) + if value is not None: + return value + # If it quacks like a RequestContext… + if hasattr(context, "user_id") or hasattr(context, "account_id"): + return context # type: ignore[return-value] + return None + + +def _apply_items_to_snapshot( + items: list[PolicyPlanItem], policy_set: PolicySet +) -> PolicySet: + policies_by_uri = {policy.uri: policy for policy in policy_set.policies} + result = list(policy_set.policies) + + for item in items: + uri = _target_uri(item, policy_set.root_uri) + + if item.kind == "delete": + existing = policies_by_uri.get(uri) or _find_policy( + policy_set, uri=None, name=item.target_name + ) + remove_uri = existing.uri if existing is not None else uri + result = [ + policy + for policy in result + if policy.uri != remove_uri and policy.name != item.target_name + ] + policies_by_uri.pop(remove_uri, None) + policies_by_uri.pop(uri, None) + continue + + if item.kind != "upsert" or item.after_content is None: + continue + existing = policies_by_uri.get(uri) or _find_policy( + policy_set, uri=None, name=item.target_name + ) + metadata = dict(existing.metadata) if existing is not None else {} + metadata.update(item.metadata.get("patch_metadata", {})) + metadata.setdefault("memory_type", item.memory_type or "skills") + version = (existing.version + 1) if existing is not None else 1 + updated = Policy( + name=item.target_name, + uri=uri, + version=version, + status=(existing.status if existing is not None else "draft"), + content=item.after_content, + metadata=metadata, + links=list(existing.links or []) if existing is not None else [], + backlinks=list(existing.backlinks or []) if existing is not None else [], + ) + if existing is None: + result.append(updated) + else: + result = [updated if policy.uri == existing.uri else policy for policy in result] + policies_by_uri[uri] = updated + + result.sort(key=lambda policy: policy.uri) + return PolicySet( + root_uri=policy_set.root_uri, + policies=result, + metadata=dict(policy_set.metadata), + viking_fs=policy_set.viking_fs, + request_context=policy_set.request_context, + ) + + +def _find_policy(policy_set: PolicySet, *, uri: str | None, name: str) -> Policy | None: + for policy in policy_set.policies: + if uri and policy.uri == uri: + return policy + if not uri and policy.name == name: + return policy + return None + + +def _target_uri(item: PolicyPlanItem, root_uri: str) -> str: + if item.target_uri: + return item.target_uri + skill_name = _safe_skill_dirname(item.target_name) + return f"{root_uri.rstrip('/')}/{skill_name}/SKILL.md" + + +def _plan_to_resolved_operations( + *, + plan: PolicyUpdatePlan, + policy_set: PolicySet, + updated_policy_set: PolicySet, +) -> ResolvedOperations: + upserts: list[ResolvedOperation] = [] + deletes: list[MemoryFile] = [] + errors: list[str] = [] + + for item in plan.items: + uri = _target_uri(item, policy_set.root_uri) + current = _find_policy(policy_set, uri=uri, name=item.target_name) + + if item.kind == "delete": + if current is not None: + deletes.append(_policy_to_memory_file(current)) + continue + + if item.kind != "upsert": + continue + if item.after_content is None: + errors.append(f"missing after_content for {item.target_name}") + continue + + updated = _find_policy(updated_policy_set, uri=uri, name=item.target_name) + if updated is None: + errors.append( + f"planned skill policy not found after simulation: {item.target_name}" + ) + continue + + old_mf = _policy_to_memory_file(current) if current is not None else None + # Build memory_fields in the shape expected by SkillOperationUpdater. + # Skill schema uses "skill_name" / "description" / "content" fields. + memory_fields: dict[str, Any] = { + "skill_name": updated.name, + "content": updated.content, + "description": updated.metadata.get("description", ""), + } + # Carry over other metadata + for key in ("allowed_tools", "tags"): + value = updated.metadata.get(key) + if value is not None: + memory_fields[key] = value + + upserts.append( + ResolvedOperation( + old_memory_file_content=old_mf, + memory_fields=memory_fields, + memory_type=SESSION_SKILL_MEMORY_TYPE, + uris=[uri], + ) + ) + + return ResolvedOperations( + upsert_operations=upserts, + delete_file_contents=deletes, + errors=errors, + resolved_links=[], + ) + + +def _policy_to_memory_file(policy: Policy | None) -> MemoryFile | None: + if policy is None: + return None + # Serialize policy content + metadata into a SKILL.md-shaped MemoryFile. + skill_dict = { + "name": policy.name, + "description": policy.metadata.get("description", ""), + "content": policy.content, + "allowed_tools": policy.metadata.get("allowed_tools", []), + "tags": policy.metadata.get("tags", []), + } + serialized = SkillLoader.to_skill_md(skill_dict) + return MemoryFile( + uri=policy.uri, + content=serialized, + links=list(policy.links or []), + backlinks=list(policy.backlinks or []), + memory_type="skills", + extra_fields={ + **dict(policy.metadata), + "memory_type": "skills", + "skill_name": policy.name, + "version": policy.version, + "status": policy.status, + }, + ) + + +def _safe_skill_dirname(name: str) -> str: + import re + + cleaned = re.sub(r"[^a-zA-Z0-9_\-一-鿿]+", "_", name.strip()).strip("._-") + return cleaned or "new_skill" + + +def _root_uri_from_skill_md(skill_md_uri: str) -> str: + suffix = "/SKILL.md" + if skill_md_uri.endswith(suffix): + return skill_md_uri[: -len(suffix)] + return skill_md_uri.rstrip("/") diff --git a/openviking/session/train/components/snapshotter.py b/openviking/session/train/components/snapshotter.py new file mode 100644 index 0000000000..27d0330916 --- /dev/null +++ b/openviking/session/train/components/snapshotter.py @@ -0,0 +1,45 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Policy snapshot helpers.""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass +from typing import Any + +from openviking.session.train.domain import ExperienceSet +from openviking.telemetry import tracer + + +@dataclass(slots=True) +class ContentHashPolicySnapshotter: + """Create deterministic policy snapshot ids from ExperienceSet content.""" + + prefix: str = "policy-snapshot" + + @tracer( + "train.policy_snapshotter.content_hash.snapshot", + ignore_result=False, + ignore_args=True, + ) + async def snapshot(self, policy_set: ExperienceSet, context: Any = None) -> str: + del context + payload = { + "root_uri": policy_set.root_uri, + "policies": [ + { + "name": policy.name, + "uri": policy.uri, + "version": policy.version, + "status": policy.status, + "content": policy.content, + "metadata": policy.metadata, + } + for policy in sorted(policy_set.policies, key=lambda p: p.uri) + ], + } + raw = json.dumps(payload, ensure_ascii=False, sort_keys=True, separators=(",", ":")) + digest = hashlib.sha256(raw.encode("utf-8")).hexdigest()[:16] + return f"{self.prefix}:{digest}" diff --git a/openviking/session/train/components/trajectory_analyzer.py b/openviking/session/train/components/trajectory_analyzer.py new file mode 100644 index 0000000000..dea375725e --- /dev/null +++ b/openviking/session/train/components/trajectory_analyzer.py @@ -0,0 +1,548 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""RolloutAnalyzer that extracts persistent trajectory memories directly.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from openviking.core.context import Context +from openviking.message import Message, TextPart +from openviking.server.identity import RequestContext +from openviking.session.memory import ExtractLoop, MemoryUpdater +from openviking.session.memory.agent_trajectory_context_provider import ( + AgentTrajectoryContextProvider, +) +from openviking.session.memory.dataclass import ( + MemoryFile, + ResolvedOperations, + StoredLink, +) +from openviking.session.memory.memory_isolation_handler import MemoryIsolationHandler +from openviking.session.memory.memory_updater import ExtractContext, MemoryUpdateResult +from openviking.session.memory.utils.memory_file_utils import MemoryFileUtils +from openviking.session.skill.session_skill_context_provider import ( + SESSION_SKILL_MEMORY_TYPE, +) +from openviking.session.train.domain import ( + CriterionResult, + Rollout, + RolloutAnalysis, + RubricEvaluation, + Trajectory, +) +from openviking.session.train.gradients import PatchSemanticGradient +from openviking.session.train.interfaces import RolloutEvaluator +from openviking.storage.viking_fs import get_viking_fs +from openviking.telemetry import tracer +from openviking_cli.utils import get_logger +from openviking_cli.utils.config import get_openviking_config + +logger = get_logger(__name__) + +_TRAJECTORY_MEMORY_TYPE = "trajectories" + + +@dataclass(slots=True) +class TrajectoryAnalyzerContext: + """Runtime context for TrajectoryRolloutAnalyzer.""" + + request_context: RequestContext + strict_extract_errors: bool = False + latest_archive_overview: str = "" + evaluator_context: Any = None + inject_evaluation_feedback: bool = True + include_session_skills: bool = False + + +@dataclass(slots=True) +class TrajectoryRolloutAnalyzer: + """Analyze rollouts by extracting persistent trajectory memory files. + + This implementation owns the trajectory extraction/apply flow directly. It + intentionally does not depend on SessionCompressorV2/V3, and it only exposes + the trajectory memory schema to ExtractLoop. + """ + + viking_fs: Any = None + vikingdb: Any = None + vlm: Any = None + evaluator: RolloutEvaluator | None = None + + @tracer("train.rollout_analyzer.trajectory.analyze", ignore_result=True, ignore_args=True) + async def analyze( + self, + rollout: Rollout, + context: TrajectoryAnalyzerContext, + ) -> RolloutAnalysis: + if context is None or context.request_context is None: + raise ValueError("TrajectoryAnalyzerContext.request_context is required") + + evaluation = await self._evaluate_rollout(rollout, context) + extraction_messages = _messages_with_evaluation_feedback( + rollout.messages, + evaluation=evaluation, + enabled=evaluation is not None and context.inject_evaluation_feedback, + ) + result = await self.extract_trajectory_memories( + messages=extraction_messages, + ctx=context.request_context, + strict_extract_errors=context.strict_extract_errors, + latest_archive_overview=context.latest_archive_overview, + include_session_skills=context.include_session_skills, + case_name=getattr(rollout.case, "name", ""), + ) + contexts = list((result or {}).get("contexts", [])) + skill_gradients = list((result or {}).get("skill_gradients", [])) + trajectory_uris = [ + item.uri + for item in contexts + if getattr(item, "category", "") == "memory_write" + and "/memories/trajectories/" in getattr(item, "uri", "") + ] + trajectory_uris = list(dict.fromkeys(trajectory_uris)) + trajectories = await self._read_trajectories( + trajectory_uris, + ctx=context.request_context, + ) + evaluation = evaluation or _evaluation_from_trajectories(trajectories) + return RolloutAnalysis( + evaluation=evaluation, + trajectories=trajectories, + gradients=skill_gradients, + metadata={ + "context_count": len(contexts), + "policy_snapshot_id": rollout.policy_snapshot_id, + "rollout_messages": rollout.messages, + "extraction_message_count": len(extraction_messages), + }, + ) + + async def _evaluate_rollout( + self, + rollout: Rollout, + context: TrajectoryAnalyzerContext, + ) -> RubricEvaluation | None: + if rollout.evaluation is not None: + return rollout.evaluation + if self.evaluator is None: + return None + return await self.evaluator.evaluate(rollout, context.evaluator_context) + + async def extract_trajectory_memories( + self, + *, + messages: list[Message], + ctx: RequestContext | None, + strict_extract_errors: bool = False, + latest_archive_overview: str = "", + include_session_skills: bool = False, + case_name: str = "", + ) -> dict[str, list[Any]]: + """Extract and persist trajectory memories from rollout messages. + + When ``include_session_skills`` is True, session skill patches are + co-extracted in the same ExtractLoop pass and returned as + ``PatchSemanticGradient`` instances in the ``"skill_gradients"`` key. + Skill patches are *not* applied to disk by this method — they are + returned as gradient signals for downstream policy training. + """ + empty_result: dict[str, list[Any]] = {"contexts": [], "skill_gradients": []} + if not messages or ctx is None: + return empty_result + + provider = AgentTrajectoryContextProvider( + messages=messages, + latest_archive_overview=latest_archive_overview, + include_trajectories=True, + include_session_skills=include_session_skills, + ) + phase_result = await self._run_trajectory_extract_phase( + provider=provider, + messages=messages, + ctx=ctx, + strict_extract_errors=strict_extract_errors, + include_session_skills=include_session_skills, + case_name=case_name, + ) + if phase_result is None: + return empty_result + + _, _, contexts, skill_gradients = phase_result + return {"contexts": contexts, "skill_gradients": skill_gradients} + + async def _run_trajectory_extract_phase( + self, + *, + provider: AgentTrajectoryContextProvider, + messages: list[Message], + ctx: RequestContext, + strict_extract_errors: bool, + include_session_skills: bool = False, + case_name: str = "", + ) -> tuple[list[str], list[str], list[Context], list[PatchSemanticGradient]] | None: + config = get_openviking_config() + vlm = self.vlm or config.vlm.get_vlm_instance() + viking_fs = self.viking_fs or get_viking_fs() + if viking_fs is None: + raise RuntimeError("VikingFS is required to extract trajectory memories") + + extract_context = provider.get_extract_context() + allowed_types: set[str] = {_TRAJECTORY_MEMORY_TYPE} + if include_session_skills: + allowed_types.add(SESSION_SKILL_MEMORY_TYPE) + isolation_handler = MemoryIsolationHandler( + ctx, + extract_context, + allowed_memory_types=allowed_types, + ) + isolation_handler.prepare_messages() + + provider._isolation_handler = isolation_handler + provider._ctx = ctx + provider._viking_fs = viking_fs + + orchestrator = ExtractLoop( + vlm=vlm, + viking_fs=viking_fs, + ctx=ctx, + context_provider=provider, + isolation_handler=isolation_handler, + thinking=True, + ) + + try: + provider._transaction_handle = None + orchestrator._transaction_handle = None + operations, _ = await orchestrator.run() + if operations is None: + tracer.info("[trajectory] No memory operations generated") + return [], [], [], [] + + _log_operations(operations) + + # Split operations into trajectory (applied to disk) and skill + # (returned as gradients). Skill ops are *not* written here — + # they flow through the patch-merge trainer. + traj_ops, skill_ops = _split_operations_by_type( + operations, target_type=_TRAJECTORY_MEMORY_TYPE + ) + skill_gradients = _skill_operations_to_gradients( + skill_ops, + viking_fs=viking_fs, + ctx=ctx, + ) + + _ensure_trajectory_case_name(traj_ops, case_name=case_name) + + memory_result = await self._apply_trajectory_operations( + operations=traj_ops, + provider=provider, + ctx=ctx, + extract_context=extract_context, + isolation_handler=isolation_handler, + ) + tracer.info( + "[trajectory] Applied memory ops: " + f"written={len(memory_result.written_uris)}, " + f"edited={len(memory_result.edited_uris)}, " + f"deleted={len(memory_result.deleted_uris)}, " + f"errors={len(memory_result.errors)}" + ) + contexts = _contexts_from_memory_result(memory_result) + return ( + list(memory_result.written_uris), + list(memory_result.edited_uris), + contexts, + skill_gradients, + ) + except Exception as exc: + logger.error("[trajectory] Failed to extract: %s", exc, exc_info=True) + if strict_extract_errors: + raise + return None + + async def _apply_trajectory_operations( + self, + *, + operations: ResolvedOperations, + provider: AgentTrajectoryContextProvider, + ctx: RequestContext, + extract_context: ExtractContext, + isolation_handler: MemoryIsolationHandler, + ) -> MemoryUpdateResult: + updater = MemoryUpdater( + registry=provider._get_registry(), + vikingdb=self.vikingdb, + transaction_handle=None, + ) + updater._viking_fs = self.viking_fs or get_viking_fs() + return await updater.apply_operations( + operations, + ctx, + extract_context=extract_context, + isolation_handler=isolation_handler, + ) + + @tracer("train.rollout_analyzer.trajectory.read_trajectories", ignore_result=True, ignore_args=True) + async def _read_trajectories( + self, + trajectory_uris: list[str], + *, + ctx: RequestContext, + ) -> list[Trajectory]: + viking_fs = self.viking_fs or get_viking_fs() + if viking_fs is None: + raise RuntimeError("VikingFS is required to read extracted trajectories") + + trajectories: list[Trajectory] = [] + for uri in dict.fromkeys(trajectory_uris): + try: + raw = await viking_fs.read_file(uri, ctx=ctx) or "" + mf = MemoryFileUtils.read(raw, uri=uri) + except Exception as exc: + logger.warning("Failed to read trajectory %s: %s", uri, exc) + continue + fields = dict(mf.extra_fields or {}) + name = str( + fields.get("trajectory_name") or uri.rstrip("/").split("/")[-1].removesuffix(".md") + ) + outcome = str(fields.get("outcome") or "unknown") + retrieval_anchor = str(fields.get("retrieval_anchor") or "") + case_name = str(fields.get("case_name") or "") + metadata = dict(fields) + metadata.setdefault("memory_type", mf.memory_type or fields.get("memory_type")) + metadata.setdefault("case_name", case_name) + trajectories.append( + Trajectory( + name=name, + uri=uri, + content=mf.plain_content(), + outcome=outcome, + retrieval_anchor=retrieval_anchor, + metadata=metadata, + ) + ) + return trajectories + + + +def _log_operations(operations: ResolvedOperations) -> None: + op_items = [ + f"{op.memory_type}(uris={op.uris!r})" + for op in getattr(operations, "upsert_operations", []) + ] + delete_uris = [dc.uri for dc in getattr(operations, "delete_file_contents", [])] + tracer.info(f"[trajectory] LLM operations: ops={op_items}, delete_uris={delete_uris}") + + +def _contexts_from_memory_result(memory_result: MemoryUpdateResult) -> list[Context]: + contexts: list[Context] = [] + for uri in memory_result.written_uris: + contexts.append(Context(uri=uri, category="memory_write", context_type="memory")) + for uri in memory_result.edited_uris: + contexts.append(Context(uri=uri, category="memory_edit", context_type="memory")) + for uri in memory_result.deleted_uris: + contexts.append(Context(uri=uri, category="memory_delete", context_type="memory")) + return contexts + + +def _evaluation_from_trajectories(trajectories: list[Trajectory]) -> RubricEvaluation: + passed = bool(trajectories) + return RubricEvaluation( + passed=passed, + score=1.0 if passed else 0.0, + criterion_results=[ + CriterionResult( + criterion_name="trajectory_extracted", + passed=passed, + score=1.0 if passed else 0.0, + feedback=[] if passed else ["No trajectory was extracted from the rollout."], + evidence=[trajectory.uri for trajectory in trajectories], + ) + ], + feedback=[] if passed else ["No trajectory was extracted from the rollout."], + metadata={"trajectory_count": len(trajectories)}, + ) + + +def _ensure_trajectory_case_name(operations: ResolvedOperations, *, case_name: str) -> None: + case_name = str(case_name or "").strip() + if not case_name: + return + for op in getattr(operations, "upsert_operations", []) or []: + if getattr(op, "memory_type", None) != _TRAJECTORY_MEMORY_TYPE: + continue + fields = getattr(op, "memory_fields", None) + if isinstance(fields, dict): + fields["case_name"] = case_name + + +def _messages_with_evaluation_feedback( + messages: list[Message], + *, + evaluation: RubricEvaluation | None, + enabled: bool, +) -> list[Message]: + result = list(messages) + if not enabled or evaluation is None: + return result + result.append(_evaluation_feedback_message(evaluation)) + return result + + +def _evaluation_feedback_message(evaluation: RubricEvaluation) -> Message: + lines = [ + "[Rollout Evaluation]", + f"passed: {evaluation.passed}", + f"score: {evaluation.score}", + ] + if evaluation.feedback: + lines.extend(["", "feedback:", *[f"- {item}" for item in evaluation.feedback]]) + criterion_lines: list[str] = [] + evidence_lines: list[str] = [] + for criterion in evaluation.criterion_results: + criterion_lines.append( + f"- {criterion.criterion_name}: " + f"passed={criterion.passed}, score={criterion.score}" + ) + criterion_lines.extend(f" feedback: {item}" for item in criterion.feedback) + evidence_lines.extend(criterion.evidence) + if criterion_lines: + lines.extend(["", "criteria:", *criterion_lines]) + if evidence_lines: + lines.extend(["", "evidence:", *[f"- {item}" for item in dict.fromkeys(evidence_lines)]]) + return Message( + id="rollout-evaluation-feedback", + role="user", + parts=[TextPart(text="\n".join(lines))], + ) + + +def _split_operations_by_type( + operations: ResolvedOperations, *, target_type: str +) -> tuple[ResolvedOperations, ResolvedOperations]: + """Split operations into (target_type_ops, other_ops).""" + target_upserts = [ + op for op in operations.upsert_operations if op.memory_type == target_type + ] + other_upserts = [ + op for op in operations.upsert_operations if op.memory_type != target_type + ] + target_deletes = [ + dc for dc in operations.delete_file_contents if dc.memory_type == target_type + ] + other_deletes = [ + dc for dc in operations.delete_file_contents if dc.memory_type != target_type + ] + target_ops = ResolvedOperations( + upsert_operations=target_upserts, + delete_file_contents=target_deletes, + errors=list(operations.errors), + resolved_links=[ + link for link in operations.resolved_links + if getattr(link, "from_uri", "").endswith("/trajectories/") + or target_type in getattr(link, "from_uri", "") + ], + ) + other_ops = ResolvedOperations( + upsert_operations=other_upserts, + delete_file_contents=other_deletes, + errors=[], + resolved_links=[], + ) + return target_ops, other_ops + + +def _skill_operations_to_gradients( + operations: ResolvedOperations, + *, + viking_fs: Any = None, + ctx: Any = None, +) -> list[PatchSemanticGradient]: + """Convert skill ResolvedOperations to PatchSemanticGradient instances. + + The resulting gradients carry the full proposed skill content in their + ``after_file`` so the patch-merge optimizer can reconcile multiple + proposals against the current policy set. + """ + gradients: list[PatchSemanticGradient] = [] + for op in operations.upsert_operations or []: + if op.memory_type != SESSION_SKILL_MEMORY_TYPE: + continue + fields = dict(op.memory_fields or {}) + skill_name = str(fields.get("skill_name") or _fallback_skill_name(op)) + target_uri = (op.uris or [None])[0] + after_content = str(fields.get("content") or "") + if not after_content.strip(): + continue + + old_file = op.old_memory_file_content + after_file = MemoryFile( + uri=target_uri, + content=after_content, + memory_type="skills", + extra_fields={ + **dict(getattr(old_file, "extra_fields", {}) or {}), + **{k: v for k, v in fields.items() if k != "content"}, + "memory_type": "skills", + "skill_name": skill_name, + }, + ) + links: list[StoredLink] = [] + # Build derived_from links from the source trajectory(s). + for link in getattr(operations, "resolved_links", []) or []: + try: + stored = link if isinstance(link, StoredLink) else StoredLink(**dict(link)) + if stored.link_type == "derived_from" and stored.to_uri: + if "/memories/trajectories/" in stored.to_uri: + links.append( + stored.model_copy(update={"from_uri": target_uri or ""}) + ) + except Exception: + continue + + gradients.append( + PatchSemanticGradient( + before_file=old_file, + after_file=after_file, + base_version=_base_version_from_old_file(old_file), + rationale=( + "Session skill patch extracted from rollout trajectory " + "by AgentTrajectoryContextProvider." + ), + links=links, + confidence=0.7, + metadata={ + "source": "trajectory_co_extract", + "memory_fields": fields, + "skill_name": skill_name, + "uris": list(op.uris or []), + }, + ) + ) + return gradients + + +def _fallback_skill_name(op: Any) -> str: + uris = getattr(op, "uris", None) or [] + if uris: + uri = str(uris[0]) + # path/to/skills/my_skill/SKILL.md → my_skill + parts = uri.rstrip("/").split("/") + if len(parts) >= 2 and parts[-1] == "SKILL.md": + return parts[-2] + return parts[-1].removesuffix(".md") + return "unknown_skill" + + +def _base_version_from_old_file(old_file: Any) -> int | None: + if old_file is None: + return None + fields = getattr(old_file, "extra_fields", {}) or {} + try: + v = int(fields.get("version")) + return v if v > 0 else None + except (TypeError, ValueError): + return None diff --git a/openviking/session/train/context.py b/openviking/session/train/context.py new file mode 100644 index 0000000000..7e597cacff --- /dev/null +++ b/openviking/session/train/context.py @@ -0,0 +1,60 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Context models for session policy training pipelines.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +from openviking.session.train.components.report_builder import PipelineReportHook +from openviking.session.train.components.reporter import ConsolePipelineReporter + +if TYPE_CHECKING: + from openviking.session.train.components.report_builder import PipelineReportBuilder + from openviking.session.train.components.reporter import PipelineLifecycleHook + from openviking.session.train.interfaces import CaseLoader + + +@dataclass(slots=True) +class PipelineHookDecision: + """Control decision returned by lifecycle hooks.""" + + stop_training: bool = False + reason: str = "" + metadata: dict[str, Any] = field(default_factory=dict) + report: dict[str, Any] | None = None + + +@dataclass(slots=True) +class PipelineContext: + """Context bundle for OfflinePolicyOptimizationPipeline. + + Context payloads are intentionally opaque and can be shaped by concrete + implementations without changing the domain interfaces. + """ + + case_load_context: Any = None + snapshot_context: Any = None + analysis_context: Any = None + gradient_context: Any = None + optimization_context: Any = None + apply_context: Any = None + execution_metadata: dict[str, Any] = field(default_factory=dict) + max_epochs: int = 1 + eval_each_epoch_case_loader: CaseLoader | None = None + eval_trials: int = 1 + train_trials: int = 1 + trial_index_key: str = "trial" + report_builder: PipelineReportBuilder | None = None + lifecycle_hooks: list[PipelineLifecycleHook] = field( + default_factory=lambda: [PipelineReportHook(), ConsolePipelineReporter()] + ) + + +@dataclass(slots=True) +class ExecutionContext: + """Runtime context passed to RolloutExecutor.""" + + policy_snapshot_id: str + metadata: dict[str, Any] = field(default_factory=dict) diff --git a/openviking/session/train/domain.py b/openviking/session/train/domain.py new file mode 100644 index 0000000000..b7348952c8 --- /dev/null +++ b/openviking/session/train/domain.py @@ -0,0 +1,332 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Domain models for session experience-policy optimization. + +This module defines the new training domain model alongside the existing +trajectory/experience memory implementation. The types here are intentionally +small and implementation-agnostic so the new framework can be built out without +changing the current extraction pipeline. +""" + +from __future__ import annotations + +from contextlib import asynccontextmanager +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Literal + +from openviking.message import Message +from openviking.session.memory.dataclass import StoredLink + +if TYPE_CHECKING: + from openviking.session.train.gradients import PatchSemanticGradient + +PolicyStatus = Literal["draft", "staging", "production", "deprecated", "archived"] +TrajectoryOutcome = Literal["success", "failure", "partial", "unfinished", "unknown"] +PolicyPlanItemKind = Literal["upsert", "delete"] + + +@dataclass(slots=True) +class Policy: + """A single policy file in a PolicySet. + + Generic policy item used for experiences, skills, and other trainable + memory types. Type-specific fields live in ``metadata``. + """ + + name: str + uri: str + version: int + status: PolicyStatus + content: str + metadata: dict[str, Any] = field(default_factory=dict) + links: list[dict[str, Any]] = field(default_factory=list) + backlinks: list[dict[str, Any]] = field(default_factory=list) + + +# Backwards-compatible alias +Experience = Policy + + +@dataclass(slots=True) +class PolicySet: + """Snapshot of all policies under a policy root directory. + + ``viking_fs`` and ``request_context`` are runtime storage dependencies used + for concurrency-safe policy updates. They are intentionally excluded from + equality/repr so the domain snapshot still behaves like policy data in + tests and diagnostics. + """ + + root_uri: str + policies: list[Policy] + metadata: dict[str, Any] = field(default_factory=dict) + viking_fs: Any | None = field(default=None, repr=False, compare=False) + request_context: Any | None = field(default=None, repr=False, compare=False) + + @asynccontextmanager + async def lock(self): + """Acquire a tree lock for the whole policy root directory. + + Policy updates serialize on this lock so concurrent realtime/batch + training jobs plan and apply against a freshly reloaded policy set. + ``timeout=None`` means wait indefinitely until the lock is available. + """ + + if self.viking_fs is None: + raise RuntimeError("PolicySet.viking_fs is required for policy locking") + if self.request_context is None: + raise RuntimeError("PolicySet.request_context is required for policy locking") + uri_to_path = getattr(self.viking_fs, "_uri_to_path", None) + if uri_to_path is None: + raise RuntimeError("PolicySet.viking_fs must provide _uri_to_path for locking") + + from openviking.storage.transaction import get_lock_manager + + lock_manager = get_lock_manager() + handle = lock_manager.create_handle() + path = uri_to_path(self.root_uri, ctx=self.request_context) + acquired = await lock_manager.acquire_tree(handle, path, timeout=None) + if not acquired: + await lock_manager.release(handle) + raise RuntimeError(f"Failed to acquire policy tree lock for {self.root_uri}") + try: + yield handle + finally: + await lock_manager.release(handle) + + async def reload(self) -> "PolicySet": + """Reload this policy set from its backing VikingFS under the same ctx.""" + + if self.viking_fs is None: + raise RuntimeError("PolicySet.viking_fs is required for policy reload") + if self.request_context is None: + raise RuntimeError("PolicySet.request_context is required for policy reload") + + from openviking.session.train.components.memory_store import ExperienceSetLoader + + return await ExperienceSetLoader(viking_fs=self.viking_fs).load( + self.root_uri, + ctx=self.request_context, + ) + + +# Backwards-compatible alias +ExperienceSet = PolicySet + + +@dataclass(slots=True) +class Trajectory: + """A distilled, trainable trajectory sample parsed from trajectory memory.""" + + name: str + uri: str + content: str + outcome: TrajectoryOutcome | str + retrieval_anchor: str + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class RubricCriterion: + """One criterion in a case rubric.""" + + name: str + description: str + required: bool + weight: float + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class Rubric: + """Acceptance criteria and scoring rules for a case.""" + + name: str + description: str + criteria: list[RubricCriterion] + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class Case: + """An executable, reproducible, evaluable training/evaluation sample.""" + + name: str + task_signature: str + input: dict[str, Any] + rubric: Rubric + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class Rollout: + """Execution record for a case under a policy-set snapshot.""" + + case: Case + messages: list[Message] + policy_snapshot_id: str + evaluation: "RubricEvaluation | None" = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class CriterionResult: + """Evaluation result for one rubric criterion.""" + + criterion_name: str + passed: bool + score: float + feedback: list[str] + evidence: list[str] + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class RubricEvaluation: + """Structured evaluation of a rollout against a rubric.""" + + passed: bool + score: float + criterion_results: list[CriterionResult] + feedback: list[str] + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class RolloutAnalysis: + """Structured analysis of a rollout. + + Contains both rubric evaluation and trajectories extracted from the same + rollout context. ``gradients`` carries any policy patches co-extracted + during analysis (e.g. session skill patches) keyed by their + ``memory_type``; these bypass the gradient estimator and are fed directly + into the corresponding policy trainer. + """ + + evaluation: RubricEvaluation + trajectories: list[Trajectory] + gradients: list["PatchSemanticGradient"] = field(default_factory=list) + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class PolicyPlanItem: + """One executable item in a PolicyUpdatePlan. + + Supports multiple policy memory types (experiences, skills, ...) via the + ``memory_type`` field. Each item represents an upsert or delete operation + against a single target policy file. + """ + + kind: PolicyPlanItemKind + memory_type: str + target_name: str + target_uri: str | None + before_content: str | None + after_content: str | None + base_version: int | None = None + confidence: float | None = None + links: list[StoredLink] = field(default_factory=list) + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class PolicyUpdatePlan: + """Planned update for a PolicySet. + + ``items`` is the executable part consumed by PolicyUpdater implementations. + ``metadata`` keeps optimizer diagnostics such as grouping, conflicts, and + unresolved gradients. + """ + + items: list[PolicyPlanItem] = field(default_factory=list) + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class PolicyApplyResult: + """Result of applying a PolicyUpdatePlan.""" + + updated_policy_set: PolicySet + written_uris: list[str] = field(default_factory=list) + deleted_uris: list[str] = field(default_factory=list) + errors: list[str] = field(default_factory=list) + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class PipelineEpochResult: + """Result of one rollout/evaluate/train epoch. + + One epoch runs the current policy snapshot on case batches, analyzes the + resulting rollouts, estimates semantic gradients, plans a policy update, and + applies it. Repeating this structure models the offline equivalent of + rollout -> evaluation -> update -> rollout -> evaluation. + """ + + epoch: int + analyses: list[RolloutAnalysis] + gradients: list[Any] + plan: PolicyUpdatePlan + apply_result: PolicyApplyResult + policy_snapshot_ids: list[str] = field(default_factory=list) + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class PipelineEvaluationResult: + """Evaluation-only rollout result for a policy snapshot. + + This is typically used as the final after-training evaluation pass. It + intentionally does not include gradients or policy updates. + """ + + epoch: int + analyses: list[RolloutAnalysis] + policy_snapshot_ids: list[str] = field(default_factory=list) + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class PipelineResult: + """End-to-end result of a policy optimization pipeline train call.""" + + analyses: list[RolloutAnalysis] + gradients: list[Any] + plan: PolicyUpdatePlan + apply_result: PolicyApplyResult + epochs: list[PipelineEpochResult] = field(default_factory=list) + evaluation_passes: list[PipelineEvaluationResult] = field(default_factory=list) + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class RolloutTrainingResult: + """Result of training directly from externally produced rollouts. + + This is the online/realtime counterpart of one offline pipeline training + epoch. The caller owns rollout execution; the training framework owns + analysis, gradient estimation, policy planning, and policy update. + """ + + analyses: list[RolloutAnalysis] + gradients: list[Any] + plan: PolicyUpdatePlan + apply_result: PolicyApplyResult + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class ScopedRolloutTrainingResult: + """Training result for one submitter inside a shared streaming batch. + + ``batch_result`` keeps the full flush outcome for diagnostics, while the + top-level fields are scoped to the submitter's rollout/gradient provenance. + """ + + analyses: list[RolloutAnalysis] + gradients: list[Any] + plan: PolicyUpdatePlan + apply_result: PolicyApplyResult + batch_result: RolloutTrainingResult + metadata: dict[str, Any] = field(default_factory=dict) diff --git a/openviking/session/train/engine.py b/openviking/session/train/engine.py new file mode 100644 index 0000000000..21d8c24d43 --- /dev/null +++ b/openviking/session/train/engine.py @@ -0,0 +1,100 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Shared training engine for rollout-driven policy updates.""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass + +from openviking.session.train.context import PipelineContext +from openviking.session.train.domain import ( + ExperienceSet, + PolicyApplyResult, + PolicyUpdatePlan, + Rollout, + RolloutAnalysis, +) +from openviking.session.train.interfaces import ( + GradientEstimator, + PolicyOptimizer, + PolicyUpdater, + RolloutAnalyzer, + SemanticGradient, +) + + +@dataclass(slots=True) +class PolicyTrainingEngine: + """Shared implementation of analyze -> estimate -> plan -> apply.""" + + rollout_analyzer: RolloutAnalyzer + gradient_estimator: GradientEstimator + policy_optimizer: PolicyOptimizer + policy_updater: PolicyUpdater + + async def analyze_estimate_plan_apply( + self, + *, + rollouts: list[Rollout], + policy_set: ExperienceSet, + ctx: PipelineContext, + ) -> tuple[list[RolloutAnalysis], list[SemanticGradient], PolicyUpdatePlan, PolicyApplyResult]: + analyses = await self.analyze_rollouts(rollouts, ctx) + gradients = await self.estimate_gradients(analyses, policy_set, ctx) + plan, apply_result = await self.plan_and_apply( + gradients=gradients, + policy_set=policy_set, + ctx=ctx, + ) + return analyses, gradients, plan, apply_result + + async def analyze_rollouts( + self, + rollouts: list[Rollout], + ctx: PipelineContext, + ) -> list[RolloutAnalysis]: + analyses = await asyncio.gather( + *[self.rollout_analyzer.analyze(rollout, ctx.analysis_context) for rollout in rollouts] + ) + return list(analyses) + + async def estimate_gradients( + self, + analyses: list[RolloutAnalysis], + policy_set: ExperienceSet, + ctx: PipelineContext, + ) -> list[SemanticGradient]: + gradient_batches = await asyncio.gather( + *[ + self.gradient_estimator.estimate( + analysis, + policy_set, + ctx.gradient_context, + ) + for analysis in analyses + ] + ) + return [gradient for batch in gradient_batches for gradient in batch] + + async def plan_and_apply( + self, + *, + gradients: list[SemanticGradient], + policy_set: ExperienceSet, + ctx: PipelineContext, + ) -> tuple[PolicyUpdatePlan, PolicyApplyResult]: + async with policy_set.lock() as transaction_handle: + latest_policy_set = await policy_set.reload() + plan = await self.policy_optimizer.plan( + gradients, + latest_policy_set, + ctx.optimization_context, + ) + apply_result = await self.policy_updater.apply( + plan, + latest_policy_set, + ctx.apply_context or latest_policy_set.request_context, + transaction_handle=transaction_handle, + ) + return plan, apply_result diff --git a/openviking/session/train/gradients.py b/openviking/session/train/gradients.py new file mode 100644 index 0000000000..9e0a015e96 --- /dev/null +++ b/openviking/session/train/gradients.py @@ -0,0 +1,46 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Semantic gradient implementations for policy optimization.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from openviking.session.memory.dataclass import MemoryFile, StoredLink + + +@dataclass(slots=True) +class PatchSemanticGradient: + """Patch-based semantic gradient for one target policy. + + A semantic gradient is represented as a typed before/after memory file pair. + The concrete patch text is a rendering concern owned by merge context + providers; the gradient itself carries structured memory-file state. + """ + + before_file: MemoryFile | None + after_file: MemoryFile + base_version: int | None + rationale: str + links: list[StoredLink] + confidence: float + metadata: dict[str, Any] = field(default_factory=dict) + + @property + def target_name(self) -> str: + fields = self.after_file.extra_fields or {} + memory_type = self.after_file.memory_type or fields.get("memory_type") or "experiences" + name = ( + fields.get("experience_name") + or fields.get("name") + or fields.get(f"{str(memory_type).rstrip('s')}_name") + ) + if name: + return str(name) + uri = self.target_uri + return uri.rstrip("/").split("/")[-1].removesuffix(".md") if uri else "unknown_policy" + + @property + def target_uri(self) -> str | None: + return self.after_file.uri or (self.before_file.uri if self.before_file is not None else None) diff --git a/openviking/session/train/interfaces.py b/openviking/session/train/interfaces.py new file mode 100644 index 0000000000..494b9b2f09 --- /dev/null +++ b/openviking/session/train/interfaces.py @@ -0,0 +1,161 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Protocol interfaces for the session training framework.""" + +from __future__ import annotations + +from collections.abc import AsyncIterator +from typing import Any, Protocol + +from openviking.session.memory.dataclass import MemoryFile, StoredLink +from openviking.session.train.context import ExecutionContext +from openviking.session.train.domain import ( + Case, + PipelineEvaluationResult, + PipelineResult, + PolicyApplyResult, + PolicySet, + PolicyUpdatePlan, + Rollout, + RolloutAnalysis, + RolloutTrainingResult, + RubricEvaluation, +) + + +class SemanticGradient(Protocol): + """A semantic update signal for one target policy.""" + + @property + def before_file(self) -> MemoryFile | None: ... + + @property + def after_file(self) -> MemoryFile: ... + + @property + def target_name(self) -> str: ... + + @property + def target_uri(self) -> str | None: ... + + @property + def base_version(self) -> int | None: ... + + @property + def rationale(self) -> str: ... + + @property + def links(self) -> list[StoredLink]: ... + + @property + def confidence(self) -> float: ... + + @property + def metadata(self) -> dict[str, Any]: ... + + +class PolicyOptimizer(Protocol): + """Plans policy-set updates from semantic gradients.""" + + async def plan( + self, + gradients: list[SemanticGradient], + policy_set: PolicySet, + context: Any, + ) -> PolicyUpdatePlan: ... + + +class PolicyUpdater(Protocol): + """Applies a policy update plan to a PolicySet.""" + + async def apply( + self, + plan: PolicyUpdatePlan, + policy_set: PolicySet, + context: Any, + *, + transaction_handle: Any = None, + ) -> PolicyApplyResult: ... + + +class CaseLoader(Protocol): + """Loads case batches for policy optimization.""" + + async def batches(self, context: Any) -> AsyncIterator[list[Case]]: ... + + +class RolloutExecutor(Protocol): + """Executes cases against a policy set and produces rollouts.""" + + async def execute( + self, + cases: list[Case], + policy_set: PolicySet, + context: ExecutionContext, + ) -> list[Rollout]: ... + + +class PolicySnapshotter(Protocol): + """Creates a snapshot identifier for a PolicySet.""" + + async def snapshot(self, policy_set: PolicySet, context: Any) -> str: ... + + +class RolloutAnalyzer(Protocol): + """Analyzes a rollout and extracts learning signals.""" + + async def analyze(self, rollout: Rollout, context: Any) -> RolloutAnalysis: ... + + +class RolloutEvaluator(Protocol): + """Evaluates a rollout before learning-signal extraction.""" + + async def evaluate(self, rollout: Rollout, context: Any) -> RubricEvaluation: ... + + +class GradientEstimator(Protocol): + """Estimates semantic gradients from rollout analysis.""" + + async def estimate( + self, + analysis: RolloutAnalysis, + policy_set: PolicySet, + context: Any, + ) -> list[SemanticGradient]: ... + + +class PolicyTrainer(Protocol): + """Trains a policy from rollout batches, optionally using precomputed analyses.""" + + async def train_rollouts( + self, + rollouts: list[Rollout], + policy_set: PolicySet, + context: Any, + analyses: list[RolloutAnalysis] | None = None, + ) -> RolloutTrainingResult: ... + + +class PolicyOptimizationPipeline(Protocol): + """Runs end-to-end policy optimization over case batches.""" + + async def train( + self, + case_loader: CaseLoader, + policy_set: PolicySet, + context: Any, + ) -> PipelineResult: ... + + async def eval( + self, + case_loader: CaseLoader, + policy_set: PolicySet, + context: Any, + ) -> PipelineEvaluationResult: ... + + async def train_from_rollouts( + self, + rollouts: list[Rollout], + policy_set: PolicySet, + context: Any, + ) -> RolloutTrainingResult: ... diff --git a/openviking/session/train/pipeline.py b/openviking/session/train/pipeline.py new file mode 100644 index 0000000000..f23c87fb0c --- /dev/null +++ b/openviking/session/train/pipeline.py @@ -0,0 +1,664 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Default orchestration for the session training framework.""" + +from __future__ import annotations + +import inspect +import time +from typing import Any + +from openviking.session.train.components.case_loader import make_trial_case_loader +from openviking.session.train.components.policy_trainer import BatchPolicyTrainer +from openviking.session.train.context import ( + ExecutionContext, + PipelineContext, + PipelineHookDecision, +) +from openviking.session.train.domain import ( + ExperienceSet, + PipelineEpochResult, + PipelineEvaluationResult, + PipelineResult, + PolicyApplyResult, + PolicyUpdatePlan, + RolloutAnalysis, + RolloutTrainingResult, +) +from openviking.session.train.interfaces import ( + CaseLoader, + GradientEstimator, + PolicyOptimizer, + PolicySnapshotter, + PolicyTrainer, + PolicyUpdater, + RolloutAnalyzer, + RolloutExecutor, + SemanticGradient, +) +from openviking.session.train.utils import average_score, validate_rollouts_have_cases +from openviking.telemetry import tracer + + +class OfflinePolicyOptimizationPipeline: + """Composable offline train/eval pipeline for case-driven policy optimization. + + This class wires the protocol interfaces together. It does not implement + rollout execution, LLM analysis, gradient estimation, optimization, or file + updates itself. + + ``train`` updates the policy set from case rollouts. ``eval`` only executes + and analyzes rollouts; it never estimates gradients or writes policy files. + Benchmark runners should explicitly compose them, for example: + ``eval(test) -> train(train) -> eval(test)``. + """ + + def __init__( + self, + *, + snapshotter: PolicySnapshotter, + rollout_executor: RolloutExecutor, + rollout_analyzer: RolloutAnalyzer, + gradient_estimator: GradientEstimator, + policy_optimizer: PolicyOptimizer, + policy_updater: PolicyUpdater, + policy_trainer: PolicyTrainer | None = None, + ) -> None: + self.snapshotter = snapshotter + self.rollout_executor = rollout_executor + self.policy_trainer = policy_trainer or BatchPolicyTrainer( + rollout_analyzer=rollout_analyzer, + gradient_estimator=gradient_estimator, + policy_optimizer=policy_optimizer, + policy_updater=policy_updater, + ) + + @tracer("train.pipeline.train", ignore_result=True, ignore_args=True) + async def train( + self, + case_loader: CaseLoader, + policy_set: ExperienceSet, + context: PipelineContext | Any, + ) -> PipelineResult: + ctx = context if isinstance(context, PipelineContext) else PipelineContext() + max_epochs = max(1, int(ctx.max_epochs or 1)) + case_loader = _train_case_loader(case_loader, ctx) + current_policy_set = policy_set + epoch_results: list[PipelineEpochResult] = [] + evaluation_passes: list[PipelineEvaluationResult] = [] + train_epoch_reports: list[dict[str, Any]] = [] + stop_decision: PipelineHookDecision | None = None + + for epoch in range(max_epochs): + ctx.execution_metadata["epoch"] = epoch + await _emit_epoch_start(ctx, epoch) + epoch_result = await self._run_training_epoch( + epoch=epoch, + case_loader=case_loader, + policy_set=current_policy_set, + ctx=ctx, + ) + epoch_results.append(epoch_result) + current_policy_set = epoch_result.apply_result.updated_policy_set + hook_decision = await _emit_epoch_end( + epoch_result=epoch_result, + policy_set=current_policy_set, + ctx=ctx, + ) + train_report = hook_decision.report if hook_decision is not None else None + if train_report is not None: + train_epoch_reports.append(train_report) + await _emit_train_report(ctx, train_report) + if hook_decision is not None and hook_decision.stop_training: + stop_decision = hook_decision + break + epoch_eval = await self._run_epoch_evaluation_pass( + epoch=epoch, + policy_set=current_policy_set, + ctx=ctx, + ) + if epoch_eval is not None: + evaluation_passes.append(epoch_eval) + + all_analyses = [analysis for epoch in epoch_results for analysis in epoch.analyses] + all_gradients: list[SemanticGradient] = [ + gradient for epoch in epoch_results for gradient in epoch.gradients + ] + + if epoch_results: + last_plan = epoch_results[-1].plan + last_apply_result = epoch_results[-1].apply_result + else: + last_plan = PolicyUpdatePlan(metadata={"empty": True}) + last_apply_result = PolicyApplyResult(updated_policy_set=current_policy_set) + + first_score = _first_epoch_score(epoch_results) + final_score = _final_epoch_score(epoch_results) + metadata: dict[str, Any] = { + "policy_set_root_uri": current_policy_set.root_uri, + "max_epochs": max_epochs, + "completed_epochs": len(epoch_results), + "evaluation_pass_count": len(evaluation_passes), + "train_reports": train_epoch_reports, + "stopped_early": stop_decision is not None, + } + if stop_decision is not None: + metadata["stop_reason"] = stop_decision.reason + if stop_decision.metadata: + metadata["stop_metadata"] = dict(stop_decision.metadata) + if first_score is not None: + metadata["first_score"] = first_score + if final_score is not None: + metadata["final_score"] = final_score + if first_score is not None and final_score is not None: + metadata["score_delta"] = final_score - first_score + + return PipelineResult( + analyses=all_analyses, + gradients=list(all_gradients), + plan=last_plan, + apply_result=last_apply_result, + epochs=epoch_results, + evaluation_passes=evaluation_passes, + metadata=metadata, + ) + + @tracer("train.pipeline.eval", ignore_result=True, ignore_args=True) + async def eval( + self, + case_loader: CaseLoader, + policy_set: ExperienceSet, + context: PipelineContext | Any, + ) -> PipelineEvaluationResult: + ctx = context if isinstance(context, PipelineContext) else PipelineContext() + eval_case_loader = _eval_case_loader(case_loader, ctx) + result = await self._run_evaluation_pass( + epoch=int(ctx.execution_metadata.get("epoch", 0) or 0), + case_loader=eval_case_loader, + policy_set=policy_set, + ctx=ctx, + ) + eval_report = await _emit_eval_end( + evaluation_result=result, + policy_set=policy_set, + ctx=ctx, + ) + if eval_report is None: + raise RuntimeError( + "pipeline eval requires a lifecycle hook to provide an evaluation report" + ) + result.metadata["report"] = eval_report + await _emit_eval_report(ctx, eval_report) + return result + + @tracer("train.pipeline.train_from_rollouts", ignore_result=True, ignore_args=True) + async def train_from_rollouts( + self, + rollouts, + policy_set: ExperienceSet, + context: PipelineContext | Any, + ) -> RolloutTrainingResult: + """Train directly from externally produced rollout records. + + This path is intended for realtime/online collection where another + component has already executed an agent loop and produced ``Rollout``s. + It deliberately skips ``CaseLoader``, ``PolicySnapshotter`` and + ``RolloutExecutor`` while reusing the same downstream training stages as + offline optimization: + + The configured PolicyTrainer owns downstream training semantics. The + default BatchPolicyTrainer analyzes rollouts locally; remote trainers + may submit raw rollouts to a server-side analyzer. + """ + + ctx = context if isinstance(context, PipelineContext) else PipelineContext() + rollout_list = list(rollouts) + validate_rollouts_have_cases(rollout_list) + result = await self.policy_trainer.train_rollouts( + rollout_list, + policy_set, + ctx, + ) + result.metadata["source"] = "external_rollouts" + return result + + async def _run_training_epoch( + self, + *, + epoch: int, + case_loader: CaseLoader, + policy_set: ExperienceSet, + ctx: PipelineContext, + ) -> PipelineEpochResult: + all_analyses: list[RolloutAnalysis] = [] + all_gradients: list[SemanticGradient] = [] + last_plan: PolicyUpdatePlan | None = None + last_apply_result: PolicyApplyResult | None = None + current_policy_set = policy_set + snapshot_ids: list[str] = [] + rollout_report: dict[str, Any] | None = None + + epoch_started_at = time.monotonic() + async for cases in case_loader.batches(ctx.case_load_context): + rollout_started_at = time.monotonic() + rollouts, snapshot_id = await self._rollout_batch( + cases=cases, + policy_set=current_policy_set, + ctx=ctx, + epoch=epoch, + training=True, + ) + rollout_cost_seconds = time.monotonic() - rollout_started_at + snapshot_ids.append(snapshot_id) + hook_rollout_report = await _emit_train_rollout_end( + epoch=epoch, + rollouts=rollouts, + snapshot_id=snapshot_id, + policy_set=current_policy_set, + ctx=ctx, + ) + rollout_report = _with_cost(hook_rollout_report, rollout_cost_seconds) + await _emit_train_rollout_report(ctx, rollout_report) + training_result = await self.policy_trainer.train_rollouts( + rollouts, + current_policy_set, + ctx, + ) + gradients = list(training_result.gradients) + all_analyses.extend(training_result.analyses) + last_plan = training_result.plan + last_apply_result = training_result.apply_result + all_gradients.extend(gradients) + current_policy_set = last_apply_result.updated_policy_set + + epoch_cost_seconds = time.monotonic() - epoch_started_at + if last_plan is None or last_apply_result is None: + last_plan = PolicyUpdatePlan(metadata={"empty": True, "epoch": epoch}) + last_apply_result = PolicyApplyResult(updated_policy_set=current_policy_set) + + return PipelineEpochResult( + epoch=epoch, + analyses=all_analyses, + gradients=list(all_gradients), + plan=last_plan, + apply_result=last_apply_result, + policy_snapshot_ids=snapshot_ids, + metadata={ + "score": average_score(all_analyses), + "analysis_count": len(all_analyses), + "gradient_count": len(all_gradients), + "train_rollout_report": rollout_report, + "cost_seconds": epoch_cost_seconds, + }, + ) + + async def _run_evaluation_pass( + self, + *, + epoch: int, + case_loader: CaseLoader, + policy_set: ExperienceSet, + ctx: PipelineContext, + ) -> PipelineEvaluationResult: + all_analyses: list[RolloutAnalysis] = [] + snapshot_ids: list[str] = [] + + started_at = time.monotonic() + async for cases in case_loader.batches(ctx.case_load_context): + rollouts, snapshot_id = await self._rollout_batch( + cases=cases, + policy_set=policy_set, + ctx=ctx, + epoch=epoch, + training=False, + ) + snapshot_ids.append(snapshot_id) + all_analyses.extend(_analyses_from_rollout_evaluations(rollouts)) + cost_seconds = time.monotonic() - started_at + + return PipelineEvaluationResult( + epoch=epoch, + analyses=all_analyses, + policy_snapshot_ids=snapshot_ids, + metadata={ + **dict(ctx.execution_metadata), + "score": average_score(all_analyses), + "analysis_count": len(all_analyses), + "evaluation_only": True, + "cost_seconds": cost_seconds, + }, + ) + + async def _run_epoch_evaluation_pass( + self, + *, + epoch: int, + policy_set: ExperienceSet, + ctx: PipelineContext, + ) -> PipelineEvaluationResult | None: + if ctx.eval_each_epoch_case_loader is None: + return None + eval_ctx = _epoch_eval_context(ctx, epoch=epoch) + eval_case_loader = _eval_case_loader(ctx.eval_each_epoch_case_loader, eval_ctx) + result = await self._run_evaluation_pass( + epoch=epoch, + case_loader=eval_case_loader, + policy_set=policy_set, + ctx=eval_ctx, + ) + eval_report = await _emit_eval_end( + evaluation_result=result, + policy_set=policy_set, + ctx=eval_ctx, + ) + if eval_report is None: + raise RuntimeError( + "pipeline eval requires a lifecycle hook to provide an evaluation report" + ) + result.metadata["report"] = eval_report + await _emit_eval_report(eval_ctx, eval_report) + return result + + async def _rollout_batch( + self, + *, + cases, + policy_set: ExperienceSet, + ctx: PipelineContext, + epoch: int, + training: bool, + ) -> tuple[list[Any], str]: + snapshot_id = await self.snapshotter.snapshot( + policy_set, + ctx.snapshot_context, + ) + stage = _rollout_stage(epoch=epoch, training=training) + if not training: + stage = ctx.execution_metadata.get("rollout_stage") or stage + execution_metadata = { + **dict(ctx.execution_metadata), + "epoch": epoch, + "training": training, + "stage": stage, + } + # ponytail: train rollouts must never inherit an eval rollout_stage — + # _stage_from_execution_metadata checks rollout_stage before stage, so + # a leaked value would mis-route artifacts into eval directories. + if training: + execution_metadata.pop("rollout_stage", None) + execution_context = ExecutionContext( + policy_snapshot_id=snapshot_id, + metadata=execution_metadata, + ) + rollouts = await self.rollout_executor.execute( + cases, + policy_set, + execution_context, + ) + return rollouts, snapshot_id + + +async def _emit_train_rollout_end( + *, + epoch: int, + rollouts: list[Any], + snapshot_id: str, + policy_set: ExperienceSet, + ctx: PipelineContext, +) -> dict[str, Any] | None: + hook_report: dict[str, Any] | None = None + for hook in ctx.lifecycle_hooks: + result = await _call_hook( + hook.on_train_rollout_end, + epoch=epoch, + rollouts=rollouts, + snapshot_id=snapshot_id, + policy_set=policy_set, + context=ctx, + ) + hook_report = _merge_report_hook_result( + hook_report, + result, + hook_name="on_train_rollout_end", + ) + return hook_report + + +def _merge_report_hook_result( + current: dict[str, Any] | None, + result: Any, + *, + hook_name: str, +) -> dict[str, Any] | None: + if result is None: + return current + if not isinstance(result, dict): + raise TypeError(f"{hook_name} must return dict or None, got {type(result).__name__}") + return result + + +async def _call_hook(method: Any, **kwargs: Any) -> Any: + result = method(**kwargs) + if inspect.isawaitable(result): + return await result + return result + + +async def _call_event_hook(method: Any, **kwargs: Any) -> None: + await _call_hook(method, **kwargs) + + +async def _emit_epoch_end( + *, + epoch_result: PipelineEpochResult, + policy_set: ExperienceSet, + ctx: PipelineContext, +) -> PipelineHookDecision | None: + hook_decision: PipelineHookDecision | None = None + for hook in ctx.lifecycle_hooks: + result = await _call_hook( + hook.on_epoch_end, + epoch_result=epoch_result, + policy_set=policy_set, + context=ctx, + ) + if result is None: + continue + if not isinstance(result, PipelineHookDecision): + raise TypeError( + "on_epoch_end must return PipelineHookDecision or None, " + f"got {type(result).__name__}" + ) + hook_decision = _merge_hook_decision(hook_decision, result) + return hook_decision + + +async def _emit_eval_end( + *, + evaluation_result: PipelineEvaluationResult, + policy_set: ExperienceSet, + ctx: PipelineContext, +) -> dict[str, Any] | None: + hook_report: dict[str, Any] | None = None + for hook in ctx.lifecycle_hooks: + result = await _call_hook( + hook.on_eval_end, + evaluation_result=evaluation_result, + policy_set=policy_set, + context=ctx, + ) + hook_report = _merge_report_hook_result( + hook_report, + result, + hook_name="on_eval_end", + ) + return hook_report + + +def _epoch_eval_context(ctx: PipelineContext, *, epoch: int) -> PipelineContext: + inherited_metadata = dict(ctx.execution_metadata) + execution_metadata = { + **inherited_metadata, + "epoch": epoch, + "training": False, + "rollout_stage": inherited_metadata.get("rollout_stage") or "test_rollout", + "eval_split": inherited_metadata.get("eval_split") or "test", + } + return PipelineContext( + case_load_context=ctx.case_load_context, + snapshot_context=ctx.snapshot_context, + analysis_context=ctx.analysis_context, + gradient_context=ctx.gradient_context, + optimization_context=ctx.optimization_context, + apply_context=ctx.apply_context, + execution_metadata=execution_metadata, + max_epochs=1, + eval_trials=ctx.eval_trials, + train_trials=ctx.train_trials, + trial_index_key=ctx.trial_index_key, + report_builder=ctx.report_builder, + lifecycle_hooks=list(ctx.lifecycle_hooks), + ) + + +def _train_case_loader(case_loader: CaseLoader, ctx: PipelineContext) -> CaseLoader: + train_trials = int(ctx.train_trials or 1) + if train_trials <= 1: + return case_loader + return make_trial_case_loader( + case_loader, + train_trials, + trial_input_key="train_trial", + ) + + +def _eval_case_loader(case_loader: CaseLoader, ctx: PipelineContext) -> CaseLoader: + eval_trials = int(ctx.eval_trials or 1) + if eval_trials <= 1: + return case_loader + return make_trial_case_loader( + case_loader, + eval_trials, + trial_input_key=ctx.trial_index_key, + ) + + +def _merge_hook_decision( + current: PipelineHookDecision | None, + incoming: PipelineHookDecision, +) -> PipelineHookDecision: + if current is None: + return incoming + return PipelineHookDecision( + stop_training=current.stop_training or incoming.stop_training, + reason=incoming.reason or current.reason, + metadata={**current.metadata, **incoming.metadata}, + report=incoming.report if incoming.report is not None else current.report, + ) + + +async def _emit_epoch_start(ctx: PipelineContext, epoch: int) -> None: + for hook in ctx.lifecycle_hooks: + await _call_event_hook(hook.on_epoch_start, epoch=epoch, context=ctx) + + +async def _emit_train_rollout_report( + ctx: PipelineContext, + report: dict[str, Any] | None, +) -> None: + if report is None: + return + for hook in ctx.lifecycle_hooks: + await _call_event_hook( + hook.on_train_rollout_report, + report=report, + context=ctx, + ) + + +async def _emit_train_report( + ctx: PipelineContext, + report: dict[str, Any] | None, +) -> None: + if report is None: + return + for hook in ctx.lifecycle_hooks: + await _call_event_hook(hook.on_train_report, report=report, context=ctx) + + +async def _emit_eval_report(ctx: PipelineContext, report: dict[str, Any] | None) -> None: + if report is None: + return + label = str( + report.get("label") + or ctx.execution_metadata.get("report_label") + or ctx.execution_metadata.get("rollout_stage") + or _rollout_stage( + epoch=int(ctx.execution_metadata.get("epoch", 0) or 0), + training=False, + ).split(maxsplit=1)[0] + ) + for hook in ctx.lifecycle_hooks: + await _call_event_hook( + hook.on_eval_report, + label=label, + report=report, + context=ctx, + ) + + +def _with_cost(report: dict[str, Any] | None, cost_seconds: float) -> dict[str, Any] | None: + if report is None: + return None + updated = dict(report) + updated["cost_seconds"] = max(0.0, float(cost_seconds)) + return updated + + +def _rollout_stage(*, epoch: int, training: bool) -> str: + if training: + return f"train_rollout epoch={epoch}" + if epoch < 0: + return "baseline_test_rollout" + return f"test_rollout epoch={epoch}" + + +def _analyses_from_rollout_evaluations(rollouts) -> list[RolloutAnalysis]: + analyses: list[RolloutAnalysis] = [] + for idx, rollout in enumerate(rollouts): + if rollout.evaluation is None: + raise ValueError( + "pipeline eval requires RolloutExecutor to provide rollout.evaluation; " + f"missing index={idx}, case={rollout.case.name}" + ) + analyses.append( + RolloutAnalysis( + evaluation=rollout.evaluation, + trajectories=[], + metadata={ + "rollout": rollout, + "rollout_messages": rollout.messages, + "policy_snapshot_id": rollout.policy_snapshot_id, + "evaluation_source": "rollout_executor", + }, + ) + ) + return analyses + + +def _first_epoch_score(epochs: list[PipelineEpochResult]) -> float | None: + for epoch in epochs: + score = average_score(epoch.analyses) + if score is not None: + return score + return None + + +def _final_epoch_score( + epochs: list[PipelineEpochResult], +) -> float | None: + for epoch in reversed(epochs): + score = average_score(epoch.analyses) + if score is not None: + return score + return None diff --git a/openviking/session/train/run_batch_train_eval.py b/openviking/session/train/run_batch_train_eval.py new file mode 100644 index 0000000000..1f40cadb4b --- /dev/null +++ b/openviking/session/train/run_batch_train_eval.py @@ -0,0 +1,224 @@ +#!/usr/bin/env python3 +"""CLI for remote benchmark batch policy train/eval.""" + +from __future__ import annotations + +import argparse +import asyncio +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[3] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Run remote benchmark batch policy train/eval") + parser.add_argument("--dataset", required=True, help="Remote benchmark dataset") + parser.add_argument("--domain", required=True, help="Benchmark domain") + parser.add_argument("--epochs", type=int, default=1, help="Training epochs (default: 1)") + parser.add_argument( + "--batch-size", + type=int, + default=None, + help="Train/eval batch size. Default uses the whole split as one batch.", + ) + parser.add_argument( + "--concurrency", + type=int, + default=200, + help="Concurrent rollout executions for train and eval (default: 200)", + ) + parser.add_argument( + "--commit-concurrency", + type=int, + default=200, + help="Concurrent OpenViking session.commit submissions during train (default: 200)", + ) + parser.add_argument("--config", default=None, help="ov.conf path (optional)") + parser.add_argument("--server-url", default=None, help="OpenViking server URL. Defaults to ov.conf/ovcli.conf") + parser.add_argument("--api-key", default=None, help="OpenViking API key. Defaults to ov.conf/ovcli.conf") + parser.add_argument("--account-id", default="default", help="OpenViking trusted account id. Default: default") + parser.add_argument("--user-id", default="default", help="OpenViking trusted user id. Default: default") + parser.add_argument("--output", default=None, help="JSON report output path") + parser.add_argument( + "--events-output", + default=None, + help="Streaming JSONL event output path. Defaults to report directory/events.jsonl.", + ) + parser.add_argument( + "--result-dir-name", + default="train", + help="Result subdirectory under result/{dataset}/ (default: train).", + ) + parser.add_argument( + "--benchmark-service-url", + default=None, + help="Benchmark runtime service URL, e.g. http://127.0.0.1:1944", + ) + parser.add_argument( + "--max-iterations", + type=int, + default=30, + help="Max steps/iterations per rollout (default: 30)", + ) + parser.add_argument( + "--train-index", + default=None, + help=( + "Run train sample(s) at 0-based split index/indices. " + "Accepts one index or comma-separated indices, e.g. 7 or 1,5,6." + ), + ) + parser.add_argument( + "--eval-index", + default=None, + help=( + "Run eval/test sample(s) at 0-based split index/indices. " + "Accepts one index or comma-separated indices, e.g. 3 or 10,14,18." + ), + ) + parser.add_argument( + "--force-baseline-recompute", + action="store_true", + help=( + "Recompute the cached pre-training baseline instead of reusing an " + "existing cache file." + ), + ) + parser.add_argument( + "--skip-baseline-eval", + action="store_true", + help="Skip the pre-training baseline eval/cache step.", + ) + parser.add_argument( + "--eval-split", + choices=("train", "test", "none"), + default="test", + help="Split used for baseline, per-epoch, and final eval (default: test; none disables eval).", + ) + parser.add_argument( + "--eval-each-epoch", + action="store_true", + help="Run held-out eval after every training epoch. Disabled by default.", + ) + parser.add_argument( + "--skip-final-eval", + action="store_true", + help=( + "Skip the final held-out eval pass. When --eval-each-epoch is enabled, " + "the last epoch eval is reused as final_eval in the report." + ), + ) + parser.add_argument( + "--trials", + type=int, + default=8, + help="Run each eval split N times and aggregate (default: 8).", + ) + parser.add_argument( + "--train-trials", + type=int, + default=1, + help="Run each train case N times per epoch (default: 1).", + ) + parser.add_argument( + "--reuse-train-rollout-cache", + action="store_true", + help=( + "Reuse cached epoch-0 train rollouts when available. Default is off; " + "cache misses still execute rollouts and populate the cache." + ), + ) + parser.add_argument( + "--keep-recent-results", + type=int, + default=5, + help=( + "When --clean-result is enabled, keep the most recent N default run_ " + "directories for the same domain while preserving cache/ and non-run_ " + "directories (default: 5)." + ), + ) + clean_group = parser.add_mutually_exclusive_group() + clean_group.add_argument( + "--clean-result", + dest="clean_result", + action="store_true", + default=True, + help="Clean previous default result/{dataset}/train artifacts before running (default).", + ) + clean_group.add_argument( + "--no-clean-result", + dest="clean_result", + action="store_false", + help="Keep previous result artifacts.", + ) + return parser.parse_args() + + +def _parse_indices_arg(value: str | None) -> list[int] | None: + if value is None or not str(value).strip(): + return None + indices: list[int] = [] + for part in str(value).split(","): + item = part.strip() + if not item: + continue + index = int(item) + if index < 0: + raise ValueError("indices must be >= 0") + indices.append(index) + return indices or None + + +async def main_async() -> int: + args = parse_args() + from openviking.session.train.batch_runner import ( + BatchTrainEvalConfig, + run_batch_train_eval, + ) + + report = await run_batch_train_eval( + BatchTrainEvalConfig( + dataset=args.dataset, + domain=args.domain, + epochs=args.epochs, + batch_size=args.batch_size, + concurrency=args.concurrency, + commit_concurrency=args.commit_concurrency, + config_path=str(Path(args.config).expanduser()) if args.config else None, + server_url=args.server_url, + api_key=args.api_key, + account_id=args.account_id, + user_id=args.user_id, + output_path=args.output, + events_path=args.events_output, + result_dir_name=args.result_dir_name, + keep_default_tools=True, + max_iterations=args.max_iterations, + train_index=_parse_indices_arg(args.train_index), + eval_index=_parse_indices_arg(args.eval_index), + benchmark_service_url=args.benchmark_service_url, + baseline_force_recompute=args.force_baseline_recompute, + eval_each_epoch=args.eval_each_epoch, + skip_final_eval=args.skip_final_eval, + skip_baseline_eval=args.skip_baseline_eval, + eval_split=args.eval_split, + trials=args.trials, + train_trials=args.train_trials, + reuse_train_rollout_cache=args.reuse_train_rollout_cache, + clean_result=args.clean_result, + keep_recent_results=args.keep_recent_results, + ) + ) + return 1 if any(epoch.get("errors") for epoch in report.train_epochs) else 0 + + +def main() -> None: + raise SystemExit(asyncio.run(main_async())) + + +if __name__ == "__main__": + main() diff --git a/openviking/session/train/run_batch_train_eval.sh b/openviking/session/train/run_batch_train_eval.sh new file mode 100755 index 0000000000..80a150b441 --- /dev/null +++ b/openviking/session/train/run_batch_train_eval.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Generic launcher for the OpenViking session/train remote benchmark batch pipeline. + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" + +load_user_env_file() { + local env_file="${OPENVIKING_ENV_FILE:-${HOME}/.openviking_benchmark_env}" + if [[ -z "${env_file}" || ! -f "${env_file}" ]]; then + return + fi + + local -a preserved_env=() + local entry + while IFS= read -r -d '' entry; do + if [[ "${entry}" != *= ]]; then + preserved_env+=("${entry}") + fi + done < <(env -0) + + echo "[batch-train-eval] loading env file ${env_file}" + set +u + set -a + # shellcheck source=/dev/null + source "${env_file}" + set +a + set -euo pipefail + + for entry in "${preserved_env[@]}"; do + export "${entry}" + done +} + +load_user_env_file + +PYTHON_BIN="${PYTHON_BIN:-python}" + +export PYTHONPATH="${REPO_ROOT}:${PYTHONPATH:-}" +export OPENVIKING_CONFIG_FILE="${OPENVIKING_CONFIG_FILE:-${HOME}/.openviking/ov.conf}" + +cd "${REPO_ROOT}" +exec "${PYTHON_BIN}" -m openviking.session.train.run_batch_train_eval "$@" diff --git a/openviking/session/train/utils.py b/openviking/session/train/utils.py new file mode 100644 index 0000000000..ece5786c8f --- /dev/null +++ b/openviking/session/train/utils.py @@ -0,0 +1,47 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Small shared helpers for the train pipeline.""" + +from __future__ import annotations + +from typing import Any + +from openviking.session.train.domain import Rollout, RolloutAnalysis + + +def average_score(analyses: list[RolloutAnalysis | None]) -> float | None: + """Return the mean evaluation score across analyses, ignoring None entries.""" + scores = [ + float(analysis.evaluation.score) + for analysis in analyses + if analysis is not None and analysis.evaluation is not None + ] + if not scores: + return None + return sum(scores) / len(scores) + + +def validate_rollouts_have_cases(rollouts: list[Rollout]) -> None: + """Raise ``ValueError`` if any rollout is missing its ``case``.""" + missing = [ + idx for idx, rollout in enumerate(rollouts) if getattr(rollout, "case", None) is None + ] + if missing: + raise ValueError( + "rollout training requires Rollout.case for all rollouts; " + f"missing indices={missing}" + ) + + +def safe_int(value: Any) -> int | None: + """Parse ``value`` as a positive integer, returning ``None`` on failure/zero.""" + try: + parsed = int(value) + except (TypeError, ValueError): + return None + return parsed if parsed > 0 else None + + +def first_uri(uris: list[str]) -> str | None: + """Return the first URI from a list, or ``None`` if empty.""" + return uris[0] if uris else None diff --git a/openviking/storage/transaction/lock_manager.py b/openviking/storage/transaction/lock_manager.py index a7e1a91852..3ef5e937f1 100644 --- a/openviking/storage/transaction/lock_manager.py +++ b/openviking/storage/transaction/lock_manager.py @@ -512,6 +512,10 @@ async def _redo_session_memory(self, info: Dict[str, Any]) -> None: ), timeout=60.0, ) + # extract_long_term_memories may return either a list[Context] + # (v2) or a {"contexts": [...], "session_skills": [...]} dict (v3). + if isinstance(memories, dict): + memories = memories.get("contexts", []) logger.info(f"Redo: extracted {len(memories)} memories from {archive_uri}") except Exception as e: logger.warning(f"Redo: memory extraction failed ({e}), falling back to queue") diff --git a/openviking/telemetry/__init__.py b/openviking/telemetry/__init__.py index 80a9ff2ddb..5405d8097c 100644 --- a/openviking/telemetry/__init__.py +++ b/openviking/telemetry/__init__.py @@ -13,7 +13,7 @@ from .registry import register_telemetry, resolve_telemetry, unregister_telemetry from .request import TelemetryRequest, TelemetrySelection, normalize_telemetry_request from .runtime import get_telemetry_runtime, set_telemetry_runtime -from .tracer import tracer +from .tracer import start_current_span, tracer __all__ = [ "OperationTelemetry", @@ -29,6 +29,7 @@ "register_telemetry", "resolve_telemetry", "set_telemetry_runtime", + "start_current_span", "tracer", "tracer_module", "unregister_telemetry", diff --git a/openviking/telemetry/tracer.py b/openviking/telemetry/tracer.py index bca47a37aa..5180f485ca 100644 --- a/openviking/telemetry/tracer.py +++ b/openviking/telemetry/tracer.py @@ -6,6 +6,7 @@ import inspect import json import logging +from contextlib import contextmanager from typing import Any, Callable, Optional from loguru import logger @@ -346,6 +347,14 @@ def from_trace_info(trace_info: str) -> Optional[Any]: return None +@contextmanager +def start_current_span(name: str, *, trace_id: Optional[str] = None): + """Start a span as the current context for an explicit code block.""" + + with tracer.start_as_current_span(name=name, trace_id=trace_id) as span: + yield span + + def start_span( name: str, trace_id: Optional[str] = None, @@ -512,22 +521,12 @@ def start_as_current_span(cls, name: str, context=None, trace_id=None): @staticmethod def get_trace_id() -> str: """Get the current trace ID as a hex string.""" - if _otel_tracer is None: - return "" - - try: - current_span = otel_trace.get_current_span() - if current_span is not None and hasattr(current_span, "context"): - trace_id = "{:032x}".format(current_span.context.trace_id) - return trace_id - except Exception: - _log_trace_internal_failure("[TRACER] failed to resolve decorator trace id") - return "" + return get_trace_id() @staticmethod def is_enabled() -> bool: """Check if tracer is enabled.""" - return _otel_tracer is not None + return is_enabled() @staticmethod def set(key: str, value: Any) -> None: diff --git a/openviking/utils/model_retry.py b/openviking/utils/model_retry.py index 9efe911038..8dd9bcbc15 100644 --- a/openviking/utils/model_retry.py +++ b/openviking/utils/model_retry.py @@ -87,6 +87,25 @@ "connection reset", ) +RETRYABLE_RATE_LIMIT_MARKERS = ( + "TooManyRequests", + "RateLimitExceeded", + "ModelAccountTpmRateLimitExceeded", + "TPM (Tokens Per Minute) limit", + "RPM (Requests Per Minute) limit", + "rate limit", + "rate_limit", +) + +_RATE_LIMIT_STATUS_RE = re.compile( + r"(?:\b(?:error\s*code|status(?:\s*code)?|http(?:\s*status)?|code)" + r"\s*[:=]?\s*429(?!\w)|(? bool: return classify_api_error(error) == ERROR_CLASS_TRANSIENT +def _load_rate_limit_error_classes() -> tuple[type[BaseException], ...]: + classes: list[type[BaseException]] = [] + try: + import openai + + classes.append(openai.RateLimitError) + except Exception: + pass + try: + from volcenginesdkarkruntime._exceptions import ArkRateLimitError + + classes.append(ArkRateLimitError) + except Exception: + pass + return tuple(classes) + + +def _iter_exception_chain(exc: BaseException) -> list[BaseException]: + chain: list[BaseException] = [] + seen: set[int] = set() + cur: BaseException | None = exc + while cur is not None and id(cur) not in seen: + chain.append(cur) + seen.add(id(cur)) + cur = cur.__cause__ or cur.__context__ + return chain + + +def _structured_rate_limit_match(exc: BaseException) -> bool: + global _RATE_LIMIT_ERROR_CLASSES + if not _RATE_LIMIT_ERROR_CLASSES: + _RATE_LIMIT_ERROR_CLASSES = _load_rate_limit_error_classes() + + for item in _iter_exception_chain(exc): + if _RATE_LIMIT_ERROR_CLASSES and isinstance(item, _RATE_LIMIT_ERROR_CLASSES): + return True + status_code = getattr(item, "status_code", None) + if status_code == 429 or str(status_code) == "429": + return True + code = getattr(item, "code", None) + error_type = getattr(item, "type", None) + if any( + isinstance(value, str) + and any(marker.lower() in value.lower() for marker in RETRYABLE_RATE_LIMIT_MARKERS) + for value in (code, error_type) + ): + return True + body = getattr(item, "body", None) + if isinstance(body, dict): + values = [body.get("code"), body.get("type"), body.get("message")] + if isinstance(body.get("error"), dict): + error = body["error"] + values.extend([error.get("code"), error.get("type"), error.get("message")]) + if any( + isinstance(value, str) + and any(marker.lower() in value.lower() for marker in RETRYABLE_RATE_LIMIT_MARKERS) + for value in values + ): + return True + return False + + +def is_retryable_rate_limit_error(exc: BaseException) -> bool: + """Return True for SDK/text-shaped LLM rate-limit errors. + + This intentionally lives in a lightweight OpenViking utility module so both + VikingBot provider adapters and benchmark integrations can share the same + classifier without importing each other's heavier runtime dependencies. + """ + if _structured_rate_limit_match(exc): + return True + text = str(exc or "") + if not text: + return False + lower_text = text.lower() + return any(marker.lower() in lower_text for marker in RETRYABLE_RATE_LIMIT_MARKERS) or bool( + _RATE_LIMIT_STATUS_RE.search(text) + ) + + +def rate_limit_retry_delay(attempt: int) -> float: + """Exponential backoff delay with jitter for LLM rate-limit retries.""" + delay = min( + RATE_LIMIT_RETRY_MAX_DELAY_SECONDS, + RATE_LIMIT_RETRY_BASE_DELAY_SECONDS * (2 ** max(0, attempt - 1)), + ) + return delay * random.uniform(0.8, 1.2) + + def _compute_delay( attempt: int, *, diff --git a/openviking_cli/client/_http_compat.py b/openviking_cli/client/_http_compat.py index 7d63768587..0a9e6b3058 100644 --- a/openviking_cli/client/_http_compat.py +++ b/openviking_cli/client/_http_compat.py @@ -6,6 +6,8 @@ from typing import Any, Dict +import httpx + from openviking_cli._sdk_import import import_openviking_sdk from openviking_cli.exceptions import ( AbortedError, @@ -86,6 +88,61 @@ def _raise_legacy_exception(error: Dict[str, Any]) -> None: class AsyncHTTPClient(import_openviking_sdk().AsyncHTTPClient): + def __init__(self, *args, **kwargs): + # Heavy local benchmark runs can keep OpenViking search requests queued + # behind embedding/vector work. Use a larger default read timeout than + # the upstream SDK's 60s while still respecting explicit caller values, + # including the SDK default when callers pass no timeout argument. + if "timeout" not in kwargs: + try: + import inspect + + sig = inspect.signature(import_openviking_sdk().AsyncHTTPClient.__init__) + params = [ + name + for name, param in sig.parameters.items() + if name != "self" + and param.kind + in (param.POSITIONAL_ONLY, param.POSITIONAL_OR_KEYWORD) + ] + timeout_index = params.index("timeout") + except Exception: + timeout_index = 7 + if len(args) <= timeout_index: + kwargs["timeout"] = 180.0 + super().__init__(*args, **kwargs) + + async def initialize(self) -> None: + # The upstream SDK uses httpx defaults (max_connections=100). High-parallel + # tau2 rollouts can exceed that from one shared client and hit PoolTimeout + # while waiting for a free connection, so raise the pool ceiling. + headers: Dict[str, str] = {} + if getattr(self, "_api_key", None): + headers["X-API-Key"] = self._api_key + if getattr(self, "_account", None): + headers["X-OpenViking-Account"] = self._account + if getattr(self, "_user_id", None): + headers["X-OpenViking-User"] = self._user_id + if getattr(self, "_actor_peer_id", None): + headers["X-OpenViking-Actor-Peer"] = self._actor_peer_id + headers.update(getattr(self, "_extra_headers", {}) or {}) + + max_connections = 512 + max_keepalive = 128 + self._http = httpx.AsyncClient( + base_url=self._url, + headers=headers, + timeout=self._timeout, + params={"profile": "1"} if self._profile_enabled else None, + limits=httpx.Limits( + max_connections=max_connections, + max_keepalive_connections=max_keepalive, + ), + ) + observer_cls = getattr(import_openviking_sdk().client, "_HTTPObserver", None) + if observer_cls is not None: + self._observer = observer_cls(self) + def _raise_exception(self, error: Dict[str, Any]) -> None: _raise_legacy_exception(error) diff --git a/openviking_cli/utils/config/memory_config.py b/openviking_cli/utils/config/memory_config.py index 7f33785ba5..c1caca226f 100644 --- a/openviking_cli/utils/config/memory_config.py +++ b/openviking_cli/utils/config/memory_config.py @@ -13,8 +13,8 @@ class MemoryConfig(BaseModel): """Memory configuration for OpenViking.""" version: str = Field( - default="v2", - description="Memory implementation version. Only 'v2' is supported.", + default="v3", + description="Deprecated and ignored. Memory extraction always uses v3.", ) custom_templates_dir: str = Field( default="", @@ -90,22 +90,31 @@ class MemoryConfig(BaseModel): @model_validator(mode="before") @classmethod - def drop_deprecated_agent_memory_enabled(cls, data: Any) -> Any: - if isinstance(data, dict) and "agent_memory_enabled" in data: + def drop_deprecated_memory_fields(cls, data: Any) -> Any: + if isinstance(data, dict): data = dict(data) - data.pop("agent_memory_enabled", None) - logger.warning( - "memory.agent_memory_enabled is deprecated and ignored; " - "use session memory_policy.memory_types to control trajectory/experience extraction" - ) + if "agent_memory_enabled" in data: + data.pop("agent_memory_enabled", None) + logger.debug( + "memory.agent_memory_enabled is deprecated and ignored; " + "use session memory_policy.memory_types to control trajectory/experience extraction" + ) + if "working_memory_enabled" in data: + data.pop("working_memory_enabled", None) + logger.debug( + "memory.working_memory_enabled is deprecated and ignored; " + "use session memory_policy.working_memory.enabled to control archive summaries" + ) return data - @field_validator("version") + @field_validator("version", mode="before") @classmethod - def validate_version(cls, value: str) -> str: - if value != "v2": - raise ValueError("memory.version only supports 'v2'; legacy memory v1 has been removed") - return value + def accept_deprecated_version(cls, value: Any) -> str: + if value not in (None, ""): + logger.debug( + "memory.version is deprecated and ignored; memory extraction always uses v3" + ) + return "v3" @classmethod def from_dict(cls, config: Dict[str, Any]) -> "MemoryConfig": diff --git a/openviking_cli/utils/config/vlm_config.py b/openviking_cli/utils/config/vlm_config.py index be8b5efb31..2243740ec5 100644 --- a/openviking_cli/utils/config/vlm_config.py +++ b/openviking_cli/utils/config/vlm_config.py @@ -64,7 +64,7 @@ class VLMConfig(BaseModel): temperature: float = Field(default=0.0, description="Generation temperature") max_retries: int = Field(default=3, description="Maximum retry attempts") timeout: float = Field( - default=60.0, + default=600.0, gt=0.0, description=( "Per-request HTTP timeout in seconds for VLM API calls. Applied to " diff --git a/pyproject.toml b/pyproject.toml index e522930089..be88afb16e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -258,6 +258,9 @@ python_files = ["test_*.py"] python_classes = ["Test*"] python_functions = ["test_*"] asyncio_mode = "auto" +markers = [ + "integration: tests that require external services or real model configuration", +] addopts = "-v --cov=openviking --cov-report=term-missing" [tool.ruff] diff --git a/tests/api_test/health_check/test_server_health_check.py b/tests/api_test/health_check/test_server_health_check.py index 04daa2e687..37460235df 100644 --- a/tests/api_test/health_check/test_server_health_check.py +++ b/tests/api_test/health_check/test_server_health_check.py @@ -26,7 +26,6 @@ def test_server_health_check(self, api_client): ) assert data.get("healthy") is not None, "'healthy' field should not be null" assert "version" in data, "'version' field should exist" - assert "user_id" in data, "'user_id' field should exist" except requests.exceptions.ConnectionError: pytest.fail("Could not connect to server service - service is not running") diff --git a/tests/integration/openviking_live_auth.py b/tests/integration/openviking_live_auth.py new file mode 100644 index 0000000000..07528781d1 --- /dev/null +++ b/tests/integration/openviking_live_auth.py @@ -0,0 +1,53 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Shared auth helpers for live OpenViking HTTP integration scripts.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path + +from openviking_cli.utils.config.config_loader import resolve_config_path +from openviking_cli.utils.config.consts import DEFAULT_OV_CONF, OPENVIKING_CONFIG_ENV + + +def root_api_key_from_ov_conf() -> str | None: + """Read ``server.root_api_key`` from ov.conf, like benchmark preflight scripts.""" + try: + config_path = resolve_config_path(None, OPENVIKING_CONFIG_ENV, DEFAULT_OV_CONF) + if config_path is None: + config_path = Path.home() / ".openviking" / DEFAULT_OV_CONF + path = Path(config_path).expanduser() + if not path.exists(): + return None + data = json.loads(path.read_text(encoding="utf-8-sig")) + key = str((data.get("server") or {}).get("root_api_key") or "").strip() + return key or None + except Exception: + return None + + +def resolve_api_key(cli_api_key: str | None = None) -> str | None: + """Resolve the API key without requiring scripts to hardcode one. + + Priority: + 1. Explicit CLI value + 2. OPENVIKING_API_KEY + 3. OPENVIKING_ROOT_API_KEY + 4. server.root_api_key from ov.conf + 5. None, allowing SyncHTTPClient to load ~/.openviking/ovcli.conf + """ + return ( + cli_api_key + or os.environ.get("OPENVIKING_API_KEY") + or os.environ.get("OPENVIKING_ROOT_API_KEY") + or root_api_key_from_ov_conf() + ) + + +API_KEY_HELP = ( + "API key. Defaults to OPENVIKING_API_KEY / OPENVIKING_ROOT_API_KEY / " + "server.root_api_key in ov.conf; when omitted, SyncHTTPClient may load " + "~/.openviking/ovcli.conf." +) diff --git a/tests/integration/test_agent_memory_e2e.py b/tests/integration/test_agent_memory_e2e.py index 8cb87219d5..8e81dcf2a2 100644 --- a/tests/integration/test_agent_memory_e2e.py +++ b/tests/integration/test_agent_memory_e2e.py @@ -15,8 +15,7 @@ Prerequisites ------------- -- ~/.openviking/ov.conf has: - "memory": { "version": "v2" } +- ~/.openviking/ov.conf has usable VLM/embedding configuration for memory extraction. Run --- @@ -309,7 +308,7 @@ def test_trajectory_and_experience_extraction( class TestAgentMemorySchemas: """Unit tests for agent memory schema filtering — no integration environment needed.""" - def test_no_agent_only_schemas_in_user_memory(self): + def test_no_agent_stage_schemas_in_user_memory(self): """ Verify that trajectory/experience schemas are filtered out from SessionExtractContextProvider (user memory path). diff --git a/tests/integration/test_compressor_v2_event_span_multiple_turns.py b/tests/integration/test_compressor_v2_event_span_multiple_turns.py index adf0f8d09a..ffba2a42f6 100644 --- a/tests/integration/test_compressor_v2_event_span_multiple_turns.py +++ b/tests/integration/test_compressor_v2_event_span_multiple_turns.py @@ -13,13 +13,17 @@ from rich.table import Table import openviking as ov +try: + from openviking_live_auth import API_KEY_HELP, resolve_api_key +except ModuleNotFoundError: # pytest/package import path + from tests.integration.openviking_live_auth import API_KEY_HELP, resolve_api_key # ── 常量 ─────────────────────────────────────────────────────────────────── DISPLAY_NAME = "小明" DEFAULT_URL = "http://localhost:1934" PANEL_WIDTH = 78 -DEFAULT_API_KEY = "1cf407c39990e5dc874ccc697942da4892208a86a44c4781396dfdc57aa5c98d" +DEFAULT_API_KEY = None DEFAULT_SESSION_ID = "event-span-multiple-turns" @@ -232,7 +236,7 @@ def main(): """入口函数""" parser = argparse.ArgumentParser(description=f"OpenViking 记忆演示 — {DISPLAY_NAME}") parser.add_argument("--url", default=DEFAULT_URL, help=f"Server URL (默认: {DEFAULT_URL})") - parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key") + parser.add_argument("--api-key", default=DEFAULT_API_KEY, help=API_KEY_HELP) parser.add_argument( "--phase", choices=["all", "ingest", "verify"], @@ -246,7 +250,7 @@ def main(): args = parser.parse_args() - client = ov.SyncHTTPClient(url=args.url, api_key=args.api_key, timeout=180) + client = ov.SyncHTTPClient(url=args.url, api_key=resolve_api_key(args.api_key), timeout=180) try: client.initialize() diff --git a/tests/integration/test_compressor_v2_tool_skill_memory.py b/tests/integration/test_compressor_v2_tool_skill_memory.py index 0c1828e22e..f3172e34bd 100644 --- a/tests/integration/test_compressor_v2_tool_skill_memory.py +++ b/tests/integration/test_compressor_v2_tool_skill_memory.py @@ -14,13 +14,17 @@ from rich.table import Table import openviking as ov +try: + from openviking_live_auth import API_KEY_HELP, resolve_api_key +except ModuleNotFoundError: # pytest/package import path + from tests.integration.openviking_live_auth import API_KEY_HELP, resolve_api_key # ── 常量 ─────────────────────────────────────────────────────────────────── DISPLAY_NAME = "测试用户" DEFAULT_URL = "http://localhost:1934" PANEL_WIDTH = 78 -DEFAULT_API_KEY = "1cf407c39990e5dc874ccc697942da4892208a86a44c4781396dfdc57aa5c98d" +DEFAULT_API_KEY = None DEFAULT_SESSION_ID = "tool-skill-memory-test" @@ -333,7 +337,7 @@ def main(): """入口函数""" parser = argparse.ArgumentParser(description="OpenViking 记忆演示 — 工具调用和Skill调用") parser.add_argument("--url", default=DEFAULT_URL, help=f"Server URL (默认: {DEFAULT_URL})") - parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key") + parser.add_argument("--api-key", default=DEFAULT_API_KEY, help=API_KEY_HELP) parser.add_argument( "--phase", choices=["all", "ingest", "verify"], @@ -347,7 +351,7 @@ def main(): args = parser.parse_args() - client = ov.SyncHTTPClient(url=args.url, api_key=args.api_key, timeout=180) + client = ov.SyncHTTPClient(url=args.url, api_key=resolve_api_key(args.api_key), timeout=180) try: client.initialize() diff --git a/tests/integration/test_compressor_v2_xiaomei.py b/tests/integration/test_compressor_v2_xiaomei.py index fdcd2bd452..71608aa5d5 100644 --- a/tests/integration/test_compressor_v2_xiaomei.py +++ b/tests/integration/test_compressor_v2_xiaomei.py @@ -13,13 +13,17 @@ from rich.table import Table import openviking as ov +try: + from openviking_live_auth import API_KEY_HELP, resolve_api_key +except ModuleNotFoundError: # pytest/package import path + from tests.integration.openviking_live_auth import API_KEY_HELP, resolve_api_key # ── 常量 ─────────────────────────────────────────────────────────────────── DISPLAY_NAME = "小美" DEFAULT_URL = "http://localhost:1934" PANEL_WIDTH = 78 -DEFAULT_API_KEY = "1cf407c39990e5dc874ccc697942da4892208a86a44c4781396dfdc57aa5c98d" +DEFAULT_API_KEY = None DEFAULT_SESSION_ID = "xiaomei-demo" @@ -253,7 +257,7 @@ def run_verify(client: ov.SyncHTTPClient): def main(): parser = argparse.ArgumentParser(description=f"OpenViking 记忆演示 — {DISPLAY_NAME}") parser.add_argument("--url", default=DEFAULT_URL, help=f"Server URL (默认: {DEFAULT_URL})") - parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key") + parser.add_argument("--api-key", default=DEFAULT_API_KEY, help=API_KEY_HELP) parser.add_argument( "--phase", choices=["all", "ingest", "verify"], @@ -275,7 +279,7 @@ def main(): ) ) - client = ov.SyncHTTPClient(url=args.url, api_key=args.api_key, timeout=180) + client = ov.SyncHTTPClient(url=args.url, api_key=resolve_api_key(args.api_key), timeout=180) try: client.initialize() diff --git a/tests/integration/test_compressor_v2_xiaowang.py b/tests/integration/test_compressor_v2_xiaowang.py index 84ef8df713..b81e13ae3d 100644 --- a/tests/integration/test_compressor_v2_xiaowang.py +++ b/tests/integration/test_compressor_v2_xiaowang.py @@ -13,13 +13,17 @@ from rich.table import Table import openviking as ov +try: + from openviking_live_auth import API_KEY_HELP, resolve_api_key +except ModuleNotFoundError: # pytest/package import path + from tests.integration.openviking_live_auth import API_KEY_HELP, resolve_api_key # ── 常量 ─────────────────────────────────────────────────────────────────── DISPLAY_NAME = "小王" DEFAULT_URL = "http://localhost:1934" PANEL_WIDTH = 78 -DEFAULT_API_KEY = "1cf407c39990e5dc874ccc697942da4892208a86a44c4781396dfdc57aa5c98d" +DEFAULT_API_KEY = None DEFAULT_SESSION_ID = "xiaowang-demo" @@ -208,7 +212,7 @@ def run_verify(client: ov.SyncHTTPClient): def main(): parser = argparse.ArgumentParser(description=f"OpenViking 记忆演示 — {DISPLAY_NAME}") parser.add_argument("--url", default=DEFAULT_URL, help=f"Server URL (默认: {DEFAULT_URL})") - parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key") + parser.add_argument("--api-key", default=DEFAULT_API_KEY, help=API_KEY_HELP) parser.add_argument( "--phase", choices=["all", "ingest", "verify"], @@ -230,7 +234,7 @@ def main(): ) ) - client = ov.SyncHTTPClient(url=args.url, api_key=args.api_key, timeout=180) + client = ov.SyncHTTPClient(url=args.url, api_key=resolve_api_key(args.api_key), timeout=180) try: client.initialize() diff --git a/tests/integration/test_compressor_v3_case_extraction.py b/tests/integration/test_compressor_v3_case_extraction.py new file mode 100644 index 0000000000..65033ee04f --- /dev/null +++ b/tests/integration/test_compressor_v3_case_extraction.py @@ -0,0 +1,362 @@ +#!/usr/bin/env python3 +""" +OpenViking V3 cases memory live integration test. + +This follows the style of ``test_compressor_v2_xiaomei.py``: it connects to a +running local OpenViking HTTP service, writes dialogue via add_message, commits +the session, and verifies that V3 extracts a trainable ``cases`` memory. + +Prerequisites for a full pass: +- OpenViking server is running locally (default: http://localhost:1933) +- Server has a usable VLM/embedding configuration for memory extraction +""" + +from __future__ import annotations + +import argparse +import json +import time +from datetime import datetime +from typing import Any + +import httpx +import pytest +from rich.console import Console +from rich.panel import Panel +from rich.table import Table + +import openviking as ov + +try: + from openviking_live_auth import API_KEY_HELP, resolve_api_key +except ModuleNotFoundError: # pytest/package import path + from tests.integration.openviking_live_auth import API_KEY_HELP, resolve_api_key + +# ── Constants ──────────────────────────────────────────────────────────────── + +DISPLAY_NAME = "V3 Case Extraction" +DEFAULT_URL = "http://localhost:1933" +DEFAULT_API_KEY = None +DEFAULT_SESSION_ID = "v3-case-extraction-demo" +PANEL_WIDTH = 90 +TASK_POLL_INTERVAL_SECONDS = 1.0 +TASK_TIMEOUT_SECONDS = 240.0 + +console = Console() + +# ── Dialogue that should trigger cases extraction ─────────────────────────── + +CONVERSATION = [ + { + "user": "请帮我处理酒店重复预订,只取消确认是重复的那一单,保留有效订单。", + "assistant": ( + "我先核对两个预订候选:订单 A 是用户主动确认保留的有效订单;订单 B 的入住日期、" + "房型、入住人都与 A 相同,并且系统标记为 duplicate_candidate。" + ), + }, + { + "user": "对,只能取消重复的那个,别影响我真正要住的订单。", + "assistant": ( + "已确认 B 是重复订单,A 是有效订单。我只取消订单 B,并保留订单 A。" + "取消后我再次检查,A 仍为 confirmed,B 已为 cancelled。" + ), + }, +] + +VERIFY_KEYWORDS = ["重复", "预订", "取消", "保留", "订单"] +CASES_DIR = "viking://user/default/memories/cases" + + +# ── Helpers ───────────────────────────────────────────────────────────────── + + +def _log(message: str) -> None: + console.print(f"[cyan][v3-case-extraction][/cyan] {message}") + + + +def _local_server_available(url: str = DEFAULT_URL) -> bool: + try: + response = httpx.get(f"{url.rstrip('/')}/health", timeout=2.0) + return response.status_code == 200 + except Exception: + return False + + +def _wait_for_task(client: ov.SyncHTTPClient, task_id: str, timeout: float) -> dict[str, Any]: + started = time.time() + while time.time() - started < timeout: + task = client.get_task(task_id) + status = task.get("status") if task else "not_found" + _log(f"poll task: task_id={task_id} status={status}") + if task and status in ("completed", "failed"): + elapsed = time.time() - started + _log(f"task finished: status={status} elapsed={elapsed:.2f}s result={task.get('result')}") + return task + time.sleep(TASK_POLL_INTERVAL_SECONDS) + raise TimeoutError(f"Task {task_id} did not finish within {timeout}s") + + +def _entry_uri(entry: Any, parent: str) -> str: + if isinstance(entry, dict): + uri = entry.get("uri") + if uri: + return str(uri) + name = entry.get("name") + if name: + return f"{parent.rstrip('/')}/{name}" + uri = getattr(entry, "uri", None) + if uri: + return str(uri) + name = getattr(entry, "name", None) + if name: + return f"{parent.rstrip('/')}/{name}" + return str(entry) + + + +def _read_memory_diff(client: ov.SyncHTTPClient, archive_uri: str | None) -> dict[str, Any]: + if not archive_uri: + return {} + diff_uri = f"{archive_uri.rstrip('/')}/memory_diff.json" + try: + raw = client.read(diff_uri) + data = json.loads(raw) + ops = data.get("operations") or {} + for kind in ("adds", "updates", "deletes"): + for item in ops.get(kind, []) or []: + _log( + "memory_diff " + f"{kind[:-1] if kind.endswith('s') else kind}: " + f"type={item.get('memory_type')} uri={item.get('uri')}" + ) + return data + except Exception as exc: + _log(f"memory_diff not readable: uri={diff_uri} error={exc}") + return {} + + +def _case_entries_from_memory_diff(diff: dict[str, Any]) -> list[str]: + entries: list[str] = [] + operations = diff.get("operations") or {} + for kind in ("adds", "updates"): + for item in operations.get(kind, []) or []: + if item.get("memory_type") == "cases" and item.get("uri"): + entries.append(str(item["uri"])) + return entries + + +def _read_case_memories(client: ov.SyncHTTPClient) -> list[tuple[str, str]]: + try: + entries = client.ls(CASES_DIR) + except Exception as exc: + _log(f"cases dir not readable yet: {CASES_DIR} error={exc}") + return [] + + memories: list[tuple[str, str]] = [] + for entry in entries or []: + uri = _entry_uri(entry, CASES_DIR) + if not uri.endswith(".md") or uri.endswith("/.overview.md") or uri.endswith("/.abstract.md"): + continue + try: + content = client.read(uri) + except Exception as exc: + _log(f"skip unreadable case memory: uri={uri} error={exc}") + continue + memories.append((uri, content)) + return memories + + +# ── Phase 1: ingest dialogue and commit session ───────────────────────────── + + +def run_ingest( + client: ov.SyncHTTPClient, + session_id: str = DEFAULT_SESSION_ID, + *, + wait_processed: bool = True, + task_timeout: float = TASK_TIMEOUT_SECONDS, +) -> dict[str, Any]: + console.rule(f"[bold]Phase 1: 写入对话并提交 Session — {DISPLAY_NAME}[/bold]") + + create_result = client.create_session(session_id=session_id) + session_id = create_result.get("session_id", session_id) + _log(f"session created: session_id={session_id}") + + session_time = datetime(2026, 6, 7, 9, 30) + session_time_str = session_time.isoformat() + + for index, turn in enumerate(CONVERSATION, 1): + _log(f"add dialogue turn {index}/{len(CONVERSATION)}") + client.add_message( + session_id, + role="user", + parts=[{"type": "text", "text": turn["user"]}], + created_at=session_time_str, + ) + client.add_message( + session_id, + role="assistant", + parts=[{"type": "text", "text": turn["assistant"]}], + created_at=session_time_str, + ) + + _log(f"added messages: count={len(CONVERSATION) * 2}") + _log("commit session: trigger archive + V3 long-term extraction") + commit_result = client.commit_session( + session_id, + memory_policy={"self": {"enabled": True}, "peer": {"enabled": False}}, + ) + _log( + "commit accepted: " + f"task_id={commit_result.get('task_id')} archive_uri={commit_result.get('archive_uri')} " + f"trace_id={commit_result.get('trace_id')}" + ) + + task = None + task_id = commit_result.get("task_id") + if task_id: + task = _wait_for_task(client, task_id, task_timeout) + if task.get("status") == "failed": + raise AssertionError(f"session commit task failed: {task}") + + if wait_processed: + _log("wait_processed: drain vectorization/semantic queues") + client.wait_processed(timeout=task_timeout) + + session_info = client.get_session(session_id) + _log(f"session info: {session_info}") + return {"session_id": session_id, "commit": commit_result, "task": task} + + +# ── Phase 2: verify cases memory ──────────────────────────────────────────── + + +def run_verify(client: ov.SyncHTTPClient, archive_uri: str | None = None) -> list[tuple[str, str]]: + console.rule(f"[bold]Phase 2: 验证 cases 记忆写入 — {DISPLAY_NAME}[/bold]") + + diff = _read_memory_diff(client, archive_uri) + diff_case_uris = _case_entries_from_memory_diff(diff) + memories = _read_case_memories(client) + table = Table(title="V3 cases memories", show_header=True, header_style="bold") + table.add_column("#", width=4) + table.add_column("URI", style="cyan", max_width=56) + table.add_column("命中关键词", style="green") + table.add_column("片段", max_width=80) + + for index, (uri, content) in enumerate(memories, 1): + hits = [keyword for keyword in VERIFY_KEYWORDS if keyword in content] + snippet = content.replace("\n", " ")[:160] + table.add_row(str(index), uri, ", ".join(hits), snippet) + _log(f"case memory: uri={uri} chars={len(content)} hits={hits}") + + console.print(table) + + matching = [ + (uri, content) + for uri, content in memories + if "重复" in content and "预订" in content and ("取消" in content or "保留" in content) + ] + if not matching and diff_case_uris: + # Directory listing can lag/shape-shift across deployments; if memory_diff says + # cases were written, read those URIs directly. + for uri in diff_case_uris: + try: + content = client.read(uri) + except Exception as exc: + _log(f"case uri from memory_diff unreadable: uri={uri} error={exc}") + continue + if "重复" in content and "预订" in content and ("取消" in content or "保留" in content): + matching.append((uri, content)) + + if not matching: + diff_types = [] + operations = diff.get("operations") or {} + for kind in ("adds", "updates", "deletes"): + diff_types.extend( + str(item.get("memory_type")) + for item in operations.get(kind, []) or [] + if item.get("memory_type") + ) + raise AssertionError( + "No V3 cases memory matched duplicate-booking evidence. " + f"cases_dir={CASES_DIR} memory_count={len(memories)} " + f"memory_diff_types={diff_types}" + ) + return matching + + +# ── Pytest entry: requires local server ───────────────────────────────────── + + +@pytest.mark.integration +@pytest.mark.skipif( + not _local_server_available(DEFAULT_URL), + reason=f"OpenViking local server is not running at {DEFAULT_URL}", +) +def test_local_service_add_dialogue_commit_triggers_v3_case_extraction(): + client = ov.SyncHTTPClient(url=DEFAULT_URL, api_key=resolve_api_key(), timeout=300) + try: + client.initialize() + ingest = run_ingest(client, session_id=f"{DEFAULT_SESSION_ID}-{int(time.time())}") + matches = run_verify(client, archive_uri=ingest["commit"].get("archive_uri")) + assert matches + finally: + client.close() + + +# ── CLI entry, like test_compressor_v2_xiaomei.py ─────────────────────────── + + +def main() -> None: + parser = argparse.ArgumentParser(description=f"OpenViking V3 cases live test — {DISPLAY_NAME}") + parser.add_argument("--url", default=DEFAULT_URL, help=f"Server URL (default: {DEFAULT_URL})") + parser.add_argument( + "--api-key", + default=None, + help=API_KEY_HELP, + ) + parser.add_argument("--session-id", default=DEFAULT_SESSION_ID, help="Session ID") + parser.add_argument( + "--phase", + choices=["all", "ingest", "verify"], + default="all", + help="all=ingest+verify, ingest=only submit session, verify=only read cases", + ) + parser.add_argument("--timeout", type=float, default=TASK_TIMEOUT_SECONDS, help="Task timeout") + args = parser.parse_args() + + console.print( + Panel( + f"[bold]OpenViking V3 cases live test — {DISPLAY_NAME}[/bold]\n" + f"Server: {args.url} | Phase: {args.phase}", + style="magenta", + width=PANEL_WIDTH, + ) + ) + + api_key = resolve_api_key(args.api_key) + if api_key: + _log("api key resolved from --api-key/env/ov.conf (not printed)") + else: + _log("no api key resolved; relying on SyncHTTPClient ovcli.conf auto-load") + client = ov.SyncHTTPClient(url=args.url, api_key=api_key, timeout=max(args.timeout, 60)) + try: + client.initialize() + _log(f"connected: {args.url}") + ingest = None + if args.phase in ("all", "ingest"): + ingest = run_ingest(client, session_id=args.session_id, task_timeout=args.timeout) + if args.phase in ("all", "verify"): + archive_uri = ingest["commit"].get("archive_uri") if ingest else None + run_verify(client, archive_uri=archive_uri) + console.print(Panel("[bold green]V3 cases live test completed[/bold green]", style="green")) + except Exception as exc: + console.print(Panel(f"[bold red]Error:[/bold red] {exc}", style="red")) + raise + finally: + client.close() + + +if __name__ == "__main__": + main() diff --git a/tests/models/vlm/test_timeout_config.py b/tests/models/vlm/test_timeout_config.py index e6bc71ca3d..b9dcc0bcbb 100644 --- a/tests/models/vlm/test_timeout_config.py +++ b/tests/models/vlm/test_timeout_config.py @@ -4,7 +4,7 @@ Before this was wired through, ``_build_openai_client_kwargs`` exposed a ``timeout`` parameter (#1208) but callers never passed it, so the default -60.0s was always used and end users could not override it via ``ov.conf``. +timeout was always used and end users could not override it via ``ov.conf``. These tests lock in that the config value flows through to the underlying OpenAI and LiteLLM clients. """ @@ -26,9 +26,9 @@ def test_vlm_config_accepts_timeout(): assert cfg.timeout == 120.0 -def test_vlm_config_timeout_defaults_to_60(): +def test_vlm_config_timeout_defaults_to_600(): cfg = VLMConfig(model="gpt-4o-mini", api_key="sk-x") - assert cfg.timeout == 60.0 + assert cfg.timeout == 600.0 def test_vlm_config_rejects_non_positive_timeout(): @@ -38,7 +38,7 @@ def test_vlm_config_rejects_non_positive_timeout(): def test_build_openai_client_kwargs_default_timeout(): kwargs = _build_openai_client_kwargs("openai", "sk-x", "https://example.invalid", None, None) - assert kwargs["timeout"] == 60.0 + assert kwargs["timeout"] == 600.0 def test_build_openai_client_kwargs_custom_timeout(): @@ -70,7 +70,7 @@ def test_openai_vlm_propagates_config_timeout(): assert fake.call_args.kwargs.get("timeout") == 120.0 -def test_openai_vlm_defaults_to_60_timeout_when_config_omits_it(): +def test_openai_vlm_defaults_to_600_timeout_when_config_omits_it(): vlm = OpenAIVLM( { "provider": "openai", @@ -79,11 +79,11 @@ def test_openai_vlm_defaults_to_60_timeout_when_config_omits_it(): "api_base": "https://example.invalid", } ) - assert vlm.timeout == 60.0 + assert vlm.timeout == 600.0 with mock.patch("openviking.models.vlm.backends.openai_vlm.openai.OpenAI") as fake: vlm.get_client() - assert fake.call_args.kwargs.get("timeout") == 60.0 + assert fake.call_args.kwargs.get("timeout") == 600.0 def test_litellm_build_kwargs_includes_timeout(): diff --git a/tests/server/test_server_health.py b/tests/server/test_server_health.py index f7f1fac3df..b8296ea5fb 100644 --- a/tests/server/test_server_health.py +++ b/tests/server/test_server_health.py @@ -72,6 +72,31 @@ async def test_health_endpoint(client: httpx.AsyncClient): assert body["status"] == "ok" +async def test_health_endpoint_never_resolves_identity(caplog): + app = create_app( + config=ServerConfig( + auth_mode="trusted", + host="127.0.0.1", + root_api_key="test-root-key", + ), + service=SimpleNamespace(), + ) + transport = httpx.ASGITransport(app=app) + + with caplog.at_level("WARNING", logger="openviking.server.routers.system"): + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + resp = await client.get("/health", headers={"X-API-Key": "test-root-key"}) + + assert resp.status_code == 200 + body = resp.json() + assert body["status"] == "ok" + assert body["auth_mode"] == "trusted" + assert "account_id" not in body + assert "user_id" not in body + assert "role" not in body + assert "Failed to resolve identity" not in caplog.text + + async def test_system_status(client: httpx.AsyncClient): resp = await client.get("/api/v1/system/status") assert resp.status_code == 200 diff --git a/tests/session/memory/test_agent_experience_context_provider.py b/tests/session/memory/test_agent_experience_context_provider.py index 4309014e7c..0f3ce2f65c 100644 --- a/tests/session/memory/test_agent_experience_context_provider.py +++ b/tests/session/memory/test_agent_experience_context_provider.py @@ -10,6 +10,14 @@ from openviking.session.memory.agent_experience_context_provider import ( AgentExperienceContextProvider, ) +from openviking.session.memory.agent_trajectory_context_provider import ( + AgentTrajectoryContextProvider, +) +from openviking.session.memory.session_extract_context_provider import ( + SessionExtractContextProvider, +) +from openviking.message import Message +from openviking.message.part import TextPart from openviking_cli.session.user_id import UserIdentifier @@ -30,6 +38,18 @@ def test_create_tool_context_uses_extract_context_page_id_map(): assert tool_ctx.page_id_map is extract_context.page_id_map +def test_user_memory_provider_splits_but_trajectory_provider_keeps_messages_whole(): + text = "第一句很长很长很长很长很长很长很长很长很长很长很长。" * 8 + messages = [Message(id="1", role="user", parts=[TextPart(text=text)])] + + user_provider = SessionExtractContextProvider(messages=messages) + trajectory_provider = AgentTrajectoryContextProvider(messages=messages) + + assert len(user_provider.get_extract_context().messages) > 1 + assert len(trajectory_provider.get_extract_context().messages) == 1 + assert trajectory_provider.get_extract_context().messages[0] is messages[0] + + @pytest.mark.asyncio async def test_agent_experience_prefetch_starts_with_conversation_and_new_trajectory_read(): provider = AgentExperienceContextProvider( @@ -108,3 +128,33 @@ async def test_agent_experience_prefetch_includes_structured_read_results(): == "candidate_experience" ) assert add_tool_call_pair.call_args_list[1].kwargs["result"]["page_id"] == 1 + + +@pytest.mark.asyncio +async def test_agent_experience_prefetch_missing_experience_dir_returns_empty_candidates(): + provider = AgentExperienceContextProvider( + messages=[], + trajectory_summary="new execution", + trajectory_uri="viking://user/user_1/memories/trajectories/new_execution.md", + ) + provider._ctx = RequestContext( + user=UserIdentifier(account_id="acc", user_id="user_1"), + role=Role.USER, + ) + provider._viking_fs = AsyncMock() + provider._viking_fs.ls = AsyncMock( + side_effect=Exception("Directory not found: viking://user/user_1/memories/experiences") + ) + provider._transaction_handle = None + provider.search_files = AsyncMock(return_value=[]) + + with patch( + "openviking.session.memory.agent_experience_context_provider.tracer.error" + ) as tracer_error, patch( + "openviking.session.memory.agent_experience_context_provider.add_tool_call_pair_to_messages" + ) as add_tool_call_pair: + messages = await provider.prefetch() + + assert messages[-1]["role"] == "user" + assert add_tool_call_pair.call_count == 1 + tracer_error.assert_not_called() diff --git a/tests/session/memory/test_compressor_v2.py b/tests/session/memory/test_compressor_v2.py index 79143a8073..e02f233213 100644 --- a/tests/session/memory/test_compressor_v2.py +++ b/tests/session/memory/test_compressor_v2.py @@ -339,7 +339,7 @@ async def run(self): SimpleNamespace( write_uris=[], edit_uris=[], - delete_uris=[], + delete_ids=[], ), [], ) @@ -569,7 +569,7 @@ async def run(self): SimpleNamespace( write_uris=[], edit_uris=[], - delete_uris=[], + delete_ids=[], ), [], ) @@ -870,12 +870,12 @@ class DummyVLM: def __init__(self): self.responses = [ - '{"profile":[{"page_id":1,"content":{"blocks":[{"search":"- Has been reading as usual","replace":"- Has been reading as usual (as of 2023-11-11)"}]} }],"delete_uris":[]}', - '{"profile":[{"page_id":1,"content":{"blocks":[{"search":"- Likes reading","replace":"- Likes reading\n- Has been reading as usual (as of 2023-11-11)"}]} }],"delete_uris":[]}', + '{"profile":[{"page_id":1,"content":{"blocks":[{"search":"- Has been reading as usual","replace":"- Has been reading as usual (as of 2023-11-11)"}]} }],"delete_ids":[]}', + '{"profile":[{"page_id":1,"content":{"blocks":[{"search":"- Likes reading","replace":"- Likes reading\n- Has been reading as usual (as of 2023-11-11)"}]} }],"delete_ids":[]}', ] self.messages = [] - async def get_completion_async(self, messages, tools=None, tool_choice=None): + async def get_completion_async(self, messages, tools=None, tool_choice=None, thinking=False): self.messages.append(list(messages)) return self.responses.pop(0) @@ -970,12 +970,12 @@ class DummyVLM: def __init__(self): self.responses = [ - '{"profile":[{"page_id":1,"content":{"blocks":[{"search":"- Missing one","replace":"- Fixed one"}]} }],"delete_uris":[]}', - '{"profile":[{"page_id":1,"content":{"blocks":[{"search":"- Missing two","replace":"- Fixed two"}]} }],"delete_uris":[]}', + '{"profile":[{"page_id":1,"content":{"blocks":[{"search":"- Missing one","replace":"- Fixed one"}]} }],"delete_ids":[]}', + '{"profile":[{"page_id":1,"content":{"blocks":[{"search":"- Missing two","replace":"- Fixed two"}]} }],"delete_ids":[]}', ] self.messages = [] - async def get_completion_async(self, messages, tools=None, tool_choice=None): + async def get_completion_async(self, messages, tools=None, tool_choice=None, thinking=False): self.messages.append(list(messages)) return self.responses.pop(0) @@ -1069,11 +1069,11 @@ class DummyVLM: def __init__(self): self.responses = [ - '{"profile":[{"page_id":1,"content":{"blocks":[{"search":"- Likes reading","replace":"- Likes reading every night (as of 2023-11-11)"}]} }],"delete_uris":[]}', + '{"profile":[{"page_id":1,"content":{"blocks":[{"search":"- Likes reading","replace":"- Likes reading every night (as of 2023-11-11)"}]} }],"delete_ids":[]}', ] self.messages = [] - async def get_completion_async(self, messages, tools=None, tool_choice=None): + async def get_completion_async(self, messages, tools=None, tool_choice=None, thinking=False): self.messages.append(list(messages)) return self.responses.pop(0) diff --git a/tests/session/memory/test_json_stability.py b/tests/session/memory/test_json_stability.py index ca9a6150de..6ffe9c06ac 100644 --- a/tests/session/memory/test_json_stability.py +++ b/tests/session/memory/test_json_stability.py @@ -287,7 +287,7 @@ class SimpleOperations(BaseModel): write_uris: List["TestMemoryOperationsIntegration.SimpleWriteOperation"] = Field( default_factory=list ) - delete_uris: List[str] = Field(default_factory=list) + delete_ids: List[str] = Field(default_factory=list) def test_parses_nested_write_operations(self): """Test nested write operations parse correctly.""" @@ -305,16 +305,16 @@ def test_parses_nested_write_operations(self): assert data.write_uris[0].topic == "theme" def test_handles_string_instead_of_list_for_delete(self): - """Test single string for delete_uris wraps to list via tolerance.""" + """Test single string for delete_ids wraps to list via tolerance.""" # Note: This would need field-level tolerance applied content = """{ "reasonning": "Removed old memory", - "delete_uris": "viking://user/default/memories/old.md" + "delete_ids": "viking://user/default/memories/old.md" }""" # First parse as raw dict data, error = parse_json_with_stability(content) assert error is None - assert data["delete_uris"] == "viking://user/default/memories/old.md" + assert data["delete_ids"] == "viking://user/default/memories/old.md" def test_recoverable_invalid_list_item_logs_below_error(self): """Test recoverable invalid list items do not emit error-level logs.""" diff --git a/tests/session/memory/test_memory_diff.py b/tests/session/memory/test_memory_diff.py index 270bbdd2cf..e9f8c1a176 100644 --- a/tests/session/memory/test_memory_diff.py +++ b/tests/session/memory/test_memory_diff.py @@ -474,6 +474,8 @@ async def test_build_memory_diff_empty(self, compressor, mock_viking_fs, mock_ct assert diff["summary"]["total_updates"] == 0 assert diff["summary"]["total_deletes"] == 0 assert "extracted_at" in diff + assert "trace_id" in diff + assert diff["trace_id"] is None assert diff["archive_uri"] == "" @pytest.mark.asyncio @@ -496,13 +498,14 @@ class TestMemoryDiffStructure: def test_memory_diff_structure(self): """Verify memory_diff.json structure.""" # This test validates the expected structure - expected_keys = ["archive_uri", "extracted_at", "operations", "summary"] + expected_keys = ["archive_uri", "trace_id", "extracted_at", "operations", "summary"] # We verify this through the actual implementation tests above # This is a placeholder for documentation assert set(expected_keys).issubset( { "archive_uri", + "trace_id", "extracted_at", "operations", "summary", diff --git a/tests/session/memory/test_memory_isolation_handler.py b/tests/session/memory/test_memory_isolation_handler.py index 19cd661ed0..612737755b 100644 --- a/tests/session/memory/test_memory_isolation_handler.py +++ b/tests/session/memory/test_memory_isolation_handler.py @@ -109,6 +109,74 @@ def test_empty_messages_uses_ctx_defaults(self): assert scope.user_ids == ["default_user"] + + def test_get_read_scope_filters_self_sentinel_from_peer_scope(self): + ctx = create_ctx(user_id="support_bot") + messages = [create_message("user")] + extract_ctx = create_mock_extract_context(messages) + handler = MemoryIsolationHandler( + ctx, + extract_ctx, + allow_self=True, + allowed_peer_ids={"__self", "web-visitor-alice"}, + ) + + scope = handler.get_read_scope() + + assert scope.user_ids == ["support_bot"] + assert scope.peer_ids == ["web-visitor-alice"] + + def test_render_schema_directories_self_sentinel_maps_to_user_space(self): + from openviking.session.memory.dataclass import MemoryTypeSchema + from openviking.session.memory.memory_isolation_handler import peer_user_space + + assert peer_user_space("support_bot", "__self") == "support_bot" + + ctx = create_ctx(user_id="support_bot") + messages = [create_message("user")] + extract_ctx = create_mock_extract_context(messages) + handler = MemoryIsolationHandler( + ctx, + extract_ctx, + allow_self=True, + allowed_peer_ids={"__self", "web-visitor-alice"}, + ) + schema = MemoryTypeSchema( + memory_type="preferences", + filename_template="preferences.md", + directory="viking://user/{{ user_space }}/memories", + ) + + dirs = handler.render_schema_directories(schema) + + assert "viking://user/support_bot/memories" in dirs + assert "viking://user/support_bot/peers/__self/memories" not in dirs + assert "viking://user/support_bot/peers/web-visitor-alice/memories" in dirs + + def test_render_schema_directories_peer_enabled_false_uses_self_only(self): + from openviking.session.memory.dataclass import MemoryTypeSchema + + ctx = create_ctx(user_id="support_bot") + messages = [create_message("user", peer_id="web-visitor-alice")] + extract_ctx = create_mock_extract_context(messages) + handler = MemoryIsolationHandler( + ctx, + extract_ctx, + allow_self=True, + allowed_peer_ids={"web-visitor-alice"}, + ) + schema = MemoryTypeSchema( + memory_type="cases", + filename_template="{{ case_name }}.md", + directory="viking://user/{{ user_space }}/memories/cases", + peer_enabled=False, + ) + + dirs = handler.render_schema_directories(schema) + + assert dirs == ["viking://user/support_bot/memories/cases"] + + class TestFillIdentityFields: """Tests for fill_identity_fields.""" @@ -413,9 +481,18 @@ def test_calculate_memory_uris_routes_ranges_to_self_and_peer(self, mock_generat assert "peer_id" not in operation.memory_fields @patch("openviking.session.memory.memory_isolation_handler.generate_uri") - def test_calculate_memory_uris_rejects_unallowed_peer_id(self, mock_generate_uri): + def test_calculate_memory_uris_unallowed_peer_id_falls_back_to_first_peer( + self, mock_generate_uri + ): + mock_generate_uri.side_effect = lambda **kwargs: ( + f"viking://user/{kwargs.get('user_space')}/memories/preferences" + ) + ctx = create_ctx(user_id="support_bot") - messages = [create_message("user", peer_id="web-visitor-alice")] + messages = [ + create_message("user", peer_id="web-visitor-bob"), + create_message("user", peer_id="web-visitor-alice"), + ] extract_ctx = create_mock_extract_context(messages) handler = MemoryIsolationHandler( ctx, @@ -439,11 +516,13 @@ def test_calculate_memory_uris_rejects_unallowed_peer_id(self, mock_generate_uri uris = handler.calculate_memory_uris(schema, operation, extract_ctx) - assert uris == [] - mock_generate_uri.assert_not_called() + assert uris == [ + "viking://user/support_bot/peers/web-visitor-bob/memories/preferences" + ] + assert operation.memory_fields["peer_id"] == "web-visitor-bob" @patch("openviking.session.memory.memory_isolation_handler.generate_uri") - def test_calculate_memory_uris_missing_peer_id_prefers_self_when_self_user_message_exists( + def test_calculate_memory_uris_missing_peer_id_prefers_first_peer_even_when_self_exists( self, mock_generate_uri ): mock_generate_uri.side_effect = lambda **kwargs: ( @@ -480,9 +559,11 @@ def test_calculate_memory_uris_missing_peer_id_prefers_self_when_self_user_messa uris = handler.calculate_memory_uris(schema, operation, extract_ctx) - assert uris == ["viking://user/support_bot/memories/preferences"] + assert uris == [ + "viking://user/support_bot/peers/web-visitor-alice/memories/preferences" + ] assert operation.memory_fields["user_id"] == "support_bot" - assert "peer_id" not in operation.memory_fields + assert operation.memory_fields["peer_id"] == "web-visitor-alice" @patch("openviking.session.memory.memory_isolation_handler.generate_uri") def test_calculate_memory_uris_missing_peer_id_falls_back_to_first_peer_when_self_absent( @@ -525,3 +606,159 @@ def test_calculate_memory_uris_missing_peer_id_falls_back_to_first_peer_when_sel assert uris == ["viking://user/support_bot/peers/web-visitor-bob/memories/preferences"] assert operation.memory_fields["user_id"] == "support_bot" assert operation.memory_fields["peer_id"] == "web-visitor-bob" + + + @patch("openviking.session.memory.memory_isolation_handler.generate_uri") + def test_peer_enabled_false_forces_self_scope_and_strips_peer_id(self, mock_generate_uri): + mock_generate_uri.side_effect = lambda **kwargs: ( + f"viking://user/{kwargs.get('user_space')}/memories/cases/demo" + ) + + ctx = create_ctx(user_id="support_bot") + messages = [create_message("user", peer_id="web-visitor-alice")] + extract_ctx = create_mock_extract_context(messages) + handler = MemoryIsolationHandler( + ctx, + extract_ctx, + allow_self=True, + allowed_peer_ids={"web-visitor-alice"}, + ) + + from openviking.session.memory.dataclass import MemoryTypeSchema, ResolvedOperation + + schema = MemoryTypeSchema( + memory_type="cases", + filename_template="demo.md", + directory="viking://user/{user_space}/memories/cases", + peer_enabled=False, + ) + operation = ResolvedOperation( + old_memory_file_content=None, + memory_fields={"case_name": "demo", "peer_id": "web-visitor-alice"}, + memory_type="cases", + uris=[], + ) + + uris = handler.calculate_memory_uris(schema, operation, extract_ctx) + + assert uris == ["viking://user/support_bot/memories/cases/demo"] + assert operation.memory_fields["user_id"] == "support_bot" + assert "peer_id" not in operation.memory_fields + + @patch("openviking.session.memory.memory_isolation_handler.generate_uri") + def test_peer_enabled_false_uses_self_scope_even_when_self_disabled(self, mock_generate_uri): + mock_generate_uri.side_effect = lambda **kwargs: ( + f"viking://user/{kwargs.get('user_space')}/memories/cases/demo" + ) + + ctx = create_ctx(user_id="support_bot") + messages = [create_message("user", peer_id="web-visitor-alice")] + extract_ctx = create_mock_extract_context(messages) + handler = MemoryIsolationHandler( + ctx, + extract_ctx, + allow_self=False, + allowed_peer_ids={"web-visitor-alice"}, + ) + + from openviking.session.memory.dataclass import MemoryTypeSchema, ResolvedOperation + + schema = MemoryTypeSchema( + memory_type="cases", + filename_template="demo.md", + directory="viking://user/{user_space}/memories/cases", + peer_enabled=False, + ) + operation = ResolvedOperation( + old_memory_file_content=None, + memory_fields={"case_name": "demo", "peer_id": "web-visitor-alice"}, + memory_type="cases", + uris=[], + ) + + uris = handler.calculate_memory_uris(schema, operation, extract_ctx) + + assert uris == ["viking://user/support_bot/memories/cases/demo"] + assert operation.memory_fields["user_id"] == "support_bot" + assert "peer_id" not in operation.memory_fields + + @patch("openviking.session.memory.memory_isolation_handler.generate_uri") + def test_peer_enabled_false_ignores_ranges_peer_targets(self, mock_generate_uri): + mock_generate_uri.side_effect = lambda **kwargs: ( + f"viking://user/{kwargs.get('user_space')}/memories/cases/demo" + ) + + ctx = create_ctx(user_id="support_bot") + messages = [ + create_message("user", "self turn"), + create_message("user", "peer turn", peer_id="web-visitor-alice"), + ] + extract_ctx = create_mock_extract_context(messages) + mock_range = MagicMock() + mock_range.elements = [messages] + extract_ctx.read_message_ranges.return_value = mock_range + handler = MemoryIsolationHandler( + ctx, + extract_ctx, + allow_self=True, + allowed_peer_ids={"web-visitor-alice"}, + ) + + from openviking.session.memory.dataclass import MemoryTypeSchema, ResolvedOperation + + schema = MemoryTypeSchema( + memory_type="cases", + filename_template="demo.md", + directory="viking://user/{user_space}/memories/cases", + peer_enabled=False, + ) + operation = ResolvedOperation( + old_memory_file_content=None, + memory_fields={"case_name": "demo", "ranges": "0-1"}, + memory_type="cases", + uris=[], + ) + + uris = handler.calculate_memory_uris(schema, operation, extract_ctx) + + assert uris == ["viking://user/support_bot/memories/cases/demo"] + assert operation.memory_fields["user_id"] == "support_bot" + assert "peer_id" not in operation.memory_fields + + @patch("openviking.session.memory.memory_isolation_handler.generate_uri") + def test_calculate_memory_uris_invalid_peer_id_falls_back_to_first_peer( + self, mock_generate_uri + ): + mock_generate_uri.side_effect = lambda **kwargs: ( + f"viking://user/{kwargs.get('user_space')}/memories/preferences" + ) + + ctx = create_ctx(user_id="support_bot") + messages = [create_message("user", peer_id="web-visitor-bob")] + extract_ctx = create_mock_extract_context(messages) + handler = MemoryIsolationHandler( + ctx, + extract_ctx, + allowed_peer_ids={"web-visitor-bob"}, + ) + + from openviking.session.memory.dataclass import MemoryTypeSchema, ResolvedOperation + + schema = MemoryTypeSchema( + memory_type="preferences", + filename_template="preferences.md", + directory="viking://user/{user_space}/memories", + ) + operation = ResolvedOperation( + old_memory_file_content=None, + memory_fields={"peer_id": "web/visitor/alice"}, + memory_type="preferences", + uris=[], + ) + + uris = handler.calculate_memory_uris(schema, operation, extract_ctx) + + assert uris == [ + "viking://user/support_bot/peers/web-visitor-bob/memories/preferences" + ] + assert operation.memory_fields["peer_id"] == "web-visitor-bob" diff --git a/tests/session/memory/test_memory_react.py b/tests/session/memory/test_memory_react.py index fd2856ae6d..4889c90f82 100644 --- a/tests/session/memory/test_memory_react.py +++ b/tests/session/memory/test_memory_react.py @@ -184,21 +184,21 @@ def test_get_allowed_directories_list(self, mock_vlm, mock_viking_fs): class TestExtractLoopFinalJsonRetry: def test_final_instruction_includes_schema_aware_empty_json(self): extract_loop = object.__new__(ExtractLoop) - extract_loop._expected_fields = ["delete_uris", "preferences", "tools"] + extract_loop._expected_fields = ["preferences", "tools"] instruction = extract_loop._build_final_operations_instruction() assert "ONLY a valid JSON object" in instruction - assert '"delete_uris": []' in instruction + assert '"delete_ids": []' in instruction assert '"preferences": []' in instruction assert '"tools": []' in instruction - def test_final_skeleton_always_includes_delete_uris(self): + def test_final_skeleton_always_includes_delete_ids(self): extract_loop = object.__new__(ExtractLoop) extract_loop._expected_fields = ["preferences"] assert extract_loop._build_final_operations_skeleton() == { - "delete_uris": [], + "delete_ids": [], "preferences": [], } @@ -251,8 +251,9 @@ async def prefetch(self): max_iterations=1, ) - with pytest.raises(RuntimeError, match="final response could not be parsed"): - await extract_loop.run() + result, _ = await extract_loop.run() + assert result.errors + assert "Final response could not be parsed" in result.errors[0] final_prompts = [ message["content"] @@ -262,5 +263,5 @@ async def prefetch(self): and "maximum number of tool call iterations" in message.get("content", "") ] assert final_prompts - assert '"delete_uris": []' in final_prompts[-1] + assert '"delete_ids": []' in final_prompts[-1] assert '"preferences": []' in final_prompts[-1] diff --git a/tests/session/memory/test_memory_react_system_prompt.py b/tests/session/memory/test_memory_react_system_prompt.py index 32b2c3c43b..a37bdc02e6 100644 --- a/tests/session/memory/test_memory_react_system_prompt.py +++ b/tests/session/memory/test_memory_react_system_prompt.py @@ -86,8 +86,8 @@ def test_instruction_includes_resource_uri_handling_for_user_scoped_resource_uri ) -class TestSkillToolCallExposure: - def test_assemble_conversation_includes_skill_tool_call(self): +class TestSessionConversationToolFiltering: + def test_session_conversation_omits_skill_tool_call(self): messages = [ Message( id="m1", @@ -111,10 +111,11 @@ def test_assemble_conversation_includes_skill_tool_call(self): conversation = provider._assemble_conversation(messages) - assert "[ToolCall]" in conversation - assert '"skill_name": "create_presentation"' in conversation + assert "ToolCall:" not in conversation + assert "create_presentation" not in conversation + assert "Running a skill." in conversation - def test_assemble_conversation_without_skill_tool_call_has_no_skill_name(self): + def test_session_conversation_omits_regular_tool_call(self): messages = [ Message( id="m1", @@ -137,9 +138,38 @@ def test_assemble_conversation_without_skill_tool_call_has_no_skill_name(self): conversation = provider._assemble_conversation(messages) - assert "[ToolCall]" in conversation - assert '"tool_name": "read"' in conversation - assert '"skill_name":' not in conversation + assert "ToolCall:" not in conversation + assert "tool_name=read" not in conversation + assert "Running a tool." in conversation + + def test_agent_provider_conversation_includes_tool_call_evidence(self): + from openviking.session.memory.agent_trajectory_context_provider import ( + AgentTrajectoryContextProvider, + ) + + messages = [ + Message( + id="m1", + role="assistant", + parts=[ + TextPart("Checking the reservation."), + ToolPart( + tool_id="tool_1", + tool_name="get_reservation_details", + tool_input={"reservation_id": "EHGLP3"}, + tool_output="available", + tool_status="completed", + ), + ], + ) + ] + provider = AgentTrajectoryContextProvider(messages=messages) + + conversation = provider._assemble_conversation(messages) + + assert "ToolCall: tool_name=get_reservation_details" in conversation + assert "input={'reservation_id': 'EHGLP3'}" in conversation + assert "output=available" in conversation def test_assemble_conversation_uses_peer_id_when_present(self): messages = [ @@ -202,6 +232,8 @@ def test_detect_language_prefers_user_text_over_assistant_text(self): assert provider._detect_language() == "zh-CN" + + async def test_prepare_extraction_messages_replaces_image_part_with_vlm_description(self): class FakeVisionVLM: def __init__(self): @@ -308,3 +340,9 @@ async def test_prepare_extraction_messages_does_not_replace_caller_message_list( assert len(messages) == 1 assert any(isinstance(part, ImagePart) for part in messages[0].parts) assert provider.messages is not messages + +def test_session_provider_empty_messages_still_uses_environment_fallback(monkeypatch): + monkeypatch.setenv("TZ", "Asia/Shanghai") + provider = SessionExtractContextProvider(messages=[]) + + assert provider.get_output_language() == "zh-CN" diff --git a/tests/session/memory/test_memory_read_tool_strip_links.py b/tests/session/memory/test_memory_read_tool_strip_links.py index 08626f5405..bbe56f9dfa 100644 --- a/tests/session/memory/test_memory_read_tool_strip_links.py +++ b/tests/session/memory/test_memory_read_tool_strip_links.py @@ -40,7 +40,7 @@ async def test_read_tool_strips_local_memory_links_from_llm_content(): uri="viking://user/default/memories/experiences/test.md", ) - assert result["content"] == "1 | Gina values emotional support with Jon." + assert result["content"] == "1\tGina values emotional support with Jon." @pytest.mark.asyncio @@ -69,4 +69,4 @@ def tracking_plain_content(self): ) assert called is True - assert result["content"] == "1 | Gina values emotional support with Jon." + assert result["content"] == "1\tGina values emotional support with Jon." diff --git a/tests/session/memory/test_memory_updater.py b/tests/session/memory/test_memory_updater.py index 9e1c034ca5..4c5de70859 100644 --- a/tests/session/memory/test_memory_updater.py +++ b/tests/session/memory/test_memory_updater.py @@ -14,6 +14,7 @@ from openviking.session.memory.dataclass import ( MemoryField, MemoryFile, + MemoryOperationSource, MemoryTypeSchema, ResolvedOperation, ResolvedOperations, @@ -102,6 +103,18 @@ def test_extract_context_initializes_page_id_map(self): page_id = extract_context.page_id_map.get_page_id("viking://user/a/memories/profile.md") assert page_id == 1 + def test_extract_context_can_disable_long_text_message_split(self): + text = "第一句很长很长很长很长很长很长很长很长很长很长很长。" * 8 + messages = [Message(id="1", role="user", parts=[TextPart(text=text)])] + + split_context = ExtractContext(messages) + unsplit_context = ExtractContext(messages, split_long_text_messages=False) + + assert len(split_context.messages) > 1 + assert len(unsplit_context.messages) == 1 + assert unsplit_context.messages[0] is messages[0] + assert unsplit_context.chunk_meta == {} + def test_extract_context_resource_event_content_hides_add_resource_fields(self): resource_uri = "viking://resources/images/2026/06/12/yueqian_jpeg" extract_context = ExtractContext( @@ -573,6 +586,146 @@ async def mock_apply_delete(uri, ctx): call.args[0] for call in mock_viking_fs.read_file.await_args_list ] + @pytest.mark.asyncio + async def test_apply_operations_skips_case_only_delete_conflicting_with_upsert(self): + written_uri = "viking://user/conv-26/memories/entities/person/melanie.md" + deleted_uri = "viking://user/conv-26/memories/entities/person/Melanie.md" + + schema = MemoryTypeSchema( + memory_type="entities", + description="entity memory", + directory="viking://user/{{ user_space }}/memories/entities", + filename_template="{{ category }}/{{ name }}.md", + fields=[], + overview_template="overview", + ) + registry = MagicMock() + registry.get.return_value = schema + + updater = MemoryUpdater(registry=registry) + updater._get_viking_fs = MagicMock(return_value=MagicMock()) + updater._apply_upsert = AsyncMock(return_value=None) + updater._apply_delete = AsyncMock() + updater._vectorize_memories = AsyncMock() + updater.generate_overview = AsyncMock() + + resolved = ResolvedOperations( + upsert_operations=[ + ResolvedOperation( + memory_fields={"category": "person", "name": "melanie"}, + memory_type="entities", + uris=[written_uri], + ) + ], + delete_file_contents=[ + MemoryFile(uri=deleted_uri, extra_fields={"memory_type": "entities"}) + ], + errors=[], + ) + ctx = RequestContext(user=UserIdentifier("acme", "conv-26"), role=Role.USER) + + result = await updater.apply_operations(operations=resolved, ctx=ctx) + + assert result.written_uris == [written_uri] + assert result.deleted_uris == [] + updater._apply_delete.assert_not_awaited() + + @pytest.mark.asyncio + async def test_apply_operations_remaps_deleted_links_to_replacement(self): + deleted_uri = "viking://user/u/memories/preferences/Evan/hobby_preferences.md" + replacement_uri = "viking://user/u/memories/preferences/Evan/hobbies.md" + profile_uri = "viking://user/u/memories/profile.md" + + schema = MemoryTypeSchema( + memory_type="preferences", + description="preference memory", + directory="viking://user/{{ user_space }}/memories/preferences", + filename_template="{{ name }}.md", + fields=[], + overview_template="overview", + ) + registry = MagicMock() + registry.get.return_value = schema + + deleted_file = MemoryFile( + uri=deleted_uri, + content="old hobby preferences", + memory_type="preferences", + links=[ + { + "from_uri": deleted_uri, + "to_uri": profile_uri, + "link_type": "related_to", + "weight": 0.8, + "match_text": "hobby", + "description": "old link", + } + ], + ) + replacement_file = MemoryFile( + uri=replacement_uri, + content="new hobbies", + memory_type="preferences", + ) + profile_file = MemoryFile( + uri=profile_uri, + content="profile", + memory_type="profile", + ) + files = { + deleted_uri: MemoryFileUtils.write(deleted_file), + replacement_uri: MemoryFileUtils.write(replacement_file), + profile_uri: MemoryFileUtils.write(profile_file), + } + + mock_viking_fs = MagicMock() + mock_viking_fs.read_file = AsyncMock(side_effect=lambda uri, ctx=None: files[uri]) + + async def write_file(uri, content, ctx=None): + files[uri] = content + + mock_viking_fs.write_file = AsyncMock(side_effect=write_file) + updater = MemoryUpdater(registry=registry) + updater._get_viking_fs = MagicMock(return_value=mock_viking_fs) + updater._apply_upsert = AsyncMock(return_value=None) + updater._apply_delete = AsyncMock() + updater._vectorize_memories = AsyncMock() + updater.generate_overview = AsyncMock() + + resolved = ResolvedOperations( + upsert_operations=[ + ResolvedOperation( + memory_fields={"name": "hobbies"}, + memory_type="preferences", + uris=[replacement_uri], + ) + ], + delete_file_contents=[deleted_file], + errors=[], + resolved_links=[ + StoredLink( + from_uri=deleted_uri, + to_uri=profile_uri, + link_type="related_to", + weight=0.9, + match_text="hobby", + description="in-flight link", + ) + ], + delete_replacements={deleted_uri: replacement_uri}, + ) + + ctx = RequestContext(user=UserIdentifier("acme", "alice"), role=Role.USER) + + result = await updater.apply_operations(operations=resolved, ctx=ctx) + + assert result.written_uris == [replacement_uri] + assert result.deleted_uris == [deleted_uri] + assert resolved.resolved_links[0].from_uri == replacement_uri + profile = MemoryFileUtils.read(files[profile_uri], uri=profile_uri) + assert profile.backlinks[0]["from_uri"] == replacement_uri + assert profile.backlinks[0]["to_uri"] == profile_uri + @pytest.mark.asyncio async def test_apply_operations_routes_backlinks_to_matching_uri_only(self): caroline_uri = ( @@ -867,6 +1020,33 @@ def _make_updater_with_registry(self): registry.register(schema) return MemoryUpdater(registry=registry) + @pytest.mark.asyncio + async def test_apply_upsert_persists_last_update_trace_id(self): + updater = self._make_updater_with_registry() + mock_viking_fs = MagicMock() + mock_viking_fs.read_file = AsyncMock(side_effect=FileNotFoundError("missing")) + written_content = None + + async def mock_write_file(uri, content, **kwargs): + nonlocal written_content + written_content = content + + mock_viking_fs.write_file = mock_write_file + updater._get_viking_fs = MagicMock(return_value=mock_viking_fs) + + op = ResolvedOperation( + memory_fields={"content": "Line 1"}, + memory_type="test", + uris=["viking://test/test.md"], + source=MemoryOperationSource(extraction_id="extract_1", trace_id="trace_1"), + ) + await updater._apply_upsert(op, MagicMock()) + + assert written_content is not None + result = MemoryFileUtils.read(written_content) + assert result.extra_fields["source_extraction_id"] == "extract_1" + assert result.extra_fields["last_update_trace_id"] == "trace_1" + @pytest.mark.asyncio async def test_apply_edit_with_str_patch_instance(self): """Test _apply_edit with StrPatch instance.""" @@ -1086,6 +1266,58 @@ async def mock_write_file(uri, content, **kwargs): final_content = store[uri] parsed = parse_memory_file_with_fields(final_content) assert parsed["content"] == "Step B" + assert parsed["version"] == 2 + + @pytest.mark.asyncio + async def test_apply_upsert_strips_user_id_and_sets_version(self): + memory_type = "notes" + uri = "viking://user/alice/memories/notes.md" + schema = MemoryTypeSchema( + memory_type=memory_type, + description="notes", + fields=[ + MemoryField( + name="content", + field_type=FieldType.STRING, + merge_op=MergeOp.PATCH, + ), + ], + ) + registry = MemoryTypeRegistry() + registry.register(schema) + + store: dict[str, str] = {} + mock_viking_fs = MagicMock() + + async def mock_read_file(uri, **kwargs): + return store.get(uri) + + async def mock_write_file(uri, content, **kwargs): + store[uri] = content + + mock_viking_fs.read_file = mock_read_file + mock_viking_fs.write_file = mock_write_file + + updater = MemoryUpdater(registry=registry) + updater._get_viking_fs = MagicMock(return_value=mock_viking_fs) + + op = ResolvedOperation( + old_memory_file_content=None, + memory_fields={ + "content": "Step A", + "user_id": "alice", + "user_ids": ["alice", "bob"], + }, + memory_type=memory_type, + uris=[uri], + ) + await updater._apply_upsert(op, MagicMock()) + + parsed = parse_memory_file_with_fields(store[uri]) + assert parsed["content"] == "Step A" + assert parsed["version"] == 1 + assert "user_id" not in parsed + assert "user_ids" not in parsed @pytest.mark.asyncio async def test_apply_upsert_skips_failed_field_and_keeps_other_fields(self, monkeypatch): diff --git a/tests/session/memory/test_memory_utils.py b/tests/session/memory/test_memory_utils.py index bfa5ad859c..5ae21225ad 100644 --- a/tests/session/memory/test_memory_utils.py +++ b/tests/session/memory/test_memory_utils.py @@ -354,8 +354,29 @@ def test_write_preserves_memory_type_in_memory_fields_comment(self): assert parsed["memory_type"] == "preferences" assert parsed["topic"] == "code_style" + assert parsed["version"] == 1 assert parsed["content"] == "Prefers concise responses." + def test_write_strips_user_identity_fields_from_memory_fields_comment(self): + memory_file = MemoryFile( + uri="viking://user/alice/memories/preferences/code_style.md", + memory_type="preferences", + content="Prefers concise responses.", + extra_fields={ + "topic": "code_style", + "user_id": "alice", + "user_ids": ["alice", "bob"], + }, + ) + + written = MemoryFileUtils.write(memory_file) + parsed = parse_memory_file_with_fields(written) + + assert "user_id" not in parsed + assert "user_ids" not in parsed + assert parsed["version"] == 1 + assert parsed["topic"] == "code_style" + def test_read_preserves_markdown_links_in_content(self): raw_content = """2023-08-22 ChatLog\n\n[Calvin]: Worked with [Frank Ocean](../../../../entities/personal/calvin.md).\n\n""" @@ -380,3 +401,25 @@ def test_memory_file_plain_content_strips_markdown_links(self): ) assert memory_file.plain_content() == "Worked with Frank Ocean." + + def test_write_does_not_put_uri_in_memory_fields_metadata(self): + memory_file = MemoryFile( + uri="viking://user/default/memories/experiences/source.md", + memory_type="experiences", + content="Source content mentions Target.", + extra_fields={"_uri": "viking://stale/should-not-persist"}, + links=[ + { + "from_uri": "viking://user/default/memories/experiences/source.md", + "to_uri": "viking://user/default/memories/trajectories/target.md", + "link_type": "derived_from", + "match_text": "Target", + } + ], + ) + + written = MemoryFileUtils.write(memory_file) + parsed = parse_memory_file_with_fields(written) + + assert "_uri" not in parsed + assert "[Target](../trajectories/target.md)" in written diff --git a/tests/session/memory/test_patch_merge_context_provider.py b/tests/session/memory/test_patch_merge_context_provider.py new file mode 100644 index 0000000000..39c98fdb28 --- /dev/null +++ b/tests/session/memory/test_patch_merge_context_provider.py @@ -0,0 +1,407 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 + +from __future__ import annotations + +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from openviking.session.memory.dataclass import MemoryFile, MemoryTypeSchema +from openviking.session.memory.patch_merge_context_provider import ( + PatchMergeContextProvider, + PatchMergePatch, +) + + +def _memory_file( + *, + name: str, + uri: str | None, + content: str, + memory_type: str = "experiences", +) -> MemoryFile: + return MemoryFile( + uri=uri, + content=content, + memory_type=memory_type, + extra_fields={ + "memory_type": memory_type, + "experience_name": name, + "status": "production", + }, + ) + + +@pytest.mark.asyncio +async def test_patch_merge_context_provider_prefetch_reads_originals_and_renders_patch(): + uri = "viking://user/u/memories/experiences/booking.md" + provider = PatchMergeContextProvider( + memory_type="experiences", + required_file_uris=[uri], + patches=[ + PatchMergePatch( + before_file=_memory_file(name="booking", uri=uri, content="old line\nkeep line"), + after_file=_memory_file(name="booking", uri=uri, content="new line\nkeep line"), + ) + ], + ) + provider.read_file = AsyncMock( + return_value={ + "memory_type": "experiences", + "experience_name": "booking", + "content": "1\told line\n2\tkeep line", + } + ) + + messages = await provider.prefetch() + + assert provider.get_tools() == [] + assert provider.read_file.await_count == 1 + read_message = json.loads(messages[0]["content"]) + assert read_message["tool_call_name"] == "read" + assert read_message["args"] == {"uri": "viking://user/u/memories/experiences/booking.md"} + assert read_message["result"]["experience_name"] == "booking" + assert messages[1]["role"] == "user" + assert messages[1]["content"].startswith("# Memory File Patches") + assert "Patch 1" in messages[1]["content"] + # Patch headers should not include target_uri/target_name/memory_type + assert "target_uri:" not in messages[1]["content"] + assert "target_name:" not in messages[1]["content"] + assert " content:" in messages[1]["content"] + assert "--- content.before" not in messages[1]["content"] + assert "+++ content.after" not in messages[1]["content"] + assert " -old line" in messages[1]["content"] + assert " +new line" in messages[1]["content"] + assert " keep line" in messages[1]["content"] # n=1 context line + assert " status:" not in messages[1]["content"] + + +@pytest.mark.asyncio +async def test_patch_merge_context_provider_prefetch_searches_and_reads_extra_candidates(): + schema = MemoryTypeSchema( + memory_type="experiences", + description="Experiences", + directory="viking://user/{{ user_space }}/memories/experiences", + filename_template="{{ experience_name }}.md", + fields=[], + ) + provider = PatchMergeContextProvider( + memory_type="experiences", + required_file_uris=["viking://user/u/memories/experiences/book.md"], + patches=[ + PatchMergePatch( + before_file=None, + after_file=_memory_file( + name="books", + uri="viking://user/u/memories/experiences/books.md", + content="用户喜欢阅读科幻书籍,尤其是太空歌剧。", + ), + ) + ], + ) + provider._registry = SimpleNamespace(get=lambda name: schema if name == "experiences" else None) + provider._ctx = SimpleNamespace(user=SimpleNamespace(user_id="u")) + provider.search_files = AsyncMock( + return_value=[ + "viking://user/u/memories/experiences/book.md", + *[f"viking://user/u/memories/experiences/candidate_{idx}.md" for idx in range(10)], + ] + ) + provider.read_file = AsyncMock( + return_value={ + "memory_type": "experiences", + "experience_name": "candidate", + "content": "candidate content", + } + ) + + messages = await provider.prefetch() + + provider.search_files.assert_awaited_once() + _, search_kwargs = provider.search_files.await_args + assert search_kwargs["search_uris"] == ["viking://user/u/memories/experiences"] + assert search_kwargs["limit"] == 10 + assert provider.read_file.await_count == 6 + read_uris = [call.args[0] for call in provider.read_file.await_args_list] + assert read_uris[0] == "viking://user/u/memories/experiences/book.md" + assert "viking://user/u/memories/experiences/candidate_0.md" in read_uris + assert "viking://user/u/memories/experiences/candidate_4.md" in read_uris + assert "viking://user/u/memories/experiences/candidate_5.md" not in read_uris + assert messages[-1]["content"].startswith("# Memory File Patches") + + +@pytest.mark.asyncio +async def test_patch_merge_context_provider_caps_extra_candidate_reads_at_ten(): + schema = MemoryTypeSchema( + memory_type="experiences", + description="Experiences", + directory="viking://user/{{ user_space }}/memories/experiences", + filename_template="{{ experience_name }}.md", + fields=[], + ) + required_uris = [f"viking://user/u/memories/experiences/required_{idx}.md" for idx in range(12)] + provider = PatchMergeContextProvider( + memory_type="experiences", + required_file_uris=required_uris, + patches=[ + PatchMergePatch( + before_file=None, + after_file=_memory_file( + name="books", + uri="viking://user/u/memories/experiences/books.md", + content="用户喜欢阅读科幻书籍,尤其是太空歌剧。", + ), + ) + ], + ) + provider._registry = SimpleNamespace(get=lambda name: schema if name == "experiences" else None) + provider._ctx = SimpleNamespace(user=SimpleNamespace(user_id="u")) + provider.search_files = AsyncMock( + return_value=[ + *required_uris, + *[f"viking://user/u/memories/experiences/candidate_{idx}.md" for idx in range(20)], + ] + ) + provider.read_file = AsyncMock( + return_value={ + "memory_type": "experiences", + "experience_name": "candidate", + "content": "candidate content", + } + ) + + await provider.prefetch() + + _, search_kwargs = provider.search_files.await_args + assert search_kwargs["limit"] == 20 + assert provider.read_file.await_count == 22 + read_uris = [call.args[0] for call in provider.read_file.await_args_list] + assert required_uris[-1] in read_uris + assert "viking://user/u/memories/experiences/candidate_9.md" in read_uris + assert "viking://user/u/memories/experiences/candidate_10.md" not in read_uris + + +@pytest.mark.asyncio +async def test_patch_merge_context_provider_renders_compact_patch_metadata(): + provider = PatchMergeContextProvider( + memory_type="experiences", + required_file_uris=[], + patches=[ + PatchMergePatch( + before_file=None, + after_file=_memory_file( + name="new_booking", + uri=None, + content="created line", + ), + metadata={ + "base_version": 3, + "rationale": "useful reason", + "confidence": 0.9, + "links": [{"to_uri": "viking://user/u/memories/trajectories/t.md"}], + "memory_fields": {"content": "created line"}, + "uris": ["viking://user/u/memories/experiences/new_booking.md"], + "gradient_metadata": { + "trajectory_outcome": "success", + "rubric_passed": True, + "training_category": "tau2:airline:train:1", + "memory_fields": {"content": "duplicated"}, + }, + }, + ) + ], + ) + + messages = await provider.prefetch() + content = messages[0]["content"] + + assert 'meta: {"confidence": 0.9}' in content + assert "base_version" not in content + assert "rationale" not in content + assert "trajectory_outcome" not in content + assert "gradient_metadata" not in content + assert "links" not in content + assert "memory_fields" not in content + assert "duplicated" not in content + + +@pytest.mark.asyncio +async def test_patch_merge_context_provider_hides_last_update_trace_id_from_patch_diff(): + provider = PatchMergeContextProvider( + memory_type="experiences", + required_file_uris=[], + patches=[ + PatchMergePatch( + before_file=MemoryFile( + uri="viking://user/u/memories/experiences/booking.md", + content="same content", + memory_type="experiences", + extra_fields={ + "memory_type": "experiences", + "experience_name": "booking", + "last_update_trace_id": "trace_old", + }, + ), + after_file=MemoryFile( + uri="viking://user/u/memories/experiences/booking.md", + content="same content", + memory_type="experiences", + extra_fields={ + "memory_type": "experiences", + "experience_name": "booking", + "last_update_trace_id": "trace_new", + }, + ), + ) + ], + ) + + messages = await provider.prefetch() + content = messages[0]["content"] + + assert "last_update_trace_id" not in content + assert "trace_old" not in content + assert "trace_new" not in content + assert "(no changes)" in content + + +@pytest.mark.asyncio +async def test_patch_merge_context_provider_renders_create_patch_from_dev_null(): + provider = PatchMergeContextProvider( + memory_type="experiences", + required_file_uris=[], + patches=[ + PatchMergePatch( + before_file=None, + after_file=_memory_file(name="new_booking", uri=None, content="created line"), + ) + ], + ) + + messages = await provider.prefetch() + + assert len(messages) == 1 + assert "Patch 1" in messages[0]["content"] + # Patch headers should not include target_name/target_uri/memory_type + assert "target_name:" not in messages[0]["content"] + assert "target_uri:" not in messages[0]["content"] + # Field diffs may include memory_type field changes (that's expected) + assert " content:" in messages[0]["content"] + assert "--- content.before" not in messages[0]["content"] + assert "+++ content.after" not in messages[0]["content"] + assert " +created line" in messages[0]["content"] + + +def test_patch_merge_context_provider_get_memory_schema_single_type(monkeypatch): + schema = MemoryTypeSchema( + memory_type="experiences", + description="Experiences", + directory="viking://user/{{ user_space }}/memories/experiences", + filename_template="{{ experience_name }}.md", + fields=[], + ) + provider = PatchMergeContextProvider( + memory_type="experiences", + required_file_uris=[], + patches=[], + ) + provider._registry = SimpleNamespace(get=lambda name: schema if name == "experiences" else None) + + assert provider.get_memory_schemas(ctx=None) == [schema] + + +def test_patch_merge_context_provider_get_memory_schema_raises_for_missing_type(): + provider = PatchMergeContextProvider( + memory_type="missing", + required_file_uris=[], + patches=[], + ) + provider._registry = SimpleNamespace(get=lambda name: None) + + with pytest.raises(ValueError, match="Memory schema not found or disabled: missing"): + provider.get_memory_schemas(ctx=None) + + +def test_patch_merge_context_provider_instruction_mentions_path_field_normalization(): + provider = PatchMergeContextProvider( + memory_type="entities", + required_file_uris=[], + patches=[], + ) + + instruction = provider.instruction() + + assert "independent extraction patch proposals" in instruction + assert "merge duplicate/overlapping\nmemories into one canonical file patch" in instruction + assert "directory/filename fields" in instruction + assert "schema identifiers" in instruction + assert "book not books" in instruction + assert "Chinese" in instruction + assert "书籍 not 书/图书" in instruction + assert "put it in delete_ids" in instruction + + +def test_patch_merge_context_provider_detects_language_from_patch_content(monkeypatch): + monkeypatch.setenv("TZ", "Asia/Shanghai") + provider = PatchMergeContextProvider( + memory_type="preferences", + required_file_uris=[], + patches=[ + PatchMergePatch( + before_file=None, + after_file=MemoryFile( + uri=None, + content="User prefers concise implementation and minimal fallback logic.", + memory_type="preferences", + extra_fields={ + "memory_type": "preferences", + "user": "alice", + "topic": "code_style", + }, + ), + ) + ], + ) + + assert provider.get_output_language() == "en" + assert "All memory content must be written in en." in provider.instruction() + + +def test_patch_merge_context_provider_empty_patches_fallback_to_english(monkeypatch): + monkeypatch.setenv("TZ", "Asia/Shanghai") + provider = PatchMergeContextProvider( + memory_type="preferences", + required_file_uris=[], + patches=[], + ) + + assert provider.get_output_language() == "en" + + +def test_patch_merge_context_provider_ignores_before_file_language(monkeypatch): + monkeypatch.setenv("TZ", "Asia/Shanghai") + provider = PatchMergeContextProvider( + memory_type="preferences", + required_file_uris=[], + patches=[ + PatchMergePatch( + before_file=MemoryFile( + uri="viking://user/u/memories/preferences/old.md", + content="用户偏好简洁实现。", + memory_type="preferences", + extra_fields={"memory_type": "preferences", "topic": "代码风格"}, + ), + after_file=MemoryFile( + uri="viking://user/u/memories/preferences/old.md", + content="User prefers concise implementation.", + memory_type="preferences", + extra_fields={"memory_type": "preferences", "topic": "code_style"}, + ), + ) + ], + ) + + assert provider.get_output_language() == "en" diff --git a/tests/session/memory/test_schema_models.py b/tests/session/memory/test_schema_models.py index 4fd0dfbf1c..1b280d1f3c 100644 --- a/tests/session/memory/test_schema_models.py +++ b/tests/session/memory/test_schema_models.py @@ -81,14 +81,33 @@ def registry_with_sample(self, sample_memory_type): @pytest.fixture def real_registry(self): """Create a registry with real schemas.""" - schemas_dir = ( - Path(__file__).parent.parent.parent.parent - / "openviking" - / "prompts" - / "templates" - / "memory" + return create_default_registry() + + + def test_peer_enabled_false_omits_peer_id_field(self): + memory_type = MemoryTypeSchema( + memory_type="cases", + description="Case memory", + fields=[ + MemoryField( + name="case_name", + field_type=FieldType.STRING, + description="Case name", + merge_op=MergeOp.IMMUTABLE, + ) + ], + filename_template="{{ case_name }}.md", + directory="viking://user/{{ user_space }}/memories/cases", + peer_enabled=False, + ) + role_scope = type("RoleScope", (), {"peer_ids": ["web-visitor-alice"]})() + + model = SchemaModelGenerator([memory_type]).create_flat_data_model( + memory_type, + role_scope=role_scope, ) - return create_default_registry(str(schemas_dir)) + + assert "peer_id" not in model.model_fields def test_render_description_template_with_language(self): memory_type = MemoryTypeSchema( @@ -159,23 +178,13 @@ def test_create_flat_data_model(self, sample_memory_type, registry_with_sample): # Check model name assert model.__name__ == "TestTypeData" - # Check model has the memory_type field - assert "memory_type" in model.model_fields - # memory_type is a required field with literal type + # memory_type is represented by the top-level structured output field. + assert "memory_type" not in model.model_fields # Check business fields assert "field1" in model.model_fields assert "field2" in model.model_fields - # Check metadata fields are present - assert "uri" in model.model_fields - assert "name" in model.model_fields - assert "abstract" in model.model_fields - assert "overview" in model.model_fields - assert "content" in model.model_fields - assert "tags" in model.model_fields - assert "created_at" in model.model_fields - assert "updated_at" in model.model_fields def test_page_id_field_is_emitted_before_mutable_content(self, registry_with_sample): """page_id should appear before mutable fields so the model anchors target page first.""" @@ -325,14 +334,14 @@ def test_get_llm_json_schema(self, real_registry): assert "$defs" in json_schema or "definitions" in json_schema assert "properties" in json_schema - # Check it includes operations - assert "write_uris" in json_schema["properties"] - assert "edit_uris" in json_schema["properties"] - assert "delete_uris" in json_schema["properties"] + # Check it includes delete operations and per-memory-type operation fields + assert "delete_ids" in json_schema["properties"] + assert "profile" in json_schema["properties"] - # Check delete_uris is an array of strings - delete_props = json_schema["properties"]["delete_uris"] - assert delete_props.get("items", {}).get("type") == "string" + # Check delete_ids is an array of objects + delete_props = json_schema["properties"]["delete_ids"] + delete_items = delete_props.get("items", {}) + assert delete_items.get("$ref") or delete_items.get("type") == "object" def test_get_memory_data_json_schema(self, real_registry): """Test getting just the MemoryData JSON schema.""" @@ -393,8 +402,7 @@ def test_dynamic_new_schema(self): # Verify the model has the custom field assert "custom_field" in model.model_fields - assert "memory_type" in model.model_fields - assert "uri" in model.model_fields + assert "memory_type" not in model.model_fields class TestWikiLink: @@ -424,14 +432,7 @@ class TestIntegration: def test_end_to_end_model_generation_and_validation(self): """Test end-to-end: load schemas, generate models, validate data.""" - schemas_dir = ( - Path(__file__).parent.parent.parent.parent - / "openviking" - / "prompts" - / "templates" - / "memory" - ) - registry = create_default_registry(str(schemas_dir)) + registry = create_default_registry() # Create generator generator = SchemaModelGenerator(registry) diff --git a/tests/session/memory/test_streaming_memory_updater.py b/tests/session/memory/test_streaming_memory_updater.py new file mode 100644 index 0000000000..a5a1e54a55 --- /dev/null +++ b/tests/session/memory/test_streaming_memory_updater.py @@ -0,0 +1,1052 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 + +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock + +import pytest + +from openviking.message import Message, TextPart +from openviking.server.identity import RequestContext, Role +from openviking.session.memory.dataclass import ( + MemoryField, + MemoryFile, + MemoryOperationSource, + MemoryTypeSchema, + ResolvedOperation, + ResolvedOperations, + StoredLink, +) +from openviking.session.memory.memory_type_registry import MemoryTypeRegistry +from openviking.session.memory.memory_updater import ExtractContext, MemoryUpdateResult +from openviking.session.memory.merge_op.base import FieldType, MergeOp, SearchReplaceBlock, StrPatch +from openviking.session.memory.streaming_memory_updater import ( + MemoryMergeGroupKey, + MemoryUpdateRequest, + StreamingMemoryUpdater, + StreamingMemoryUpdaterConfig, + StreamingMemoryUpdateResult, + classify_memory_merge_mode, + enforce_merge_group_peer_id, + merge_one_memory_type_operations, + operation_to_patch, + render_operation_after_file_content, + split_request_by_merge_group, +) +from openviking.session.memory.utils.memory_file_utils import MemoryFileUtils +from openviking_cli.session.user_id import UserIdentifier + + +class InMemoryVikingFS: + def __init__(self, files: dict[str, str] | None = None): + self.files = dict(files or {}) + self.writes = [] + + async def ls(self, uri: str, output: str = "original", ctx=None): + del output, ctx + prefix = uri.rstrip("/") + "/" + return [ + {"name": path.removeprefix(prefix), "uri": path, "isDir": False} + for path in sorted(self.files) + if path.startswith(prefix) and "/" not in path.removeprefix(prefix) + ] + + async def read_file(self, uri: str, ctx=None): + uri = _canonical_user_uri(uri, ctx) + if uri not in self.files: + raise FileNotFoundError(uri) + return self.files[uri] + + async def write_file(self, uri: str, content: str, ctx=None): + uri = _canonical_user_uri(uri, ctx) + self.files[uri] = content + self.writes.append((uri, content, ctx)) + + async def rm(self, uri: str, recursive: bool = False, ctx=None, lock_handle=None): + del recursive, lock_handle + uri = _canonical_user_uri(uri, ctx) + self.files.pop(uri, None) + + +def _canonical_user_uri(uri: str, ctx=None) -> str: + if not uri.startswith("viking://user/memories/"): + return uri + user_id = getattr(getattr(ctx, "user", None), "user_id", None) or "u" + return uri.replace("viking://user/memories/", f"viking://user/{user_id}/memories/", 1) + + +def _ctx() -> RequestContext: + return RequestContext(user=UserIdentifier.the_default_user("u"), role=Role.ROOT) + + +def _registry() -> MemoryTypeRegistry: + registry = MemoryTypeRegistry(load_schemas=False) + registry.register( + MemoryTypeSchema( + memory_type="cases", + description="case memory", + directory="viking://user/{{ user_space }}/memories/cases", + filename_template="{{ case_name }}.md", + operation_mode="add_only", + peer_enabled=False, + fields=[ + MemoryField( + name="case_name", + field_type=FieldType.STRING, + merge_op=MergeOp.IMMUTABLE, + ), + MemoryField( + name="task_signature", + field_type=FieldType.STRING, + merge_op=MergeOp.IMMUTABLE, + ), + MemoryField( + name="input", + field_type=FieldType.STRING, + merge_op=MergeOp.IMMUTABLE, + ), + MemoryField( + name="rubric", + field_type=FieldType.STRING, + merge_op=MergeOp.IMMUTABLE, + ), + ], + ) + ) + registry.register( + MemoryTypeSchema( + memory_type="notes", + description="note memory", + directory="viking://user/{{ user_space }}/memories/notes", + filename_template="{{ note_name }}.md", + operation_mode="upsert", + fields=[ + MemoryField( + name="note_name", + field_type=FieldType.STRING, + merge_op=MergeOp.IMMUTABLE, + ), + MemoryField( + name="content", + field_type=FieldType.STRING, + merge_op=MergeOp.PATCH, + ), + ], + ) + ) + return registry + + +def _case_op(name: str) -> ResolvedOperation: + return ResolvedOperation( + old_memory_file_content=None, + memory_type="cases", + uris=[f"viking://user/u/memories/cases/{name}.md"], + memory_fields={ + "case_name": name, + "task_signature": f"{name} signature", + "input": '{"summary":"case input"}', + "rubric": '{"criteria":[{"name":"done","description":"done","required":true,"weight":1.0}]}', + }, + ) + + +def _note_op(name: str) -> ResolvedOperation: + return ResolvedOperation( + old_memory_file_content=None, + memory_type="notes", + uris=[f"viking://user/u/memories/notes/{name}.md"], + memory_fields={ + "note_name": name, + "content": f"{name} content", + }, + ) + + +def _note_op_with_source(name: str, extraction_id: str) -> ResolvedOperation: + op = _note_op(name) + op.memory_fields["source_extraction_id"] = extraction_id + return op + + +def _peer_note_op(name: str, peer_id: str) -> ResolvedOperation: + op = _note_op(name) + op.memory_fields["peer_id"] = peer_id + op.uris = [f"viking://user/u/peers/{peer_id}/memories/notes/{name}.md"] + return op + + +def test_operation_to_patch_omits_raw_operation_metadata(): + schema = _registry().get("notes") + old_file = MemoryFile( + uri="viking://user/u/memories/notes/note.md", + content="old content", + memory_type="notes", + extra_fields={"note_name": "note"}, + ) + op = ResolvedOperation( + old_memory_file_content=old_file, + memory_type="notes", + uris=["viking://user/u/memories/notes/note.md"], + memory_fields={ + "note_name": "note", + "content": StrPatch( + blocks=[SearchReplaceBlock(search="old content", replace="new content")] + ), + }, + ) + + patch = operation_to_patch(op, schema=schema, extract_context=ExtractContext([])) + + assert patch.metadata == {} + assert patch.after_file.content == "new content" + + +def test_operation_to_patch_raises_when_after_file_preview_rendering_fails(monkeypatch): + schema = _registry().get("notes") + op = _note_op("note_render_failure") + + def fail_write(*args, **kwargs): + raise RuntimeError("template render failed") + + monkeypatch.setattr( + "openviking.session.memory.streaming_memory_updater.MemoryFileUtils.write", + fail_write, + ) + + with pytest.raises(RuntimeError, match="template render failed"): + operation_to_patch(op, schema=schema, extract_context=ExtractContext([])) + + +def test_operation_to_patch_skips_failed_field_preview_update(): + schema = MemoryTypeSchema( + memory_type="notes", + description="note memory", + directory="viking://user/{{ user_space }}/memories/notes", + filename_template="{{ note_name }}.md", + operation_mode="upsert", + fields=[ + MemoryField( + name="note_name", + field_type=FieldType.STRING, + merge_op=MergeOp.IMMUTABLE, + ), + MemoryField( + name="content", + field_type=FieldType.STRING, + merge_op=MergeOp.PATCH, + ), + MemoryField( + name="summary", + field_type=FieldType.STRING, + merge_op=MergeOp.PATCH, + ), + ], + ) + old_file = MemoryFile( + uri="viking://user/u/memories/notes/note.md", + content="old content", + memory_type="notes", + extra_fields={ + "note_name": "note", + "summary": "old summary", + }, + ) + op = ResolvedOperation( + old_memory_file_content=old_file, + memory_type="notes", + uris=["viking://user/u/memories/notes/note.md"], + memory_fields={ + "note_name": "note", + "content": StrPatch( + blocks=[SearchReplaceBlock(search="old content", replace="new content")] + ), + "summary": StrPatch( + blocks=[SearchReplaceBlock(search="missing summary", replace="new summary")] + ), + }, + ) + + patch = operation_to_patch(op, schema=schema, extract_context=ExtractContext([])) + + assert patch.after_file.content == "new content" + assert patch.after_file.extra_fields["summary"] == "old summary" + assert isinstance(op.memory_fields["summary"], StrPatch) + + +@pytest.mark.asyncio +async def test_streaming_memory_updater_submit_applies_fast_path(monkeypatch): + fs = InMemoryVikingFS({}) + fs.search = AsyncMock(return_value=[]) + monkeypatch.setattr( + "openviking.session.memory.streaming_memory_updater.get_viking_fs", + lambda: fs, + ) + monkeypatch.setattr( + "openviking.session.memory.memory_updater.get_viking_fs", + lambda: fs, + ) + + updater = StreamingMemoryUpdater( + registry=_registry(), + config=StreamingMemoryUpdaterConfig( + max_operations_per_update=8, + max_wait_seconds=0.01, + timer_check_interval_seconds=0.01, + ), + ) + result = await updater.submit( + MemoryUpdateRequest( + operations=ResolvedOperations( + upsert_operations=[_case_op("重复预订处理")], + delete_file_contents=[], + errors=[], + ), + messages=[Message(id="m1", role="user", parts=[TextPart("处理重复预订")])], + ctx=_ctx(), + ) + ) + + assert result.request_count == 1 + assert result.operations.upsert_operations[0].memory_type == "cases" + assert result.apply_result.written_uris == ["viking://user/u/memories/cases/重复预订处理.md"] + assert fs.writes + written_uri, written_content, _ = fs.writes[0] + assert written_uri.endswith("/memories/cases/重复预订处理.md") + assert "重复预订处理" in written_content + + +@pytest.mark.asyncio +async def test_streaming_memory_updater_fast_path_filters_links(monkeypatch): + fs = InMemoryVikingFS( + { + "viking://user/u/memories/events/existing.md": ( + 'existing\n' + ) + } + ) + fs.search = AsyncMock(return_value=[]) + monkeypatch.setattr( + "openviking.session.memory.streaming_memory_updater.get_viking_fs", + lambda: fs, + ) + monkeypatch.setattr( + "openviking.session.memory.memory_updater.get_viking_fs", + lambda: fs, + ) + + updater = StreamingMemoryUpdater( + registry=_registry(), + config=StreamingMemoryUpdaterConfig( + max_operations_per_update=8, + max_wait_seconds=0.01, + timer_check_interval_seconds=0.01, + ), + ) + op1 = _case_op("并发案例A") + link = StoredLink( + from_uri=op1.uris[0], + to_uri="viking://user/u/memories/events/existing.md", + link_type="related_to", + weight=0.8, + match_text="并发", + description="valid link", + ) + duplicate_link = link.model_copy(update={"weight": 0.6, "description": "short"}) + missing_link = StoredLink( + from_uri=op1.uris[0], + to_uri="viking://user/u/memories/events/missing.md", + link_type="related_to", + weight=0.9, + match_text="缺失", + description="invalid link", + ) + + result = await updater.submit( + MemoryUpdateRequest( + operations=ResolvedOperations( + upsert_operations=[op1], + delete_file_contents=[], + errors=[], + resolved_links=[link, duplicate_link, missing_link], + ), + messages=[Message(id="m1", role="user", parts=[TextPart("并发A")])], + ctx=_ctx(), + ) + ) + + assert result.request_count == 1 + assert result.metadata["flush_reason"] == "append_only_fast_path" + assert len(result.operations.upsert_operations) == 1 + assert len(result.operations.resolved_links) == 1 + assert result.operations.resolved_links[0].to_uri.endswith("/events/existing.md") + assert result.apply_result.written_uris == [op1.uris[0]] + + +@pytest.mark.asyncio +async def test_streaming_memory_updater_batches_non_append_only_submits(monkeypatch): + fs = InMemoryVikingFS({}) + fs.search = AsyncMock(return_value=[]) + monkeypatch.setattr( + "openviking.session.memory.streaming_memory_updater.get_viking_fs", + lambda: fs, + ) + monkeypatch.setattr( + "openviking.session.memory.memory_updater.get_viking_fs", + lambda: fs, + ) + + updater = StreamingMemoryUpdater( + registry=_registry(), + config=StreamingMemoryUpdaterConfig( + max_operations_per_update=2, + max_wait_seconds=0.01, + timer_check_interval_seconds=0.01, + ), + ) + op1 = _note_op("note_a") + op2 = _note_op("note_b") + + result1, result2 = await asyncio.gather( + updater.submit( + MemoryUpdateRequest( + operations=ResolvedOperations( + upsert_operations=[op1], + delete_file_contents=[], + errors=[], + ), + messages=[Message(id="m1", role="user", parts=[TextPart("note A")])], + ctx=_ctx(), + ) + ), + updater.submit( + MemoryUpdateRequest( + operations=ResolvedOperations( + upsert_operations=[op2], + delete_file_contents=[], + errors=[], + ), + messages=[Message(id="m2", role="user", parts=[TextPart("note B")])], + ctx=_ctx(), + ) + ), + ) + + assert result1 is not result2 + assert result1.request_count == 1 + assert result2.request_count == 1 + assert result1.metadata["flush_reason"] == "count" + assert result1.metadata["batch_request_count"] == 2 + assert result1.metadata["scoped_to_submitter"] is True + assert result1.apply_result.written_uris == [op1.uris[0]] + assert result2.apply_result.written_uris == [op2.uris[0]] + assert sorted(result1.metadata["unscoped_written_uris"]) == sorted([op1.uris[0], op2.uris[0]]) + + +def test_scope_memory_update_result_to_submitter_filters_shared_batch_by_source(): + from openviking.session.memory.streaming_memory_updater import ( + scope_memory_update_result_to_submitter, + ) + + op_a = _note_op_with_source("scoped_a", "extract_a") + op_b = _note_op_with_source("scoped_b", "extract_b") + apply_result = MemoryUpdateResult() + apply_result.add_written(op_a.uris[0]) + apply_result.add_written(op_b.uris[0]) + batch_result = StreamingMemoryUpdateResult( + operations=ResolvedOperations( + upsert_operations=[op_a, op_b], + delete_file_contents=[], + errors=[], + ), + apply_result=apply_result, + request_count=2, + metadata={"flush_reason": "count", "operation_count": 2}, + ) + request = MemoryUpdateRequest( + operations=ResolvedOperations( + upsert_operations=[op_a], + delete_file_contents=[], + errors=[], + ), + messages=[Message(id="m1", role="user", parts=[TextPart("note A")])], + ctx=_ctx(), + metadata={"source_extraction_id": "extract_a", "session_id": "session_a"}, + ) + + scoped = scope_memory_update_result_to_submitter(batch_result, request) + + assert scoped.request_count == 1 + assert scoped.metadata["batch_request_count"] == 2 + assert scoped.metadata["scoped_to_source_extraction_id"] == "extract_a" + assert scoped.apply_result.written_uris == [op_a.uris[0]] + assert scoped.operations.upsert_operations == [op_a] + assert scoped.metadata["unscoped_written_uris"] == [op_a.uris[0], op_b.uris[0]] + + +def test_split_request_by_merge_group_groups_by_peer_and_memory_type(): + self_op = _note_op("self_note") + peer_op = _peer_note_op("peer_note", "web-visitor-alice") + case_op = _case_op("case_note") + link = StoredLink( + from_uri=self_op.uris[0], + to_uri=peer_op.uris[0], + link_type="related_to", + weight=0.8, + ) + request = MemoryUpdateRequest( + operations=ResolvedOperations( + upsert_operations=[self_op, peer_op, case_op], + delete_file_contents=[], + errors=[], + resolved_links=[link], + ), + messages=[], + ctx=_ctx(), + ) + + grouped = split_request_by_merge_group(request) + + assert [key for key, _ in grouped] == [ + MemoryMergeGroupKey(peer_id=None, memory_type="notes"), + MemoryMergeGroupKey(peer_id="web-visitor-alice", memory_type="notes"), + MemoryMergeGroupKey(peer_id=None, memory_type="cases"), + ] + assert [len(group_request.operations.upsert_operations) for _, group_request in grouped] == [ + 1, + 1, + 1, + ] + assert [len(group_request.operations.resolved_links) for _, group_request in grouped] == [ + 0, + 0, + 0, + ] + + +def test_split_request_by_merge_group_infers_peer_from_uri_when_field_missing(): + peer_uri = "viking://user/u/peers/conv-42/memories/notes/peer_note.md" + op = ResolvedOperation( + old_memory_file_content=None, + memory_fields={"note_name": "peer_note", "content": "peer content"}, + memory_type="notes", + uris=[peer_uri], + ) + request = MemoryUpdateRequest( + operations=ResolvedOperations( + upsert_operations=[op], + delete_file_contents=[], + errors=[], + ), + messages=[], + ctx=_ctx(), + ) + + grouped = split_request_by_merge_group(request) + + assert [key for key, _ in grouped] == [ + MemoryMergeGroupKey(peer_id="conv-42", memory_type="notes") + ] + + +def test_enforce_merge_group_peer_id_rewrites_merged_output_scope(): + op = ResolvedOperation( + old_memory_file_content=None, + memory_fields={"note_name": "peer_note", "content": "peer content"}, + memory_type="notes", + uris=["viking://user/u/memories/notes/peer_note.md"], + ) + + enforce_merge_group_peer_id( + [op], + peer_id="conv-42", + memory_type="notes", + registry=_registry(), + ctx=_ctx(), + ) + + assert op.memory_fields["peer_id"] == "conv-42" + assert op.uris == ["viking://user/u/peers/conv-42/memories/notes/peer_note.md"] + + +def test_enforce_merge_group_self_scope_removes_peer_id(): + op = ResolvedOperation( + old_memory_file_content=None, + memory_fields={ + "note_name": "self_note", + "content": "self content", + "peer_id": "conv-42", + }, + memory_type="notes", + uris=["viking://user/u/peers/conv-42/memories/notes/self_note.md"], + ) + + enforce_merge_group_peer_id( + [op], + peer_id=None, + memory_type="notes", + registry=_registry(), + ctx=_ctx(), + ) + + assert "peer_id" not in op.memory_fields + assert op.uris == ["viking://user/u/memories/notes/self_note.md"] + + +def test_enforce_merge_group_peer_enabled_false_keeps_self_scope(): + op = ResolvedOperation( + old_memory_file_content=None, + memory_fields={ + "case_name": "case_note", + "task_signature": "case signature", + "input": "{}", + "rubric": "{}", + "peer_id": "conv-42", + }, + memory_type="cases", + uris=["viking://user/u/peers/conv-42/memories/cases/case_note.md"], + ) + + enforce_merge_group_peer_id( + [op], + peer_id="conv-42", + memory_type="cases", + registry=_registry(), + ctx=_ctx(), + ) + + assert "peer_id" not in op.memory_fields + assert op.uris == ["viking://user/u/memories/cases/case_note.md"] + + +@pytest.mark.asyncio +async def test_streaming_memory_updater_batches_per_merge_group(monkeypatch): + fs = InMemoryVikingFS({}) + fs.search = AsyncMock(return_value=[]) + monkeypatch.setattr( + "openviking.session.memory.streaming_memory_updater.get_viking_fs", + lambda: fs, + ) + monkeypatch.setattr( + "openviking.session.memory.memory_updater.get_viking_fs", + lambda: fs, + ) + + updater = StreamingMemoryUpdater( + registry=_registry(), + config=StreamingMemoryUpdaterConfig( + max_operations_per_update=2, + max_wait_seconds=0.05, + timer_check_interval_seconds=0.01, + ), + ) + note_a = _note_op("note_group_a") + note_b = _note_op("note_group_b") + peer_note = _peer_note_op("note_peer", "web-visitor-alice") + + result1, result2, peer_result = await asyncio.gather( + updater.submit( + MemoryUpdateRequest( + operations=ResolvedOperations( + upsert_operations=[note_a], + delete_file_contents=[], + errors=[], + ), + messages=[Message(id="m1", role="user", parts=[TextPart("note A")])], + ctx=_ctx(), + ) + ), + updater.submit( + MemoryUpdateRequest( + operations=ResolvedOperations( + upsert_operations=[note_b], + delete_file_contents=[], + errors=[], + ), + messages=[Message(id="m2", role="user", parts=[TextPart("note B")])], + ctx=_ctx(), + ) + ), + updater.submit( + MemoryUpdateRequest( + operations=ResolvedOperations( + upsert_operations=[peer_note], + delete_file_contents=[], + errors=[], + ), + messages=[Message(id="m3", role="user", parts=[TextPart("peer note")])], + ctx=_ctx(), + ) + ), + ) + + assert result1 is not result2 + assert result1.request_count == 1 + assert result2.request_count == 1 + assert result1.metadata["flush_reason"] == "count" + assert result1.metadata["batch_request_count"] == 2 + assert result1.metadata["merge_group"] == "peer=self,memory_type=notes" + assert result1.apply_result.written_uris == [note_a.uris[0]] + assert result2.apply_result.written_uris == [note_b.uris[0]] + + assert peer_result is not result1 + assert peer_result.request_count == 1 + assert peer_result.metadata["flush_reason"] == "time" + assert peer_result.metadata["merge_group"] == "peer=web-visitor-alice,memory_type=notes" + assert peer_result.apply_result.written_uris == [peer_note.uris[0]] + + +@pytest.mark.asyncio +async def test_streaming_memory_updater_submit_waits_for_all_merge_groups(monkeypatch): + fs = InMemoryVikingFS({}) + fs.search = AsyncMock(return_value=[]) + monkeypatch.setattr( + "openviking.session.memory.streaming_memory_updater.get_viking_fs", + lambda: fs, + ) + monkeypatch.setattr( + "openviking.session.memory.memory_updater.get_viking_fs", + lambda: fs, + ) + + updater = StreamingMemoryUpdater( + registry=_registry(), + config=StreamingMemoryUpdaterConfig( + max_operations_per_update=8, + max_wait_seconds=0.01, + timer_check_interval_seconds=0.01, + ), + ) + self_op = _note_op("multi_self") + peer_op = _peer_note_op("multi_peer", "web-visitor-alice") + + result = await updater.submit( + MemoryUpdateRequest( + operations=ResolvedOperations( + upsert_operations=[self_op, peer_op], + delete_file_contents=[], + errors=[], + ), + messages=[Message(id="m1", role="user", parts=[TextPart("multi group")])], + ctx=_ctx(), + ) + ) + + assert result.metadata["combined_result"] is True + assert result.request_count == 1 + assert result.metadata["batch_request_count"] == 2 + assert sorted(result.apply_result.written_uris) == sorted([self_op.uris[0], peer_op.uris[0]]) + assert self_op.uris[0] in fs.files + assert peer_op.uris[0] in fs.files + + +@pytest.mark.asyncio +async def test_streaming_memory_updater_applies_cross_group_links_after_all_groups(monkeypatch): + fs = InMemoryVikingFS({}) + fs.search = AsyncMock(return_value=[]) + monkeypatch.setattr( + "openviking.session.memory.streaming_memory_updater.get_viking_fs", + lambda: fs, + ) + monkeypatch.setattr( + "openviking.session.memory.memory_updater.get_viking_fs", + lambda: fs, + ) + + updater = StreamingMemoryUpdater( + registry=_registry(), + config=StreamingMemoryUpdaterConfig( + max_operations_per_update=8, + max_wait_seconds=0.01, + timer_check_interval_seconds=0.01, + ), + ) + self_op = _note_op("linked_self") + peer_op = _peer_note_op("linked_peer", "web-visitor-alice") + link = StoredLink( + from_uri=self_op.uris[0], + to_uri=peer_op.uris[0], + link_type="related_to", + weight=0.8, + match_text="linked", + ) + + result = await updater.submit( + MemoryUpdateRequest( + operations=ResolvedOperations( + upsert_operations=[self_op, peer_op], + delete_file_contents=[], + errors=[], + resolved_links=[link], + ), + messages=[Message(id="m1", role="user", parts=[TextPart("cross group link")])], + ctx=_ctx(), + ) + ) + + self_file = MemoryFileUtils.read(fs.files[self_op.uris[0]], uri=self_op.uris[0]) + peer_file = MemoryFileUtils.read(fs.files[peer_op.uris[0]], uri=peer_op.uris[0]) + + assert len(result.operations.resolved_links) == 1 + assert self_file.links[0]["to_uri"] == peer_op.uris[0] + assert peer_file.backlinks[0]["from_uri"] == self_op.uris[0] + + +def test_classify_memory_merge_mode_forces_cross_extraction_merge(): + op1 = _note_op_with_source("note_a", "extract_a") + op2 = _note_op_with_source("note_b", "extract_b") + + fast_path, reason = classify_memory_merge_mode([op1, op2], schema=_registry().get("notes")) + + assert fast_path is False + assert reason == "cross_extraction_batch" + + +def test_classify_memory_merge_mode_treats_noop_str_patch_as_unchanged(): + old_file = MemoryFile( + uri="viking://user/u/memories/notes/note.md", + content="old content", + memory_type="notes", + extra_fields={"note_name": "note"}, + ) + op = ResolvedOperation( + old_memory_file_content=old_file, + memory_type="notes", + uris=["viking://user/u/memories/notes/note.md"], + memory_fields={ + "note_name": "note", + "content": StrPatch( + blocks=[SearchReplaceBlock(search="old content", replace="old content")] + ), + }, + ) + + fast_path, reason = classify_memory_merge_mode([op], schema=_registry().get("notes")) + + assert fast_path is True + assert reason == "single_existing_content_unchanged" + + +def test_classify_memory_merge_mode_detects_changed_str_patch_after_preview(): + old_file = MemoryFile( + uri="viking://user/u/memories/notes/note.md", + content="old content", + memory_type="notes", + extra_fields={"note_name": "note"}, + ) + op = ResolvedOperation( + old_memory_file_content=old_file, + memory_type="notes", + uris=["viking://user/u/memories/notes/note.md"], + memory_fields={ + "note_name": "note", + "content": StrPatch( + blocks=[SearchReplaceBlock(search="old content", replace="new content")] + ), + }, + ) + + fast_path, reason = classify_memory_merge_mode([op], schema=_registry().get("notes")) + + assert fast_path is False + assert reason == "single_existing_content_changed" + + +@pytest.mark.asyncio +async def test_streaming_memory_updater_persists_source_extraction_id_trace_id_and_hides_from_read( + monkeypatch, +): + fs = InMemoryVikingFS({}) + fs.search = AsyncMock(return_value=[]) + monkeypatch.setattr( + "openviking.session.memory.streaming_memory_updater.get_viking_fs", + lambda: fs, + ) + monkeypatch.setattr( + "openviking.session.memory.memory_updater.get_viking_fs", + lambda: fs, + ) + + updater = StreamingMemoryUpdater( + registry=_registry(), + config=StreamingMemoryUpdaterConfig( + max_operations_per_update=8, + max_wait_seconds=0.01, + timer_check_interval_seconds=0.01, + ), + ) + op = _note_op("note_source") + result = await updater.submit( + MemoryUpdateRequest( + operations=ResolvedOperations( + upsert_operations=[op], + delete_file_contents=[], + errors=[], + ), + messages=[Message(id="m1", role="user", parts=[TextPart("note source")])], + ctx=_ctx(), + metadata={"source_extraction_id": "extract_1", "trace_id": "trace_1"}, + ) + ) + + assert result.apply_result.written_uris == [op.uris[0]] + assert '"source_extraction_id": "extract_1"' in fs.files[op.uris[0]] + assert '"last_update_trace_id": "trace_1"' in fs.files[op.uris[0]] + + from openviking.server.identity import ToolContext + from openviking.session.memory.tools import MemoryReadTool + + read_result = await MemoryReadTool().execute( + ToolContext(viking_fs=fs, request_ctx=_ctx(), read_file_contents={}), + uri=op.uris[0], + ) + + assert "source_extraction_id" not in read_result + assert "last_update_trace_id" not in read_result + + +def test_render_operation_after_file_content_persists_source_trace_id(): + schema = _registry().get("notes") + op = _note_op("note_trace") + op.source = MemoryOperationSource(extraction_id="extract_2", trace_id="trace_2") + + rendered = render_operation_after_file_content( + op, + schema=schema, + extract_context=ExtractContext([]), + ) + + assert '"source_extraction_id": "extract_2"' in rendered + assert '"last_update_trace_id": "trace_2"' in rendered + + +@pytest.mark.asyncio +async def test_cross_extraction_merge_preserves_existing_uri_without_explicit_delete(monkeypatch): + existing_uri = "viking://user/u/memories/notes/existing.md" + winner_uri = "viking://user/u/memories/notes/winner.md" + old_file = __import__( + "openviking.session.memory.dataclass", fromlist=["MemoryFile"] + ).MemoryFile( + uri=existing_uri, + content="old", + memory_type="notes", + extra_fields={"note_name": "existing"}, + ) + existing_op = ResolvedOperation( + old_memory_file_content=old_file, + memory_type="notes", + uris=[existing_uri], + memory_fields={ + "note_name": "existing", + "content": {"blocks": [{"search": "old", "replace": "old updated"}]}, + "source_extraction_id": "extract_a", + }, + ) + new_op = ResolvedOperation( + old_memory_file_content=None, + memory_type="notes", + uris=[winner_uri], + memory_fields={ + "note_name": "winner", + "content": "merged content", + "source_extraction_id": "extract_b", + }, + ) + + async def fake_run(self): + return ( + ResolvedOperations( + upsert_operations=[new_op], + delete_file_contents=[], + errors=[], + ), + [], + ) + + monkeypatch.setattr( + "openviking.session.memory.streaming_memory_updater.ExtractLoop.run", + fake_run, + ) + fs = InMemoryVikingFS({existing_uri: "old"}) + fs.search = AsyncMock(return_value=[]) + monkeypatch.setattr( + "openviking.session.memory.streaming_memory_updater.get_viking_fs", + lambda: fs, + ) + monkeypatch.setattr( + "openviking.session.memory.memory_updater.get_viking_fs", + lambda: fs, + ) + + merged = await merge_one_memory_type_operations( + memory_type="notes", + operations=[existing_op, new_op], + messages=[], + ctx=_ctx(), + registry=_registry(), + ) + + assert [op.uris for op in merged.upsert_operations] == [[winner_uri]] + assert merged.delete_file_contents == [] + + +@pytest.mark.asyncio +async def test_patch_merge_uses_original_messages_for_output_language(monkeypatch): + existing_uri = "viking://user/u/memories/notes/code.md" + old_file = MemoryFile( + uri=existing_uri, + content="old", + memory_type="notes", + extra_fields={"memory_type": "notes", "topic": "code"}, + ) + existing_op = ResolvedOperation( + old_memory_file_content=old_file, + memory_fields={"topic": "code", "content": "older"}, + memory_type="notes", + uris=[existing_uri], + ) + new_op = ResolvedOperation( + old_memory_file_content=None, + memory_fields={"topic": "code", "content": "new"}, + memory_type="notes", + uris=["viking://user/u/memories/notes/code_new.md"], + ) + captured_languages = [] + + async def fake_run(self): + captured_languages.append(self.context_provider.get_output_language()) + return ( + ResolvedOperations( + upsert_operations=[existing_op], + delete_file_contents=[], + errors=[], + ), + [], + ) + + monkeypatch.setattr( + "openviking.session.memory.streaming_memory_updater.ExtractLoop.run", + fake_run, + ) + fs = InMemoryVikingFS({existing_uri: "old"}) + fs.search = AsyncMock(return_value=[]) + monkeypatch.setattr( + "openviking.session.memory.streaming_memory_updater.get_viking_fs", + lambda: fs, + ) + monkeypatch.setattr( + "openviking.session.memory.memory_updater.get_viking_fs", + lambda: fs, + ) + + await merge_one_memory_type_operations( + memory_type="notes", + operations=[existing_op, new_op], + messages=[Message(id="m1", role="user", parts=[TextPart("请保持中文记忆")])], + ctx=_ctx(), + registry=_registry(), + ) + + assert captured_languages == ["zh-CN"] diff --git a/tests/session/test_compressor_v3.py b/tests/session/test_compressor_v3.py new file mode 100644 index 0000000000..dc2700b475 --- /dev/null +++ b/tests/session/test_compressor_v3.py @@ -0,0 +1,1127 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from openviking.message import Message, TextPart +from openviking.server.identity import RequestContext, Role +from openviking.session import create_session_compressor +from openviking.session.compressor_v3 import SessionCompressorV3, _experience_root_uri +from openviking.session.memory.dataclass import ( + MemoryFile, + ResolvedOperation, + ResolvedOperations, + StoredLink, +) +from openviking.session.memory.memory_updater import MemoryUpdateResult +from openviking.session.memory.utils.memory_file_utils import MemoryFileUtils +from openviking.session.train import ( + Case, + ExperienceSet, + PolicyApplyResult, + PolicyPlanItem, + PolicyUpdatePlan, + Rollout, + RolloutAnalysis, + RolloutTrainingResult, + Rubric, + RubricCriterion, + RubricEvaluation, + StreamingPolicyTrainerConfig, + Trajectory, +) +from openviking.session.train.components.session_commit import _case_spec_message_to_request +from openviking_cli.session.user_id import UserIdentifier + + +def _ctx() -> RequestContext: + return RequestContext(user=UserIdentifier.the_default_user("u"), role=Role.ROOT) + + +def _messages() -> list[Message]: + return [ + Message( + id="m1", + role="user", + parts=[TextPart("请处理重复预订,只取消确认是重复的那一单。")], + ), + Message( + id="m2", + role="assistant", + parts=[TextPart("已读取两个预订,确认第二个是重复记录并取消。")], + ), + ] + + +def _case_operation() -> ResolvedOperation: + return ResolvedOperation( + old_memory_file_content=None, + memory_type="cases", + uris=["viking://user/u/memories/cases/重复预订处理.md"], + memory_fields={ + "case_name": "重复预订处理", + "task_signature": "处理重复预订并只取消确认重复的订单", + "input": '{"summary":"用户要求处理重复预订","preconditions":["存在两个相似预订"]}', + "rubric": '{"name":"重复预订处理Rubric","description":"成功且高效处理重复预订","criteria":[{"name":"先验证重复","description":"取消前必须确认哪一单是重复订单","required":true,"weight":0.6},{"name":"只取消重复项","description":"不得影响有效订单","required":true,"weight":0.4}]}', + "evidence": "助手根据读取结果确认重复项并完成取消。", + }, + ) + + +def test_factory_defaults_to_v3(): + compressor = create_session_compressor(vikingdb=None) + assert isinstance(compressor, SessionCompressorV3) + + +def test_factory_ignores_deprecated_memory_version(): + assert isinstance( + create_session_compressor(vikingdb=None, memory_version="v2"), SessionCompressorV3 + ) + assert isinstance( + create_session_compressor(vikingdb=None, memory_version="unsupported"), + SessionCompressorV3, + ) + + +def test_experience_root_uri_requires_request_user(): + with pytest.raises(ValueError, match="RequestContext.user.user_id is required"): + _experience_root_uri(SimpleNamespace(user=None)) + + +def test_case_experience_links_require_policy_root_uri(): + traj_uri = "viking://user/u/memories/trajectories/t.md" + plan = PolicyUpdatePlan( + items=[ + PolicyPlanItem( + kind="upsert", + memory_type="experiences", + target_name="exp", + target_uri=None, + before_content=None, + after_content="exp", + links=[ + StoredLink( + from_uri="", + to_uri=traj_uri, + link_type="derived_from", + weight=1.0, + ) + ], + ) + ] + ) + apply_result = PolicyApplyResult( + updated_policy_set=ExperienceSet(root_uri="", policies=[]), + written_uris=["viking://user/u/memories/experiences/exp.md"], + ) + + from openviking.session.compressor_v3 import _case_experience_links_via_trajectories + + with pytest.raises(ValueError, match="updated_policy_set.root_uri is required"): + _case_experience_links_via_trajectories( + case_uri="viking://user/u/memories/cases/case.md", + trajectory_uris={traj_uri}, + plan=plan, + apply_result=apply_result, + ) + + +@pytest.mark.asyncio +async def test_train_from_extracted_case_memories_submits_streaming_rollout(monkeypatch): + submitted_gradients = [] + submitted_analyses = [] + + class FakeTrainer: + policy_set = ExperienceSet( + root_uri="viking://user/u/memories/experiences", + policies=[], + ) + + async def submit_gradients(self, gradients, *, analysis=None, rollout=None): + submitted_gradients.append(gradients) + submitted_analyses.append(analysis) + return RolloutTrainingResult( + analyses=[analysis] if analysis else [], + gradients=list(gradients), + plan=PolicyUpdatePlan(items=[], metadata={}), + apply_result=PolicyApplyResult( + updated_policy_set=self.policy_set, + written_uris=[], + errors=[], + ), + metadata={}, + ) + + class FakeAnalyzer: + async def analyze(self, rollout, context): + return RolloutAnalysis( + evaluation=RubricEvaluation( + passed=True, + score=1.0, + criterion_results=[], + feedback=[], + ), + trajectories=[ + Trajectory( + name="duplicate_booking", + uri="viking://user/u/memories/trajectories/t1.md", + content="trajectory content", + outcome="success", + retrieval_anchor="", + ) + ], + gradients=[], + ) + + async def fake_estimate_exp_gradients(self, *args, **kwargs): + # Return one dummy gradient so we can verify submission + from openviking.session.memory.dataclass import MemoryFile + from openviking.session.train import PatchSemanticGradient + + return [ + PatchSemanticGradient( + before_file=None, + after_file=MemoryFile( + uri="viking://user/u/memories/experiences/e1.md", + content="new exp", + memory_type="experiences", + extra_fields={"experience_name": "e1"}, + ), + base_version=1, + rationale="test", + links=[], + confidence=0.9, + metadata={}, + ) + ] + + monkeypatch.setattr( + "openviking.session.compressor_v3.get_viking_fs", + lambda: SimpleNamespace(ls=AsyncMock(return_value=[])), + ) + monkeypatch.setattr( + "openviking.session.compressor_v3.get_streaming_policy_trainer", + AsyncMock(return_value=FakeTrainer()), + ) + monkeypatch.setattr( + "openviking.session.train.components.gradient_estimator.ExperienceGradientEstimator.estimate", + fake_estimate_exp_gradients, + ) + + compressor = SessionCompressorV3( + vikingdb=None, + rollout_analyzer=FakeAnalyzer(), + streaming_trainer_config=StreamingPolicyTrainerConfig( + max_wait_seconds=60, + max_gradients_per_update=8, + ), + ) + operations = ResolvedOperations( + upsert_operations=[_case_operation()], + delete_file_contents=[], + errors=[], + ) + + # The extracted case comes from the same memory operations as profile/preferences/etc.; + # no extra LLM/VLM case extractor is involved. + cases = __import__( + "openviking.session.compressor_v3", fromlist=["_operations_to_cases"] + )._operations_to_cases(operations) + result = await compressor.train_from_extracted_cases( + cases=cases, + messages=_messages(), + ctx=_ctx(), + session_id="s1", + ) + + assert result["case_count"] == 1 + assert result["submitted"] == 1 + assert len(submitted_gradients) == 1 + assert len(submitted_gradients[0]) == 1 # one exp gradient per case + # Verify analysis was used + assert submitted_analyses[0] is not None + assert submitted_analyses[0].trajectories[0].name == "duplicate_booking" + # Verify case info carried through correctly + assert cases[0].name == "重复预订处理" + assert cases[0].input["summary"] == "用户要求处理重复预订" + assert cases[0].rubric.criteria[0].name == "先验证重复" + + +@pytest.mark.asyncio +async def test_train_from_extracted_multiple_case_memories_analyzes_bound_rollouts(monkeypatch): + seen_rollouts = [] + rollout_messages = _messages() + + class FakeTrainer: + policy_set = ExperienceSet( + root_uri="viking://user/u/memories/experiences", + policies=[], + ) + + class FakeAnalyzer: + async def analyze(self, rollout, context): + del context + seen_rollouts.append(rollout) + return RolloutAnalysis( + evaluation=RubricEvaluation( + passed=True, + score=1.0, + criterion_results=[], + feedback=[], + ), + trajectories=[], + gradients=[], + ) + + async def fake_estimate_exp_gradients(self, *args, **kwargs): + return [] + + monkeypatch.setattr( + "openviking.session.compressor_v3.get_viking_fs", + lambda: SimpleNamespace(ls=AsyncMock(return_value=[])), + ) + monkeypatch.setattr( + "openviking.session.compressor_v3.get_streaming_policy_trainer", + AsyncMock(return_value=FakeTrainer()), + ) + monkeypatch.setattr( + "openviking.session.train.components.gradient_estimator.ExperienceGradientEstimator.estimate", + fake_estimate_exp_gradients, + ) + + case_a = _training_case() + case_b = Case( + name="case_b", + task_signature="Handle a second extracted case.", + input={"summary": "second case"}, + rubric=case_a.rubric, + ) + + compressor = SessionCompressorV3( + vikingdb=None, + rollout_analyzer=FakeAnalyzer(), + streaming_trainer_config=StreamingPolicyTrainerConfig( + max_wait_seconds=60, + max_gradients_per_update=8, + ), + ) + + result = await compressor.train_from_extracted_cases( + cases=[case_a, case_b], + messages=rollout_messages, + ctx=_ctx(), + session_id="s1", + ) + + assert result["case_count"] == 2 + assert result["submitted"] == 2 + assert [rollout.case.name for rollout in seen_rollouts] == [ + "duplicate_booking", + "case_b", + ] + assert [rollout.messages for rollout in seen_rollouts] == [ + rollout_messages, + rollout_messages, + ] + + +@pytest.mark.asyncio +async def test_v3_extract_uses_patch_merge_without_directory_lock(monkeypatch): + applied_operations = [] + trained_cases = [] + + class DummyRegistry: + async def initialize_memory_files(self, ctx): + return None + + class DummyOrchestrator: + async def run(self): + return ( + ResolvedOperations( + upsert_operations=[_case_operation()], + delete_file_contents=[], + errors=[], + ), + [], + ) + + class FakeStreamingUpdater: + async def submit(self, request): + applied_operations.append(request.operations) + result = MemoryUpdateResult() + result.add_written(_case_operation().uris[0]) + return SimpleNamespace(operations=request.operations, apply_result=result) + + compressor = SessionCompressorV3(vikingdb=None) + compressor._get_or_create_react = lambda **kwargs: DummyOrchestrator() + + async def fake_train_from_extracted_cases(**kwargs): + trained_cases.extend(kwargs["cases"]) + return {"case_count": len(kwargs["cases"]), "submitted": len(kwargs["cases"])} + + compressor.train_from_extracted_cases = fake_train_from_extracted_cases + + monkeypatch.setattr( + "openviking.session.compressor_v3.get_viking_fs", + lambda: SimpleNamespace(write_file=AsyncMock()), + ) + monkeypatch.setattr( + "openviking.session.compressor_v3.create_default_registry", + lambda: DummyRegistry(), + ) + monkeypatch.setattr( + "openviking.session.compressor_v3.get_streaming_memory_updater", + AsyncMock(return_value=FakeStreamingUpdater()), + ) + + contexts = await compressor.extract_long_term_memories( + messages=_messages(), + ctx=_ctx(), + allowed_memory_types={"cases", "profile"}, + ) + + assert len(applied_operations) == 1 + assert applied_operations[0].upsert_operations[0].memory_type == "cases" + assert [case.name for case in trained_cases] == ["重复预订处理"] + assert contexts[0].uri.endswith("重复预订处理.md") + + +@pytest.mark.asyncio +async def test_v3_extract_trains_only_canonical_case_after_patch_merge(monkeypatch): + trained_kwargs = [] + canonical_uri = "viking://user/u/memories/cases/duplicate_booking.md" + loser_uri = "viking://user/u/memories/cases/duplicate_booking_duplicate.md" + canonical_fields = { + "case_name": "duplicate_booking", + "task_signature": "Handle duplicate bookings safely.", + "input": '{"summary":"cancel only the confirmed duplicate booking"}', + "rubric": ( + '{"name":"duplicate_booking_rubric","description":"Verify duplicate handling",' + '"criteria":[{"name":"verify_duplicate","description":"The assistant verifies ' + 'which booking is duplicate before cancellation.","required":true,"weight":1.0}]}' + ), + "evidence": "Canonical merged case evidence.", + } + + def case_op(uri: str, name: str) -> ResolvedOperation: + fields = dict(canonical_fields) + fields["case_name"] = name + return ResolvedOperation( + old_memory_file_content=None, + memory_type="cases", + uris=[uri], + memory_fields=fields, + ) + + class DummyRegistry: + async def initialize_memory_files(self, ctx): + return None + + class DummyOrchestrator: + async def run(self): + return ( + ResolvedOperations( + upsert_operations=[ + case_op(canonical_uri, "duplicate_booking"), + case_op(loser_uri, "duplicate_booking_duplicate"), + ], + delete_file_contents=[], + errors=[], + ), + [], + ) + + class FakeFS: + async def read_file(self, uri, ctx=None): + del ctx + if uri != canonical_uri: + raise FileNotFoundError(uri) + return MemoryFileUtils.write( + MemoryFile( + uri=canonical_uri, + content="", + memory_type="cases", + extra_fields=dict(canonical_fields), + ) + ) + + async def write_file(self, uri, content, ctx=None): + del uri, content, ctx + + class FakeStreamingUpdater: + async def submit(self, request): + del request + result = MemoryUpdateResult() + result.add_written(canonical_uri) + return SimpleNamespace( + operations=ResolvedOperations( + upsert_operations=[case_op(canonical_uri, "duplicate_booking")], + delete_file_contents=[], + errors=[], + ), + apply_result=result, + ) + + compressor = SessionCompressorV3(vikingdb=None) + compressor._get_or_create_react = lambda **kwargs: DummyOrchestrator() + + async def fake_train_from_extracted_cases(**kwargs): + trained_kwargs.append(kwargs) + return {"case_count": len(kwargs["cases"]), "submitted": len(kwargs["cases"])} + + compressor.train_from_extracted_cases = fake_train_from_extracted_cases + + monkeypatch.setattr("openviking.session.compressor_v3.get_viking_fs", lambda: FakeFS()) + monkeypatch.setattr( + "openviking.session.compressor_v3.create_default_registry", + lambda: DummyRegistry(), + ) + monkeypatch.setattr( + "openviking.session.compressor_v3.get_streaming_memory_updater", + AsyncMock(return_value=FakeStreamingUpdater()), + ) + + contexts = await compressor.extract_long_term_memories( + messages=_messages(), + ctx=_ctx(), + allowed_memory_types={"cases", "profile"}, + ) + + assert [context.uri for context in contexts] == [canonical_uri] + assert len(trained_kwargs) == 1 + assert [case.name for case in trained_kwargs[0]["cases"]] == ["duplicate_booking"] + assert trained_kwargs[0]["cases"][0].metadata["case_uris"] == [canonical_uri] + assert trained_kwargs[0]["case_uri_by_name"] == {"duplicate_booking": canonical_uri} + + +def _training_case() -> Case: + return Case( + name="duplicate_booking", + task_signature="Handle duplicate bookings safely.", + input={"summary": "cancel only the duplicate booking", "task_id": "task-1"}, + rubric=Rubric( + name="duplicate_booking_rubric", + description="Verify duplicates before cancellation.", + criteria=[ + RubricCriterion( + name="verify_duplicate", + description="The assistant verifies which booking is the duplicate before acting.", + required=True, + weight=1.0, + ) + ], + ), + metadata={"evidence": "The rollout contains duplicate-booking handling evidence."}, + ) + + +def _case_spec_message(case: Case | None = None) -> Message: + rollout = Rollout( + case=case or _training_case(), + messages=[], + policy_snapshot_id="snapshot-1", + ) + request = _case_spec_message_to_request(rollout) + return Message( + id="case-spec", + role="system", + parts=[TextPart(text=request["parts"][0]["text"])], + ) + + +@pytest.mark.asyncio +async def test_v3_training_case_spec_fast_path_skips_user_memory_extraction_and_strips_control_message(): + case_spec = _case_spec_message() + rollout_messages = _messages() + written = [] + trained = [] + + compressor = SessionCompressorV3(vikingdb=None, rollout_analyzer=SimpleNamespace()) + + async def fail_extract_user_memories(**kwargs): + raise AssertionError("fast path must not run LLM user-memory extraction") + + async def fake_write_training_case_memory(**kwargs): + written.append(kwargs["case"]) + result = MemoryUpdateResult() + result.add_written("viking://user/u/memories/cases/duplicate_booking.md") + return result + + async def fake_train_from_extracted_cases(**kwargs): + trained.append(kwargs) + return {"case_count": len(kwargs["cases"]), "submitted": len(kwargs["cases"])} + + compressor._extract_user_memories = fail_extract_user_memories + compressor._write_training_case_memory = fake_write_training_case_memory + compressor.train_from_extracted_cases = fake_train_from_extracted_cases + + contexts = await compressor.extract_long_term_memories( + messages=[case_spec, *rollout_messages], + ctx=_ctx(), + session_id="s1", + archive_uri="viking://user/u/sessions/s1/history/archive_001", + allowed_memory_types={"cases", "trajectories", "experiences"}, + ) + + assert [case.name for case in written] == ["duplicate_booking"] + assert [case.name for case in trained[0]["cases"]] == ["duplicate_booking"] + assert trained[0]["messages"] == rollout_messages + assert contexts[0].uri == "viking://user/u/memories/cases/duplicate_booking.md" + + +@pytest.mark.asyncio +async def test_v3_training_case_spec_fast_path_not_used_with_user_memory_policy(): + extracted = False + trained = [] + + compressor = SessionCompressorV3(vikingdb=None, rollout_analyzer=SimpleNamespace()) + + async def fake_extract_user_memories(**kwargs): + nonlocal extracted + extracted = True + return SimpleNamespace(contexts=[], cases=[]) + + async def fake_train_from_extracted_cases(**kwargs): + trained.append(kwargs) + return {"case_count": 0, "submitted": 0} + + compressor._extract_user_memories = fake_extract_user_memories + compressor.train_from_extracted_cases = fake_train_from_extracted_cases + + contexts = await compressor.extract_long_term_memories( + messages=[_case_spec_message(), *_messages()], + ctx=_ctx(), + allowed_memory_types={"cases", "profile"}, + ) + + assert contexts == [] + assert extracted is True + assert trained and trained[0]["messages"][0].id == "case-spec" + + +@pytest.mark.asyncio +async def test_v3_training_case_spec_fast_path_rejects_invalid_protocol(): + message = _case_spec_message() + assert isinstance(message.parts[0], TextPart) + message.parts[0].text = message.parts[0].text.replace( + "openviking.batch_train.case_spec.v1", + "openviking.batch_train.case_spec.v0", + ) + compressor = SessionCompressorV3(vikingdb=None, rollout_analyzer=SimpleNamespace()) + + with pytest.raises(ValueError, match="protocol mismatch"): + await compressor.extract_long_term_memories( + messages=[message, *_messages()], + ctx=_ctx(), + allowed_memory_types={"cases", "trajectories", "experiences"}, + ) + + +def test_training_case_spec_message_uses_fast_path_protocol(): + message = _case_spec_message() + part = message.parts[0] + assert isinstance(part, TextPart) + text = part.text + + assert text.startswith("# OpenViking Batch Training CaseSpec v1") + assert "openviking.batch_train.case_spec.v1" in text + assert "duplicate_booking_rubric" in text + + +def test_training_case_spec_message_uses_original_case_name_for_trials(): + case = _training_case() + case.name = "tau2_airline_train_1_t0" + case.task_signature = "tau2:airline:train:1:trial:0" + case.input.update( + { + "domain": "airline", + "split": "train", + "data_split": "airline_train", + "task_id": "1", + "task_no": 1, + "train_trial": 0, + "original_case_name": "tau2_airline_train_1", + } + ) + message = _case_spec_message(case) + payload = __import__( + "openviking.session.compressor_v3", fromlist=["_training_case_spec_payload_from_message"] + )._training_case_spec_payload_from_message(message) + + assert payload["case"]["name"] == "tau2_airline_train_1" + assert payload["case"]["task_signature"] == "tau2:airline:train:1" + assert payload["case"]["metadata"]["rollout_case_name"] == "tau2_airline_train_1_t0" + + +@pytest.mark.asyncio +async def test_v3_fast_path_writes_final_memory_diff_with_case_traj_and_exp(monkeypatch): + archive_uri = "viking://user/u/sessions/s1/history/archive_001" + writes: dict[str, str] = {} + + class FakeFS: + async def write_file(self, uri, content, ctx=None): + del ctx + writes[uri] = content + + async def read_file(self, uri, ctx=None): + del ctx + if uri.endswith("/cases/duplicate_booking.md"): + return "# duplicate_booking\n\n" + if uri.endswith("/experiences/booking_duplicate_handling.md"): + return "new exp content\n\n" + raise FileNotFoundError(uri) + + compressor = SessionCompressorV3(vikingdb=None, rollout_analyzer=SimpleNamespace()) + + async def fake_write_training_case_memory(**kwargs): + result = MemoryUpdateResult() + result.add_written("viking://user/u/memories/cases/duplicate_booking.md") + return SimpleNamespace( + result=result, + memory_diff={ + "archive_uri": archive_uri, + "trace_id": None, + "extracted_at": "now", + "operations": { + "adds": [ + { + "uri": "viking://user/u/memories/cases/duplicate_booking.md", + "memory_type": "cases", + "after": "# duplicate_booking", + } + ], + "updates": [], + "deletes": [], + }, + "summary": {"total_adds": 1, "total_updates": 0, "total_deletes": 0}, + }, + ) + + async def fake_train_from_extracted_cases(**kwargs): + return { + "case_count": 1, + "submitted": 1, + "memory_diff": { + "archive_uri": archive_uri, + "trace_id": None, + "extracted_at": "now", + "operations": { + "adds": [ + { + "uri": "viking://user/u/memories/trajectories/duplicate_booking.md", + "memory_type": "trajectories", + "after": "trajectory content", + } + ], + "updates": [ + { + "uri": "viking://user/u/memories/experiences/booking_duplicate_handling.md", + "memory_type": "experiences", + "before": "old exp content", + "after": "new exp content", + } + ], + "deletes": [], + }, + "summary": {"total_adds": 1, "total_updates": 1, "total_deletes": 0}, + }, + } + + compressor._write_training_case_memory = fake_write_training_case_memory + compressor.train_from_extracted_cases = fake_train_from_extracted_cases + monkeypatch.setattr("openviking.session.compressor_v3.get_viking_fs", lambda: FakeFS()) + + contexts = await compressor.extract_long_term_memories( + messages=[_case_spec_message(), *_messages()], + ctx=_ctx(), + session_id="s1", + archive_uri=archive_uri, + allowed_memory_types={"cases", "trajectories", "experiences"}, + ) + + assert contexts[0].uri.endswith("/cases/duplicate_booking.md") + diff = __import__("json").loads(writes[f"{archive_uri}/memory_diff.json"]) + assert [item["memory_type"] for item in diff["operations"]["adds"]] == [ + "cases", + "trajectories", + ] + assert [item["memory_type"] for item in diff["operations"]["updates"]] == ["experiences"] + assert diff["summary"] == {"total_adds": 2, "total_updates": 1, "total_deletes": 0} + + +@pytest.mark.asyncio +async def test_v3_builds_training_memory_diff_from_streaming_result(monkeypatch): + archive_uri = "viking://user/u/sessions/s1/history/archive_001" + + class FakeFS: + async def read_file(self, uri, ctx=None): + del ctx + assert uri.endswith("/experiences/booking_duplicate_handling.md") + return "new exp content\n\n" + + compressor = SessionCompressorV3(vikingdb=None, rollout_analyzer=SimpleNamespace()) + plan = PolicyUpdatePlan( + items=[ + PolicyPlanItem( + kind="upsert", + memory_type="experiences", + target_name="booking_duplicate_handling", + target_uri="viking://user/u/memories/experiences/booking_duplicate_handling.md", + before_content="old exp content", + after_content="new exp content fallback", + links=[ + StoredLink( + from_uri="viking://user/u/memories/experiences/booking_duplicate_handling.md", + to_uri="viking://user/u/memories/trajectories/duplicate_booking.md", + link_type="derived_from", + weight=1.0, + ) + ], + ) + ] + ) + training_result = RolloutTrainingResult( + analyses=[ + RolloutAnalysis( + evaluation=RubricEvaluation( + passed=True, + score=1.0, + criterion_results=[], + feedback=[], + ), + trajectories=[ + Trajectory( + name="duplicate_booking", + uri="viking://user/u/memories/trajectories/duplicate_booking.md", + content="trajectory content", + outcome="success", + retrieval_anchor="Stage: final", + ) + ], + ) + ], + gradients=[], + plan=plan, + apply_result=PolicyApplyResult( + updated_policy_set=ExperienceSet( + root_uri="viking://user/u/memories/experiences", + policies=[], + ), + written_uris=["viking://user/u/memories/experiences/booking_duplicate_handling.md"], + ), + ) + + diff = await compressor._build_training_memory_diff( + training_result=training_result, + viking_fs=FakeFS(), + ctx=_ctx(), + archive_uri=archive_uri, + ) + + assert diff["summary"] == {"total_adds": 1, "total_updates": 1, "total_deletes": 0} + assert diff["operations"]["adds"][0]["memory_type"] == "trajectories" + update = diff["operations"]["updates"][0] + assert update["memory_type"] == "experiences" + assert update["before"] == "old exp content" + assert update["after"] == "new exp content" + + +@pytest.mark.asyncio +async def test_v3_training_memory_diff_filters_batch_items_by_current_analysis_trajectory( + monkeypatch, +): + archive_uri = "viking://user/u/sessions/s1/history/archive_001" + traj_a = "viking://user/u/memories/trajectories/traj_a.md" + traj_b = "viking://user/u/memories/trajectories/traj_b.md" + exp_a = "viking://user/u/memories/experiences/exp_a.md" + exp_b = "viking://user/u/memories/experiences/exp_b.md" + + class FakeFS: + async def read_file(self, uri, ctx=None): + del ctx + return { + exp_a: "exp a\n\n", + exp_b: "exp b\n\n", + }[uri] + + compressor = SessionCompressorV3(vikingdb=None, rollout_analyzer=SimpleNamespace()) + training_result = RolloutTrainingResult( + analyses=[ + RolloutAnalysis( + evaluation=RubricEvaluation( + passed=True, score=1.0, criterion_results=[], feedback=[] + ), + trajectories=[ + Trajectory( + name="traj_a", + uri=traj_a, + content="trajectory a", + outcome="success", + retrieval_anchor="", + ) + ], + ) + ], + gradients=[], + plan=PolicyUpdatePlan( + items=[ + PolicyPlanItem( + kind="upsert", + memory_type="experiences", + target_name="exp_a", + target_uri=exp_a, + before_content=None, + after_content="exp a fallback", + links=[ + StoredLink( + from_uri=exp_a, + to_uri=traj_a, + link_type="derived_from", + weight=1.0, + ) + ], + ), + PolicyPlanItem( + kind="upsert", + memory_type="experiences", + target_name="exp_b", + target_uri=exp_b, + before_content=None, + after_content="exp b fallback", + links=[ + StoredLink( + from_uri=exp_b, + to_uri=traj_b, + link_type="derived_from", + weight=1.0, + ) + ], + ), + ] + ), + apply_result=PolicyApplyResult( + updated_policy_set=ExperienceSet( + root_uri="viking://user/u/memories/experiences", + policies=[], + ), + written_uris=[exp_a, exp_b], + ), + ) + + diff = await compressor._build_training_memory_diff( + training_result=training_result, + viking_fs=FakeFS(), + ctx=_ctx(), + archive_uri=archive_uri, + ) + + assert diff["summary"] == {"total_adds": 2, "total_updates": 0, "total_deletes": 0} + assert [op["uri"] for op in diff["operations"]["adds"]] == [traj_a, exp_a] + + +@pytest.mark.asyncio +async def test_v3_training_links_case_to_trajectory_and_experience_via_trajectory(monkeypatch): + case_uri = "viking://user/u/memories/cases/duplicate_booking.md" + traj_uri = "viking://user/u/memories/trajectories/duplicate_booking.md" + exp_uri = "viking://user/u/memories/experiences/booking_duplicate_handling.md" + + class FakeFS: + def __init__(self): + self.files = { + case_uri: MemoryFileUtils.write( + MemoryFile( + uri=case_uri, + content="# duplicate_booking", + memory_type="cases", + extra_fields={"memory_type": "cases", "case_name": "duplicate_booking"}, + ), + content_template=( + "# {{ case_name }}\n\n" + "## Linked Experiences\n" + "{% for link in links or [] %}" + "{% set target_uri = link.to_uri or '' %}" + "{% if '/memories/experiences/' in target_uri %}" + "- [{{ uri_basename(target_uri) }}]({{ link_target(target_uri) }})\n" + "{% endif %}" + "{% endfor %}" + ), + ), + traj_uri: MemoryFileUtils.write( + MemoryFile( + uri=traj_uri, + content="trajectory content", + memory_type="trajectories", + extra_fields={ + "memory_type": "trajectories", + "trajectory_name": "duplicate_booking", + }, + backlinks=[ + StoredLink( + from_uri=exp_uri, + to_uri=traj_uri, + link_type="derived_from", + weight=1.0, + match_text=None, + description="", + ).model_dump() + ], + ) + ), + exp_uri: MemoryFileUtils.write( + MemoryFile( + uri=exp_uri, + content="old exp content", + memory_type="experiences", + extra_fields={ + "memory_type": "experiences", + "experience_name": "booking_duplicate_handling", + }, + ) + ), + } + + async def read_file(self, uri, ctx=None): + del ctx + return self.files[uri] + + async def write_file(self, uri, content, ctx=None): + del ctx + self.files[uri] = content + + async def ls(self, uri, output="original", ctx=None): + del uri, output, ctx + return [] + + class FakeTrainer: + policy_set = ExperienceSet(root_uri="viking://user/u/memories/experiences", policies=[]) + + async def submit_gradients(self, gradients, *, analysis=None, rollout=None): + del gradients, analysis, rollout + plan = PolicyUpdatePlan( + items=[ + PolicyPlanItem( + kind="upsert", + memory_type="experiences", + target_name="booking_duplicate_handling", + target_uri=exp_uri, + before_content="old exp content", + after_content="new exp content", + links=[ + StoredLink( + from_uri=exp_uri, + to_uri=traj_uri, + link_type="derived_from", + weight=1.0, + match_text=None, + description="", + ) + ], + ) + ] + ) + return RolloutTrainingResult( + analyses=[], + gradients=[], + plan=plan, + apply_result=PolicyApplyResult( + updated_policy_set=ExperienceSet( + root_uri="viking://user/u/memories/experiences", + policies=[], + ), + written_uris=[exp_uri], + errors=[], + ), + metadata={}, + ) + + class FakeAnalyzer: + async def analyze(self, rollout, context): + del rollout, context + return RolloutAnalysis( + evaluation=RubricEvaluation( + passed=True, score=1.0, criterion_results=[], feedback=[] + ), + trajectories=[ + Trajectory( + name="duplicate_booking", + uri=traj_uri, + content="trajectory content", + outcome="success", + retrieval_anchor="", + ) + ], + gradients=[], + ) + + async def fake_estimate_exp_gradients(self, *args, **kwargs): + from openviking.session.train import PatchSemanticGradient + + return [ + PatchSemanticGradient( + before_file=None, + after_file=MemoryFile( + uri=exp_uri, + content="new exp content", + memory_type="experiences", + extra_fields={"experience_name": "booking_duplicate_handling"}, + ), + base_version=1, + rationale="test", + links=[], + confidence=0.9, + metadata={}, + ) + ] + + fs = FakeFS() + monkeypatch.setattr("openviking.session.compressor_v3.get_viking_fs", lambda: fs) + monkeypatch.setattr( + "openviking.session.compressor_v3.get_streaming_policy_trainer", + AsyncMock(return_value=FakeTrainer()), + ) + monkeypatch.setattr( + "openviking.session.train.components.gradient_estimator.ExperienceGradientEstimator.estimate", + fake_estimate_exp_gradients, + ) + + compressor = SessionCompressorV3(vikingdb=None, rollout_analyzer=FakeAnalyzer()) + result = await compressor.train_from_extracted_cases( + cases=[_training_case()], + case_uri_by_name={"duplicate_booking": case_uri}, + messages=_messages(), + ctx=_ctx(), + ) + + assert result["submitted"] == 1 + case_file = MemoryFileUtils.read(fs.files[case_uri], uri=case_uri) + assert any( + link["to_uri"] == traj_uri + and link["link_type"] == "related_to" + and link.get("match_text") is None + and link.get("description") == "" + for link in case_file.links + ) + assert any( + link["to_uri"] == exp_uri + and link["link_type"] == "related_to" + and link.get("match_text") is None + and link.get("description") == "" + for link in case_file.links + ) + linked_experiences_section = ( + fs.files[case_uri].split("## Linked Experiences", 1)[1].split("" + ) + return RolloutAnalysis( + evaluation=evaluation, + trajectories=[ + Trajectory( + name="complex_duplicate_booking_case", + uri=self.trajectory_uri, + content=trajectory_content.strip(), + outcome=trajectory_outcome, + retrieval_anchor=f"阶段:{learning_target};能力:复杂重复预订处理", + metadata={"memory_type": "trajectories"}, + ) + ], + metadata={ + "rollout_messages": rollout.messages, + "policy_snapshot_id": rollout.policy_snapshot_id, + }, + ) + + +def _case() -> Case: + return Case( + name="complex_duplicate_booking_case", + task_signature="complex_booking_duplicate", + input={ + "user_request": ( + "我好像重复订了酒店,但不确定哪一笔重复。请帮我处理,如果确实重复就取消多余的那笔。" + ), + "booking_records": [ + { + "booking_id": "A-100", + "hotel": "上海中心酒店", + "check_in": "2026-07-01", + "check_out": "2026-07-03", + "room_type": "豪华大床房", + "guest": "张三", + "price_cny": 1800, + "created_at": "2026-05-01T10:00:00", + "cancellation_fee_cny": 0, + "refund_cny": 1800, + "status": "confirmed", + }, + { + "booking_id": "A-101", + "hotel": "上海中心酒店", + "check_in": "2026-07-01", + "check_out": "2026-07-03", + "room_type": "豪华大床房", + "guest": "张三", + "price_cny": 1800, + "created_at": "2026-05-01T10:05:00", + "cancellation_fee_cny": 0, + "refund_cny": 1800, + "status": "confirmed", + }, + { + "booking_id": "A-102", + "hotel": "上海中心酒店", + "check_in": "2026-07-01", + "check_out": "2026-07-03", + "room_type": "行政双床房", + "guest": "张三", + "price_cny": 2200, + "created_at": "2026-05-01T10:08:00", + "cancellation_fee_cny": 300, + "refund_cny": 1900, + "status": "confirmed", + }, + ], + "policy_constraints": [ + "没有用户二次确认前,不能真正执行取消动作。", + "如果要取消,需要先告知退款金额和取消费。", + ], + }, + rubric=Rubric( + name="complex_booking_duplicate_hidden_rubric", + description=("分阶段处理复杂重复预订:识别重复、选对目标、检查取消成本、请求确认。"), + criteria=[ + RubricCriterion( + name="identify_duplicate", + description=( + "必须核验酒店、入住/离店日期、房型、入住人、价格等关键信息," + "识别 A-100 和 A-101 是重复候选。" + ), + required=True, + weight=0.25, + ), + RubricCriterion( + name="select_correct_target", + description=( + "必须选择取消创建时间更晚的 A-101,保留 A-100," + "且不能误取消同酒店同日期但房型不同的 A-102。" + ), + required=True, + weight=0.25, + ), + RubricCriterion( + name="check_cancellation_cost", + description="取消前必须说明 A-101 的退款金额和取消费,并提醒 A-102 有取消费。", + required=True, + weight=0.25, + ), + RubricCriterion( + name="request_confirmation", + description="不能直接执行取消;必须先给出建议并请求用户确认后再取消。", + required=True, + weight=0.25, + ), + ], + ), + ) + + +def _strict_policy_prompt(case, policy_set, context) -> str: + del context + experiences = "\n\n".join( + f"### {policy.name} v{policy.version} [{policy.status}]\n{policy.content}" + for policy in policy_set.policies + ) + if not experiences: + experiences = "(no experience policies available)" + return "\n".join( + [ + "你正在执行一个受经验策略控制的离线案例。", + "必须优先遵循 Experience Policies;如果经验策略与常识冲突,也先按经验策略输出。", + "这里不会向你展示隐藏评分 Rubric。请只根据用户输入和经验策略完成任务。", + "请给出你会执行的最终处理动作,不要只复述原则。", + "", + "# Experience Policies", + experiences, + "", + "# Case Input", + json.dumps(case.input, ensure_ascii=False, indent=2, sort_keys=True), + ] + ) + + +def _passed_criterion_names(payload: dict) -> set[str]: + return { + str(item.get("criterion_name")) + for item in payload.get("criterion_results", []) + if isinstance(item, dict) and item.get("passed") + } + + +def _first_failed_criterion_feedback(payload: dict, case: Case) -> tuple[str | None, str]: + results = { + str(item.get("criterion_name")): item + for item in payload.get("criterion_results", []) + if isinstance(item, dict) + } + for criterion in case.rubric.criteria: + item = results.get(criterion.name) + if item is None or not item.get("passed"): + feedback = item.get("feedback") if isinstance(item, dict) else None + if isinstance(feedback, list) and feedback: + return criterion.name, "; ".join(str(value) for value in feedback) + return criterion.name, criterion.description + return None, "所有阶段都已通过。" + + +def _staged_score(payload: dict, case: Case) -> float: + passed = _passed_criterion_names(payload) + score = 0.0 + for criterion in case.rubric.criteria: + if criterion.name not in passed: + break + score += float(criterion.weight) + return max(0.0, min(1.0, score)) + + +async def _evaluate_rollout_with_real_llm( + *, + vlm, + case: Case, + rollout_messages: list[Message], + active_criteria: list[str] | None = None, +): + assistant_text = "\n".join( + message.content for message in rollout_messages if message.role == "assistant" + ) + active_criteria = active_criteria or [criterion.name for criterion in case.rubric.criteria] + active_criterion_set = set(active_criteria) + visible_criteria = [ + criterion for criterion in case.rubric.criteria if criterion.name in active_criterion_set + ] + prompt = "\n".join( + [ + "你是 OpenViking 离线训练的严格评估器。", + "请评估助手是否真正完成了当前已解锁的 hidden Rubric。只根据助手输出评分。", + "不要因为 Case Input 中出现了约束就给分;必须看助手是否实际执行了该要求。", + "", + "# Case Input", + json.dumps(case.input, ensure_ascii=False, indent=2, sort_keys=True), + "", + "# Active Hidden Rubric", + f"{case.rubric.name}: {case.rubric.description}", + *[ + f"- {criterion.name} ({'required' if criterion.required else 'optional'}, " + f"weight={criterion.weight}): {criterion.description}" + for criterion in visible_criteria + ], + "", + "# Assistant Output", + assistant_text, + "", + "# 评分规则", + "- 每个 active criterion 独立判断 passed/score。", + "- 总分 score = active criteria 中通过项数量 / active criteria 数量。", + "- 非 active criterion 不参与本轮评分,也不要出现在 criterion_results 中。", + "- 只输出 JSON,不要输出 markdown。", + json.dumps( + { + "passed": False, + "score": 0.0, + "feedback": ["string"], + "criterion_results": [ + { + "criterion_name": "identify_duplicate", + "passed": False, + "score": 0.0, + "feedback": ["string"], + "evidence": ["string"], + }, + { + "criterion_name": "select_correct_target", + "passed": False, + "score": 0.0, + "feedback": ["string"], + "evidence": ["string"], + }, + ], + }, + ensure_ascii=False, + indent=2, + ), + ] + ) + response = await vlm.get_completion_async(prompt=prompt, thinking=False) + payload = parse_json_from_response(response) + if not isinstance(payload, dict): + return { + "passed": False, + "score": 0.0, + "feedback": ["评估器输出无法解析为 JSON。"], + "criterion_results": [], + "raw_response": getattr(response, "content", str(response)), + } + payload.setdefault("feedback", []) + payload.setdefault("criterion_results", []) + payload["criterion_results"] = [ + item + for item in payload["criterion_results"] + if isinstance(item, dict) and item.get("criterion_name") in active_criterion_set + ] + if not payload["criterion_results"]: + payload["criterion_results"] = [ + { + "criterion_name": criterion.name, + "passed": False, + "score": 0.0, + "feedback": ["评估器没有返回该阶段的结构化结果。"], + "evidence": [], + } + for criterion in visible_criteria + ] + for item in payload["criterion_results"]: + item["passed"] = bool(item.get("passed")) + try: + item["score"] = max(0.0, min(1.0, float(item.get("score", 0.0)))) + except (TypeError, ValueError): + item["score"] = 0.0 + passed_count = sum(1 for item in payload["criterion_results"] if item["passed"]) + payload["score"] = _staged_score(payload, case) + payload["passed"] = bool(payload.get("passed")) and passed_count == len(visible_criteria) + payload["active_criteria"] = active_criteria + return payload + + +def _print_real_llm_e2e_summary( + *, + assistant_text: str, + trajectory_content: str, + gradient, + written_experience: str | None = None, +) -> None: + lines = [ + "\n========== Real LLM Policy Optimization E2E =========", + f"[TraceID] {tracer.get_trace_id()}", + "[Rollout Assistant]", + assistant_text, + "", + "[Extracted Trajectory]", + trajectory_content, + "", + "[Semantic Gradient]", + f"target_name: {gradient.target_name}", + f"target_uri: {gradient.target_uri}", + f"base_version: {gradient.base_version}", + f"confidence: {gradient.confidence}", + "", + "[Gradient before_content]", + gradient.before_file.plain_content() if gradient.before_file is not None else "None", + "", + "[Gradient after_content]", + gradient.after_file.plain_content(), + ] + if written_experience is not None: + lines.extend(["", "[Written Experience File]", written_experience]) + lines.append("=====================================================\n") + tracer.info("\n".join(lines), console=True) + + +def _print_iterative_real_llm_summary( + *, + result, + final_evaluation, + fs: InMemoryVikingFS, + experience_uri: str, +) -> None: + lines = [ + "\n========== Real LLM Iterative Policy Optimization =========", + f"[TraceID] {tracer.get_trace_id()}", + f"epochs: {len(result.epochs)}", + f"final_evaluation_score: {final_evaluation.metadata.get('score')}", + f"first_score: {result.metadata.get('first_score')}", + f"final_score: {result.metadata.get('final_score')}", + f"score_delta: {result.metadata.get('score_delta')}", + f"last_optimizer: {result.plan.metadata.get('optimizer')}", + f"last_merge_errors: {result.plan.metadata.get('merge_errors')}", + ] + for epoch in result.epochs: + lines.extend( + [ + "", + f"[Epoch {epoch.epoch}]", + f"score: {epoch.metadata.get('score')}", + f"snapshot_ids: {epoch.policy_snapshot_ids}", + f"gradient_count: {epoch.metadata.get('gradient_count')}", + f"written_uris: {epoch.apply_result.written_uris}", + f"errors: {epoch.apply_result.errors}", + ] + ) + if epoch.gradients: + for gradient_idx, gradient in enumerate(epoch.gradients): + lines.extend( + [ + f"[Epoch {epoch.epoch} Gradient {gradient_idx}]", + f"target_name: {gradient.target_name}", + f"target_uri: {gradient.target_uri}", + f"confidence: {gradient.confidence}", + ] + ) + lines.extend( + [ + "gradient_after_content:", + gradient.after_file.plain_content(), + ] + ) + for analysis in epoch.analyses: + messages = analysis.metadata.get("rollout_messages", []) + assistant_text = "\n".join( + message.content for message in messages if message.role == "assistant" + ) + lines.extend( + [ + f"case: {analysis.trajectories[0].name if analysis.trajectories else 'n/a'}", + f"passed: {analysis.evaluation.passed}", + f"feedback: {'; '.join(analysis.evaluation.feedback)}", + "assistant:", + assistant_text, + ] + ) + lines.extend( + [ + "", + f"[Final Evaluation {final_evaluation.epoch}]", + f"score: {final_evaluation.metadata.get('score')}", + f"snapshot_ids: {final_evaluation.policy_snapshot_ids}", + ] + ) + for analysis in final_evaluation.analyses: + messages = analysis.metadata.get("rollout_messages", []) + assistant_text = "\n".join( + message.content for message in messages if message.role == "assistant" + ) + lines.extend( + [ + f"passed: {analysis.evaluation.passed}", + f"feedback: {'; '.join(analysis.evaluation.feedback)}", + "assistant:", + assistant_text, + ] + ) + lines.extend(["", "[Updated Experience File]", fs.files.get(experience_uri, "")]) + lines.append("==========================================================\n") + tracer.info("\n".join(lines), console=True) + + +def _patch_experience_prefetch( + monkeypatch, fs: InMemoryVikingFS, experience_uri: str +) -> None: + async def search_files(self, query, search_uris=None, limit=5): + return [experience_uri] + + async def read_file(self, uri): + raw = fs.files.get(uri) + if raw is None: + target_name = uri.rstrip("/").split("/")[-1].removesuffix(".md") + for candidate_uri, candidate_raw in fs.files.items(): + supersedes = f'"supersedes":"{target_name}"' + if candidate_uri.endswith(".md") and supersedes in candidate_raw: + raw = candidate_raw + uri = candidate_uri + break + if raw is None: + raw = fs.files[experience_uri] + uri = experience_uri + from openviking.session.memory.utils.memory_file_utils import MemoryFileUtils + + memory_file = MemoryFileUtils.read(raw, uri=uri) + self.read_file_contents[uri] = memory_file + result = memory_file.to_metadata() + result["content"] = memory_file.content + result["uri"] = uri + return result + + monkeypatch.setattr( + "openviking.session.memory.session_extract_context_provider.SessionExtractContextProvider.search_files", + search_files, + ) + monkeypatch.setattr( + "openviking.session.memory.session_extract_context_provider.SessionExtractContextProvider.read_file", + read_file, + ) + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_policy_optimization_pipeline_real_config_llm_e2e_writes_updated_experience( + monkeypatch, +): + with start_current_span("train.test.real_llm_e2e.pipeline"): + await _run_policy_optimization_pipeline_real_config_llm_e2e_writes_updated_experience( + monkeypatch + ) + + +async def _run_policy_optimization_pipeline_real_config_llm_e2e_writes_updated_experience( + monkeypatch, +): + root = "viking://user/u/memories/experiences" + experience_uri = f"{root}/booking_duplicate_handling.md" + trajectory_uri = "viking://user/u/memories/trajectories/complex_duplicate_booking_case.md" + fs = InMemoryVikingFS( + { + experience_uri: ( + "## Situation\n- 处理用户取消预订请求\n\n" + "## Approach\n" + "- 如果用户说可能重复预订,只对比酒店和入住日期。\n" + "- 如果酒店和入住日期相同,就直接取消所有相关预订,不需要检查房型、退款或二次确认。\n\n" + "" + ) + } + ) + reset_lock_manager() + init_lock_manager(fs.agfs, redo_recovery_enabled=False) + request_context = SimpleNamespace( + user=SimpleNamespace(account_id="default", user_id="u"), + account_id="default", + ) + policy_set = await ExperienceSetLoader(viking_fs=fs).load(root, ctx=request_context) + vlm = get_openviking_config().vlm + _patch_experience_prefetch(monkeypatch, fs, experience_uri) + + pipeline = OfflinePolicyOptimizationPipeline( + snapshotter=ContentHashPolicySnapshotter(), + rollout_executor=SingleTurnLLMRolloutExecutor( + vlm=vlm, + prompt_builder=_strict_policy_prompt, + thinking=False, + ), + rollout_analyzer=RealRubricTrajectoryAnalyzer( + trajectory_uri=trajectory_uri, + viking_fs=fs, + vlm=vlm, + ), + gradient_estimator=ExperienceGradientEstimator( + viking_fs=fs, + vlm=vlm, + ), + policy_optimizer=PatchMergePolicyOptimizer( + viking_fs=fs, + vlm=vlm, + ), + policy_updater=MemoryFilePolicyUpdater(viking_fs=fs), + ) + + result = await pipeline.train( + case_loader=ListCaseLoader([_case()]), + policy_set=policy_set, + context=PipelineContext( + analysis_context=TrajectoryAnalyzerContext(request_context=request_context), + gradient_context=ExperienceGradientContext( + request_context=request_context, + messages=[], + ), + optimization_context=PatchMergePolicyOptimizerContext(request_context=request_context), + apply_context=request_context, + max_epochs=4, + ), + ) + final_evaluation = await pipeline.eval( + case_loader=ListCaseLoader([_case()]), + policy_set=result.apply_result.updated_policy_set, + context=PipelineContext( + analysis_context=TrajectoryAnalyzerContext(request_context=request_context), + execution_metadata={"epoch": 4}, + ), + ) + + rollout_messages = result.analyses[0].metadata["rollout_messages"] + assistant_text = rollout_messages[1].content + trajectory_content = result.analyses[0].trajectories[0].content + gradient = result.gradients[0] + _print_iterative_real_llm_summary( + result=result, + final_evaluation=final_evaluation, + fs=fs, + experience_uri=experience_uri, + ) + assert assistant_text.strip() + assert trajectory_content.strip() + assert gradient.after_file.plain_content().strip() + assert all(epoch.apply_result.errors == [] for epoch in result.epochs) + written_uris = [ + uri for epoch in result.epochs for uri in epoch.apply_result.written_uris + ] + assert experience_uri in written_uris + assert result.plan.metadata["optimizer"] == "patch_merge" + assert any( + epoch.plan.metadata.get("optimizer") == "patch_merge" for epoch in result.epochs + ) + assert len(result.epochs) == 4 + assert result.evaluation_passes == [] + assert final_evaluation.metadata["score"] > result.metadata["first_score"] + assert result.metadata["score_delta"] > 0 + assert len({epoch.metadata["score"] for epoch in result.epochs}) >= 3 + assert "重复" in fs.files[experience_uri] + assert "房型" in fs.files[experience_uri] + assert "确认" in fs.files[experience_uri] + + +@pytest.mark.asyncio +@pytest.mark.integration +@tracer( + "train.test.real_llm_e2e.gradient_estimator", + ignore_result=True, + ignore_args=True, + is_new_trace=True, +) +async def test_experience_gradient_estimator_real_config_llm_generates_gradient( + monkeypatch, +): + root = "viking://user/u/memories/experiences" + experience_uri = f"{root}/booking_duplicate_handling.md" + trajectory_uri = "viking://user/u/memories/trajectories/duplicate_booking_case.md" + fs = InMemoryVikingFS( + { + experience_uri: ( + "## Situation\n- 重复预订处理\n\n" + "" + ), + trajectory_uri: ( + "# 重复预订处理轨迹\n" + "用户要求取消重复预订。助手先核验两笔预订是否确实重复," + "然后只取消重复的那一笔,避免误取消原始有效预订。\n\n" + "" + ), + } + ) + reset_lock_manager() + init_lock_manager(fs.agfs, redo_recovery_enabled=False) + request_context = SimpleNamespace( + user=SimpleNamespace(account_id="default", user_id="u"), + account_id="default", + ) + policy_set = await ExperienceSetLoader(viking_fs=fs).load(root, ctx=request_context) + + _patch_experience_prefetch(monkeypatch, fs, experience_uri) + + rollout_executor = SingleTurnLLMRolloutExecutor( + vlm=get_openviking_config().vlm, + thinking=False, + ) + analyzer = TrajectoryRolloutAnalyzer(viking_fs=fs) + snapshotter = ContentHashPolicySnapshotter() + snapshot_id = await snapshotter.snapshot(policy_set) + rollouts = await rollout_executor.execute( + [_case()], + policy_set, + SimpleNamespace(policy_snapshot_id=snapshot_id, metadata={}), + ) + analysis = await analyzer.analyze( + rollouts[0], + TrajectoryAnalyzerContext(request_context=request_context), + ) + + estimator = ExperienceGradientEstimator( + viking_fs=fs, + vlm=get_openviking_config().vlm, + ) + gradients = await estimator.estimate( + analysis, + policy_set, + ExperienceGradientContext(request_context=request_context, messages=[]), + ) + + assert gradients + gradient = gradients[0] + _print_real_llm_e2e_summary( + assistant_text=analysis.metadata["rollout_messages"][1].content, + trajectory_content=analysis.trajectories[0].content, + gradient=gradient, + ) + assert gradient.target_name + assert gradient.after_file.plain_content().strip() + assert gradient.links + assert gradient.links[0].to_uri in fs.files + assert "/memories/trajectories/" in gradient.links[0].to_uri diff --git a/tests/session/train/test_rollout_executor_component.py b/tests/session/train/test_rollout_executor_component.py new file mode 100644 index 0000000000..f89b5c9879 --- /dev/null +++ b/tests/session/train/test_rollout_executor_component.py @@ -0,0 +1,917 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 + +from __future__ import annotations + +import asyncio +import threading +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from openviking.session.train import ( + Case, + ExecutionContext, + Experience, + ExperienceSet, + Rubric, + RubricCriterion, + SingleTurnLLMRolloutExecutor, + default_single_turn_prompt, +) + + +class FakeVLM: + def __init__(self, response="assistant answer"): + self.response = response + self.calls = [] + + async def get_completion_async(self, **kwargs): + self.calls.append(kwargs) + return self.response + + +def _case() -> Case: + return Case( + name="case-1", + task_signature="booking_duplicate", + input={"user_request": "cancel duplicate booking"}, + rubric=Rubric( + name="booking_rubric", + description="Cancel only the verified duplicate booking.", + criteria=[ + RubricCriterion( + name="verify_duplicate", + description="Verify duplicate status first.", + required=True, + weight=1.0, + ) + ], + ), + ) + + +def _policy_set() -> ExperienceSet: + return ExperienceSet( + root_uri="viking://user/u/memories/experiences", + policies=[ + Experience( + name="booking_policy", + uri="viking://user/u/memories/experiences/booking_policy.md", + version=2, + status="production", + content="Always verify duplicates before cancellation.", + ) + ], + ) + + +@pytest.mark.asyncio +async def test_single_turn_llm_rollout_executor_produces_rollout_messages(): + vlm = FakeVLM() + executor = SingleTurnLLMRolloutExecutor(vlm=vlm, thinking=False) + context = ExecutionContext(policy_snapshot_id="snapshot-1") + + rollouts = await executor.execute([_case()], _policy_set(), context) + + assert len(rollouts) == 1 + rollout = rollouts[0] + assert rollout.case.name == "case-1" + assert rollout.policy_snapshot_id == "snapshot-1" + assert [message.role for message in rollout.messages] == ["user", "assistant"] + assert "Always verify duplicates" in rollout.messages[0].content + assert "cancel duplicate booking" in rollout.messages[0].content + assert rollout.messages[1].content == "assistant answer" + assert vlm.calls[0]["thinking"] is False + assert vlm.calls[0]["prompt"] == rollout.messages[0].content + + +@pytest.mark.asyncio +async def test_single_turn_llm_rollout_executor_accepts_custom_prompt_builder(): + vlm = FakeVLM(response=type("Resp", (), {"content": "structured answer"})()) + + def build_prompt(case, policy_set, context): + return f"custom:{case.name}:{len(policy_set.policies)}:{context.policy_snapshot_id}" + + executor = SingleTurnLLMRolloutExecutor(vlm=vlm, prompt_builder=build_prompt) + + rollouts = await executor.execute( + [_case()], + _policy_set(), + ExecutionContext(policy_snapshot_id="snapshot-2"), + ) + + assert rollouts[0].messages[0].content == "custom:case-1:1:snapshot-2" + assert rollouts[0].messages[1].content == "structured answer" + + +def test_default_single_turn_prompt_contains_case_policy_and_rubric(): + prompt = default_single_turn_prompt( + _case(), + _policy_set(), + ExecutionContext(policy_snapshot_id="snapshot-3"), + ) + + assert "Policy snapshot: snapshot-3" in prompt + assert "booking_policy v2 [production]" in prompt + assert "cancel duplicate booking" in prompt + assert "verify_duplicate" in prompt + + +def test_dataset_service_policy_set_from_dict_preserves_policies(): + from openviking.session.train.components.dataset_service import policy_set_from_dict + + policy_set = policy_set_from_dict( + { + "root_uri": "viking://user/u/memories/experiences", + "policies": [ + { + "name": "booking_policy", + "uri": "viking://user/u/memories/experiences/booking_policy.md", + "version": 2, + "status": "production", + "content": "Always verify duplicates before cancellation.", + "metadata": {"domain": "booking"}, + } + ], + "metadata": {"snapshot": "remote"}, + } + ) + + assert policy_set.root_uri == "viking://user/u/memories/experiences" + assert policy_set.metadata == {"snapshot": "remote"} + assert len(policy_set.policies) == 1 + policy = policy_set.policies[0] + assert policy.name == "booking_policy" + assert policy.uri == "viking://user/u/memories/experiences/booking_policy.md" + assert policy.version == 2 + assert policy.status == "production" + assert policy.content == "Always verify duplicates before cancellation." + assert policy.metadata == {"domain": "booking"} + + +def test_tau2_rollout_messages_use_completed_structured_tool_parts(): + from benchmark.tau2.train.rollout_executor import _build_rollout_messages + from openviking.message import TextPart, ToolPart + + rollout_messages = _build_rollout_messages( + system_prompt="policy", + user_prompt="user request", + tools_used=[ + { + "tool_name": "get_user_details", + "args": '{"user_id": "emma_kim_9957"}', + "result": '{"membership": "gold"}', + } + ], + final_content="done", + evaluation_result=None, + reward=1.0, + ) + + assert isinstance(rollout_messages[0].parts[0], TextPart) + assert rollout_messages[0].parts[0].text.startswith("system:\npolicy") + + tool_message = rollout_messages[2] + assert tool_message.role == "user" + assert isinstance(tool_message.parts[0], ToolPart) + assert tool_message.parts[0].tool_status == "completed" + assert tool_message.parts[0].tool_input == {"user_id": "emma_kim_9957"} + assert tool_message.parts[0].tool_output == '{"membership": "gold"}' + assert not any( + isinstance(part, TextPart) and "tool-call:" in part.text + for message in rollout_messages + for part in message.parts + ) + assert not any( + isinstance(part, ToolPart) and part.tool_status == "running" + for message in rollout_messages + for part in message.parts + ) + + +def test_tau2_communicate_with_user_renders_as_dialogue(): + from benchmark.tau2.train.rollout_executor import _build_rollout_messages + from openviking.message import TextPart, ToolPart + + rollout_messages = _build_rollout_messages( + system_prompt="policy", + user_prompt="user request", + tools_used=[ + { + "tool_name": "communicate_with_user", + "args": {"content": "Could you provide your user ID?"}, + "result": "Sure, it is emma_kim_9957.", + } + ], + final_content=None, + evaluation_result=None, + reward=1.0, + ) + + assert rollout_messages[2].role == "assistant" + assert isinstance(rollout_messages[2].parts[0], TextPart) + assert rollout_messages[2].parts[0].text == "Could you provide your user ID?" + assert rollout_messages[3].role == "user" + assert isinstance(rollout_messages[3].parts[0], TextPart) + assert rollout_messages[3].parts[0].text == "Sure, it is emma_kim_9957." + assert not any( + isinstance(part, ToolPart) and part.tool_name == "communicate_with_user" + for message in rollout_messages + for part in message.parts + ) + + +def test_tau2_rollout_messages_omit_empty_final_after_done(): + from benchmark.tau2.train.rollout_executor import _build_rollout_messages + + rollout_messages = _build_rollout_messages( + system_prompt="policy", + user_prompt="user request", + tools_used=[{"tool_name": "done", "args": "{}", "result": "Task Terminated"}], + final_content=None, + evaluation_result=None, + reward=1.0, + ) + + assert "tau2-final" not in {message.id for message in rollout_messages} + assert rollout_messages[-1].id == "tau2-reward" + + +def test_tau2_reward_info_is_json_safe_in_rollout_messages_and_evaluation(): + import json + + from tau2.data_model.simulation import RewardInfo, RewardType + + from benchmark.tau2.train.rollout_executor import _build_rollout_messages, _tau2_evaluation + + reward_info = RewardInfo( + reward=1.0, + reward_basis=[RewardType.DB], + reward_breakdown={RewardType.DB: 1.0}, + ) + + rollout_messages = _build_rollout_messages( + system_prompt="policy", + user_prompt="user request", + tools_used=[], + final_content="done", + evaluation_result=reward_info, + reward=1.0, + ) + evaluation = _tau2_evaluation(reward=1.0, evaluation_result=reward_info) + + reward_message = rollout_messages[-1].content + assert "'reward': 1.0" not in reward_message + assert '"reward": 1.0' in reward_message + assert '"reward_basis": ["DB"]' in evaluation.feedback[0] + json.dumps(evaluation.metadata, sort_keys=True) + + +def test_tau2_litellm_generate_rate_limit_retry_patch(monkeypatch): + import benchmark.tau2.common.tau2_env.tau2_environment as tau2_environment + + calls = {"count": 0} + sleeps = [] + + def fake_generate(): + calls["count"] += 1 + if calls["count"] < 5: + raise RuntimeError("TPM (Tokens Per Minute) limit of the model is exceeded") + return "ok" + + class FakeLLMUtils: + generate = staticmethod(fake_generate) + + class FakeUserSimulator: + generate = staticmethod(fake_generate) + + modules = { + "tau2.utils.llm_utils": FakeLLMUtils, + "tau2.user.user_simulator": FakeUserSimulator, + } + + def fake_import_module(name): + if name in modules: + return modules[name] + raise ImportError(name) + + monkeypatch.setattr(tau2_environment.importlib, "import_module", fake_import_module) + monkeypatch.setattr(tau2_environment, "_tau2_rate_limit_retry_delay", lambda attempt: attempt) + monkeypatch.setattr(tau2_environment.time, "sleep", lambda delay: sleeps.append(delay)) + + tau2_environment._install_tau2_litellm_rate_limit_retry() + + assert FakeLLMUtils.generate() == "ok" + assert calls["count"] == 5 + assert sleeps == [1, 2, 3, 4] + assert FakeUserSimulator.generate is FakeLLMUtils.generate + + +def test_tau2_litellm_generate_retry_patch_does_not_retry_non_rate_limit(monkeypatch): + import benchmark.tau2.common.tau2_env.tau2_environment as tau2_environment + + calls = {"count": 0} + + def fake_generate(): + calls["count"] += 1 + raise RuntimeError("AuthenticationError Unauthorized") + + class FakeLLMUtils: + generate = staticmethod(fake_generate) + + def fake_import_module(name): + if name == "tau2.utils.llm_utils": + return FakeLLMUtils + raise ImportError(name) + + monkeypatch.setattr(tau2_environment.importlib, "import_module", fake_import_module) + + def fail_on_sleep(_delay): + raise AssertionError("unexpected sleep") + + monkeypatch.setattr(tau2_environment.time, "sleep", fail_on_sleep) + + tau2_environment._install_tau2_litellm_rate_limit_retry() + + with pytest.raises(RuntimeError, match="AuthenticationError"): + FakeLLMUtils.generate() + assert calls["count"] == 1 + + +def test_tau2_native_env_reward_handles_required_id_and_tool_call_ids(monkeypatch): + import benchmark.tau2.common.tau2_env.tau2_environment as tau2_environment + from benchmark.tau2.common.tau2_env.tau2_environment import Tau2BenchEnv + + monkeypatch.setattr(tau2_environment, "AgentGymEnv", None) + env = Tau2BenchEnv("airline", "1") + env.reset() + env.tool_call("get_user_details", {"user_id": "raj_sanchez_7340"}) + env.tool_call("get_reservation_details", {"reservation_id": "Q69X3R"}) + + reward, evaluation = env._impl._get_reward() + + assert reward == 1.0 + assert evaluation.reward == 1.0 + + +def test_tau2_native_env_records_communication_as_assistant_text(monkeypatch): + import benchmark.tau2.common.tau2_env.tau2_environment as tau2_environment + from benchmark.tau2.common.tau2_env.tau2_environment import Tau2BenchEnv + + monkeypatch.setattr(tau2_environment, "AgentGymEnv", None) + env = Tau2BenchEnv("airline", "3") + env.reset() + env.tool_call("communicate_with_user", {"content": "You may bring 4 suitcases."}) + + reward, evaluation = env._impl._get_reward() + + assert reward == 1.0 + assert evaluation.communicate_checks[0].met is True + + +def test_tau2_final_answer_is_appended_for_native_evaluation(monkeypatch): + import benchmark.tau2.common.tau2_env.tau2_environment as tau2_environment + from benchmark.tau2.common.tau2_env.tau2_environment import Tau2BenchEnv + from benchmark.tau2.train.rollout_executor import _append_final_answer_for_tau2_evaluation + + monkeypatch.setattr(tau2_environment, "AgentGymEnv", None) + env = Tau2BenchEnv("airline", "3") + env.reset() + _append_final_answer_for_tau2_evaluation(env, "You may bring 4 suitcases.") + + reward, evaluation = env._impl._get_reward() + + assert reward == 1.0 + assert evaluation.communicate_checks[0].met is True + + +def test_tau2_configure_tools_removes_only_openviking_tools(): + from benchmark.tau2.train.rollout_executor import _configure_tools + + class FakeTools: + def __init__(self): + self.tool_names = [ + "read_file", + "openviking_search", + "openviking_memory_commit", + "web_search", + ] + self.unregistered = [] + self.registered = [] + + def unregister(self, name): + self.unregistered.append(name) + self.tool_names.remove(name) + + def register(self, tool): + self.registered.append(tool.name) + + class FakeAgent: + def __init__(self): + self.tools = FakeTools() + + class FakeProvider: + def list_openai_tools(self): + return [ + { + "type": "function", + "function": { + "name": "get_user_details", + "description": "get user", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def call_tool(self, name, args): + return "ok" + + agent = FakeAgent() + + _configure_tools(agent, FakeProvider(), keep_default_tools=True) + + assert agent.tools.unregistered == ["openviking_search", "openviking_memory_commit"] + assert agent.tools.tool_names == ["read_file", "web_search"] + assert agent.tools.registered == [ + "search_experience", + "read_experience", + "get_user_details", + ] + + +def test_tau2_rollout_backend_factory_defaults_to_native(): + from benchmark.tau2.train.rollout_executor import ( + NativeTau2RolloutExecutor, + make_tau2_rollout_executor, + normalize_tau2_rollout_backend, + ) + + executor = make_tau2_rollout_executor( + options={"keep_default_tools": False, "max_iterations": 7}, + concurrency=3, + ) + + assert normalize_tau2_rollout_backend(None) == "native" + assert isinstance(executor, NativeTau2RolloutExecutor) + assert executor.concurrency == 3 + assert executor.memory_enabled is False + assert executor.max_steps == 7 + assert executor.show_progress is False + + +def test_tau2_native_rollout_resolves_non_empty_llms(monkeypatch): + from benchmark.tau2.train.rollout_executor_native import ( + NativeTau2RolloutExecutor, + _resolve_llm_runtime_config, + ) + + monkeypatch.delenv("TAU2_AGENT_LLM", raising=False) + monkeypatch.delenv("TAU2_USER_LLM", raising=False) + + agent_llm, agent_args, user_llm, user_args = _resolve_llm_runtime_config( + NativeTau2RolloutExecutor( + agent_llm_args={"temperature": 0.2}, + user_llm_args={"top_p": 0.9}, + ) + ) + + assert agent_llm + assert user_llm + assert agent_args["temperature"] == 0.2 + assert user_args["temperature"] == 0.0 + assert user_args["top_p"] == 0.9 + + +def test_tau2_native_rollout_uses_env_llm_when_options_omit_model(monkeypatch): + from benchmark.tau2.train.rollout_executor_native import ( + NativeTau2RolloutExecutor, + _resolve_llm_runtime_config, + ) + + monkeypatch.setenv("TAU2_AGENT_LLM", "openai/test-agent") + monkeypatch.setenv("TAU2_USER_LLM", "openai/test-user") + + agent_llm, _agent_args, user_llm, _user_args = _resolve_llm_runtime_config( + NativeTau2RolloutExecutor() + ) + + assert agent_llm == "openai/test-agent" + assert user_llm == "openai/test-user" + + +def test_tau2_rollout_backend_factory_selects_vikingbot(monkeypatch): + import benchmark.tau2.train.rollout_executor as module + + created = {} + + class FakeVikingBotExecutor: + def __init__(self, **kwargs): + created.update(kwargs) + + monkeypatch.setattr(module, "VikingBotTau2RolloutExecutor", FakeVikingBotExecutor) + + executor = module.make_tau2_rollout_executor( + backend="vikingbot", + options={ + "config_path": "/tmp/ov.conf", + "max_iterations": 9, + }, + concurrency=2, + rollout_language="zh", + ) + + assert isinstance(executor, FakeVikingBotExecutor) + assert created == { + "config_path": "/tmp/ov.conf", + "concurrency": 2, + "keep_default_tools": True, + "max_iterations": 9, + "rollout_language": "zh", + } + + +def test_tau2_service_rollout_backend_option_overrides_default(monkeypatch): + import benchmark.tau2.train.service_app as service_app + + calls = [] + + def fake_create_dataset_service_app(**kwargs): + calls.append(kwargs) + return kwargs + + class FakeExecutor: + pass + + def fake_make_tau2_rollout_executor(**kwargs): + calls.append({"factory": kwargs}) + return FakeExecutor() + + monkeypatch.setattr(service_app, "create_dataset_service_app", fake_create_dataset_service_app) + monkeypatch.setattr(service_app, "make_tau2_rollout_executor", fake_make_tau2_rollout_executor) + + app = service_app.create_app(rollout_backend="native") + executor = app["make_rollout_executor"]({"rollout_backend": "vikingbot", "max_iterations": 5}) + + assert isinstance(executor, FakeExecutor) + assert calls[-1]["factory"]["backend"] == "vikingbot" + assert calls[-1]["factory"]["options"]["max_iterations"] == 5 + assert calls[-1]["factory"]["options"]["show_progress"] is False + + app["make_rollout_executor"]({"rollout_backend": "native", "show_progress": True}) + assert calls[-1]["factory"]["options"]["show_progress"] is True + + +@pytest.mark.asyncio +async def test_tau2_vikingbot_rollout_runs_on_current_event_loop(): + from benchmark.tau2.train.rollout_executor_vikingbot import VikingBotTau2RolloutExecutor + + expected_loop = asyncio.get_running_loop() + expected_thread = threading.get_ident() + observed = {} + + class FakeVikingBotExecutor(VikingBotTau2RolloutExecutor): + async def _execute_one_async(self, case, context): + del context + observed["loop"] = asyncio.get_running_loop() + observed["thread"] = threading.get_ident() + await asyncio.sleep(0) + return case.name + + executor = FakeVikingBotExecutor() + + result = await executor._execute_one( + _case(), + ExecutionContext(policy_snapshot_id="snapshot", metadata={}), + ) + + assert result == "case-1" + assert observed["loop"] is expected_loop + assert observed["thread"] == expected_thread + + +@pytest.mark.asyncio +async def test_tau2_prepare_experience_loader_skill_writes_required_skill(tmp_path): + import benchmark.tau2.train.rollout_executor_vikingbot as module + + class FakeSandbox: + def __init__(self): + self.writes = [] + + async def write_file(self, path, content): + self.writes.append((path, content)) + target = tmp_path / path + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(content, encoding="utf-8") + + fake_sandbox = FakeSandbox() + + class FakeSandboxManager: + def get_workspace_path(self, session_key): + return tmp_path + + def to_workspace_id(self, session_key): + return "workspace" + + async def get_sandbox(self, session_key): + return fake_sandbox + + class FakeAgent: + sandbox_manager = FakeSandboxManager() + context = SimpleNamespace(workspace=tmp_path) + + context_builder = await module._prepare_experience_loader_skill( + agent=FakeAgent(), + session_key=SimpleNamespace(), + ) + + skill_path = tmp_path / "skills" / "experience_loader" / "SKILL.md" + content = skill_path.read_text(encoding="utf-8") + assert context_builder.workspace == tmp_path + assert "name: experience_loader" in content + assert "search_experience" in content + assert "read_experience" in content + assert fake_sandbox.writes + assert fake_sandbox.writes[0][0] == "skills/experience_loader/SKILL.md" + assert context_builder.latest_experience_loader_skill_content == content + + +@pytest.mark.asyncio +async def test_tau2_experience_loader_skill_is_required_with_relative_read_path(tmp_path): + from vikingbot.config.schema import SessionKey + + import benchmark.tau2.train.rollout_executor_vikingbot as module + + module._write_experience_loader_files( + workspace_path=tmp_path, + skill_content="# experience_loader\n\nUse search_experience then read_experience.", + ) + + from vikingbot.agent.context import ContextBuilder + + context_builder = ContextBuilder(tmp_path, eval=True) + system_prompt = await context_builder.build_system_prompt( + SessionKey(type="cli", channel_id="tau2", chat_id="case"), + ov_tools_enable=False, + ) + + assert "Required skill: before taking any task action" in system_prompt + assert "`skills/experience_loader/SKILL.md`" in system_prompt + assert "skills/experience_loader/SKILL.md" in system_prompt + assert f"{tmp_path}" not in system_prompt + + +@pytest.mark.asyncio +async def test_tau2_vikingbot_blocking_setup_and_reward_are_offloaded(monkeypatch): + import benchmark.tau2.train.rollout_executor_vikingbot as module + from benchmark.tau2.train.rollout_executor_vikingbot import VikingBotTau2RolloutExecutor + + event_loop_thread = threading.get_ident() + calls = [] + + class FakeEnv: + def _get_reward(self): + calls.append(("reward", threading.get_ident())) + return 1.0, {"ok": True} + + class FakeTau2BenchToolProvider: + def __init__(self, domain, task_id, data_root=None): + self.domain = domain + self.task_id = task_id + self.data_root = data_root + self.env = FakeEnv() + self.policy = "policy" + self.user_query = "user query" + + def reset(self): + calls.append(("reset", threading.get_ident())) + + def list_openai_tools(self): + return [] + + class FakeAgent: + def __init__(self): + calls.append(("build_agent", threading.get_ident())) + + async def fake_run_agent(**kwargs): + calls.append(("run_agent", threading.get_ident())) + calls.append(("case_lookup", kwargs.get("case_lookup"))) + return "final", None, [], {}, 1, None, None, None + + monkeypatch.setattr(module, "_tool_provider_cls", lambda: FakeTau2BenchToolProvider) + monkeypatch.setattr(module, "_build_agent", lambda *args, **kwargs: FakeAgent()) + monkeypatch.setattr(module, "_configure_tools", lambda *args, **kwargs: None) + monkeypatch.setattr(module, "_run_agent", fake_run_agent) + + case = Case( + name="tau2_case", + task_signature="tau2:airline:train:0", + input={ + "domain": "airline", + "split": "train", + "task_id": "0", + "task_no": 0, + "data_split": "airline_train", + }, + rubric=Rubric(name="rubric", description="", criteria=[]), + ) + executor = VikingBotTau2RolloutExecutor() + + rollout = await executor._execute_one( + case, + ExecutionContext(policy_snapshot_id="snapshot", metadata={}), + ) + + assert rollout.metadata["reward"] == 1.0 + call_values = dict(calls) + assert call_values["case_lookup"] == { + "benchmark": "tau2", + "strict": True, + "case_names": ["tau2_case", "tau2_airline_train_0"], + "domain": "airline", + "split": "train", + "data_split": "airline_train", + "task_no": 0, + "task_id": "0", + "case_name": "tau2_case", + "task_signature": "tau2:airline:train:0", + "original_case_name": None, + "expected_fields": { + "input.domain": "airline", + "input.split": "train", + "input.data_split": "airline_train", + "input.task_no": 0, + "input.task_id": "0", + }, + } + call_threads = call_values + assert call_threads["reset"] != event_loop_thread + assert call_threads["build_agent"] != event_loop_thread + assert call_threads["reward"] != event_loop_thread + assert call_threads["run_agent"] == event_loop_thread + + +@pytest.mark.asyncio +async def test_tau2_run_agent_force_loads_experience_loader_skill_before_task_actions(monkeypatch): + from vikingbot.providers.base import LLMResponse, ToolCallRequest + + import benchmark.tau2.train.rollout_executor_vikingbot as module + + observed = {} + real_imports = module._vikingbot_imports() + + class FakeSandbox: + content = "" + + async def read_file(self, path): + observed.setdefault("sandbox_reads", []).append(path) + return self.content + + async def write_file(self, path, content): + observed.setdefault("sandbox_writes", []).append((path, content)) + type(self).content = content + + class FakeSandboxManager: + def get_workspace_path(self, session_key): + return Path("/tmp/fake-workspace") + + def to_workspace_id(self, session_key): + return "workspace" + + async def get_sandbox(self, session_key): + return FakeSandbox() + + class FakeContextBuilder: + def __init__(self, workspace, *, sandbox_manager=None, eval=False, **kwargs): + self.workspace = workspace + self.sandbox_manager = sandbox_manager + self.latest_experience_loader_skill_content = "" + + async def build_messages(self, **kwargs): + return [ + {"role": "system", "content": "ctx system"}, + {"role": "user", "content": kwargs["current_message"]}, + ] + + def add_assistant_message( + self, messages, content, tool_calls=None, reasoning_content=None + ): + msg = {"role": "assistant", "content": content or "[tool call]"} + if tool_calls: + msg["tool_calls"] = tool_calls + if reasoning_content: + msg["reasoning_content"] = reasoning_content + messages.append(msg) + return messages + + def add_tool_result(self, messages, tool_call_id, tool_name, result): + messages.append( + { + "role": "tool", + "tool_call_id": tool_call_id, + "name": tool_name, + "content": result, + } + ) + return messages + + class FakeProvider: + async def chat(self, messages, tools=None, **kwargs): + observed["llm_messages"] = list(messages) + return LLMResponse( + content=None, + tool_calls=[ToolCallRequest("call-1", "done", {}, 0)], + ) + + async def chat_stream(self, **kwargs): + from vikingbot.providers.base import LLMStreamEvent + + yield LLMStreamEvent(type="response", response=await self.chat(**kwargs)) + + def get_default_model(self): + return "fake" + + class FakeAgent: + def __init__(self): + from vikingbot.agent.tools.filesystem import ReadFileTool + from vikingbot.agent.tools.registry import ToolRegistry + + self.sandbox_manager = FakeSandboxManager() + self.context = FakeContextBuilder(Path("/tmp/fake-workspace")) + self.tools = ToolRegistry() + self.tools.register(ReadFileTool()) + self.tools.register(_DoneTool()) + self.provider = FakeProvider() + self.model = "fake" + self.temperature = None + self.max_iterations = 1 + + _chat_with_stream_events = real_imports["AgentLoop"]._chat_with_stream_events + _run_agent_loop = real_imports["AgentLoop"]._run_agent_loop + + class _DoneTool: + @property + def name(self): + return "done" + + @property + def description(self): + return "done" + + @property + def parameters(self): + return {"type": "object", "properties": {}} + + def to_schema(self): + return { + "type": "function", + "function": { + "name": self.name, + "description": self.description, + "parameters": self.parameters, + }, + } + + def validate_params(self, params): + return [] + + async def execute(self, tool_context, **kwargs): + return "" + + monkeypatch.setattr( + module, + "_vikingbot_imports", + lambda: {**real_imports, "ContextBuilder": FakeContextBuilder}, + ) + + result = await module._run_agent( + agent=FakeAgent(), + system_prompt="tau2 policy", + user_prompt="user query", + session_key=SimpleNamespace(safe_name=lambda: "session"), + sender_id="tau2_user", + keep_default_tools=True, + case_lookup={"benchmark": "tau2", "strict": True, "case_name": "case"}, + ) + + tools_used = result[2] + messages = observed["llm_messages"] + read_call_index = next( + i + for i, msg in enumerate(messages) + if msg.get("role") == "assistant" and "read_file" in str(msg.get("tool_calls")) + ) + tool_result_index = next( + i + for i, msg in enumerate(messages) + if msg.get("role") == "tool" and msg.get("name") == "read_file" + ) + + assert observed["sandbox_writes"][0][0] == "skills/experience_loader/SKILL.md" + assert observed["sandbox_reads"] == ["skills/experience_loader/SKILL.md"] + assert read_call_index < tool_result_index + assert "search_experience" in messages[tool_result_index]["content"] + assert "read_experience" in messages[tool_result_index]["content"] + assert tools_used[0]["tool_name"] == "read_file" + assert tools_used[0]["required_skill"] == "experience_loader" diff --git a/tests/session/train/test_train_components.py b/tests/session/train/test_train_components.py new file mode 100644 index 0000000000..de346e0768 --- /dev/null +++ b/tests/session/train/test_train_components.py @@ -0,0 +1,841 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 + +from __future__ import annotations + +from typing import Any + +import pytest +from test_fakes import fake_request_context + +from openviking.session.memory.dataclass import MemoryFile, StoredLink +from openviking.session.memory.utils.memory_file_utils import MemoryFileUtils +from openviking.session.train import ( + ContentHashPolicySnapshotter, + DryRunPolicyUpdater, + Experience, + ExperienceSet, + ExperienceSetLoader, + MemoryFilePolicyUpdater, + PatchMergePolicyOptimizer, + PatchMergePolicyOptimizerContext, + PatchSemanticGradient, + PolicyUpdatePlan, +) + + +class FakeVikingFS: + def __init__(self, files: dict[str, str]): + self.files = files + self.rm_lock_handles = [] + + async def ls(self, uri: str, output: str = "original", ctx=None, **kwargs): + del kwargs + assert output == "original" + prefix = uri.rstrip("/") + "/" + return [ + { + "name": path.removeprefix(prefix), + "uri": path, + "isDir": False, + } + for path in sorted(self.files) + if path.startswith(prefix) and "/" not in path.removeprefix(prefix) + ] + + async def read_file(self, uri: str, ctx=None): + return self.files[uri] + + async def write_file(self, uri: str, content: str, ctx=None): + self.files[uri] = content + + async def rm(self, uri: str, recursive: bool = False, ctx=None, lock_handle=None): + del recursive, ctx + self.rm_lock_handles.append(lock_handle) + self.files.pop(uri, None) + return {"estimated_deleted_count": 1} + + +class FakeVikingDB: + def __init__(self): + self.embedding_messages = [] + + async def enqueue_embedding_msg(self, embedding_msg): + self.embedding_messages.append(embedding_msg) + return True + + +def _experience_set() -> ExperienceSet: + return ExperienceSet( + root_uri="viking://user/u/memories/experiences", + policies=[ + Experience( + name="booking_duplicate_handling", + uri="viking://user/u/memories/experiences/booking_duplicate_handling.md", + version=1, + status="production", + content="content", + ) + ], + ) + + +def _memory_file( + *, + name: str, + uri: str | None, + content: str, + version: int | None = 1, + status: str = "production", +) -> MemoryFile: + fields: dict[str, Any] = { + "memory_type": "experiences", + "experience_name": name, + "status": status, + } + if version is not None: + fields["version"] = version + return MemoryFile( + uri=uri, + content=content, + memory_type="experiences", + extra_fields=fields, + ) + + +def _patch_gradient( + *, + name: str = "booking_duplicate_handling", + uri: str | None = "viking://user/u/memories/experiences/booking_duplicate_handling.md", + before: str | None = "content", + after: str = "new content", + base_version: int | None = 1, + rationale: str = "r", + links: list[StoredLink] | None = None, + confidence: float = 0.8, + metadata: dict[str, Any] | None = None, +) -> PatchSemanticGradient: + return PatchSemanticGradient( + before_file=( + _memory_file(name=name, uri=uri, content=before, version=base_version) + if before is not None + else None + ), + after_file=_memory_file(name=name, uri=uri, content=after, version=base_version), + base_version=base_version, + rationale=rationale, + links=links or [ + StoredLink( + from_uri=uri or "", + to_uri="viking://user/u/memories/trajectories/traj1.md", + link_type="derived_from", + weight=1.0, + ) + ], + confidence=confidence, + metadata=metadata or {}, + ) + + +def _plan_from_gradient(gradient: PatchSemanticGradient) -> PolicyUpdatePlan: + return PolicyUpdatePlan( + items=[ + _plan_item_from_gradient(gradient), + ] + ) + + +def _plan_item_from_gradient(gradient: PatchSemanticGradient): + from openviking.session.train import PolicyPlanItem + + return PolicyPlanItem( + kind="upsert", + memory_type="experiences", + target_name=gradient.target_name, + target_uri=gradient.target_uri, + before_content=( + gradient.before_file.plain_content() if gradient.before_file is not None else None + ), + after_content=gradient.after_file.plain_content(), + base_version=gradient.base_version, + confidence=gradient.confidence, + links=list(gradient.links), + metadata={"rationale": gradient.rationale}, + ) + + +def _delete_plan(*, uri: str, before_content: str = "content") -> PolicyUpdatePlan: + from openviking.session.train import PolicyPlanItem + + return PolicyUpdatePlan( + items=[ + PolicyPlanItem( + kind="delete", + memory_type="experiences", + target_name="booking_duplicate_handling", + target_uri=uri, + before_content=before_content, + after_content=None, + base_version=1, + confidence=0.8, + links=[ + StoredLink( + from_uri=uri, + to_uri="viking://user/u/memories/trajectories/traj1.md", + link_type="derived_from", + weight=1.0, + ) + ], + metadata={"rationale": "delete duplicate experience"}, + ) + ] + ) + + +@pytest.mark.asyncio +async def test_experience_set_loader_reads_memory_files(): + root = "viking://user/u/memories/experiences" + fs = FakeVikingFS( + { + f"{root}/booking_duplicate_handling.md": '## Situation\n- test\n\n', + f"{root}/.overview.md": "hidden", + } + ) + + ctx = fake_request_context() + loaded = await ExperienceSetLoader(viking_fs=fs).load(root, ctx=ctx) + + assert loaded.root_uri == root + assert loaded.viking_fs is fs + assert loaded.request_context is ctx + assert len(loaded.policies) == 1 + policy = loaded.policies[0] + assert policy.name == "booking_duplicate_handling" + assert policy.version == 3 + assert policy.status == "staging" + assert policy.content == "## Situation\n- test" + assert policy.metadata["memory_type"] == "experiences" + + +@pytest.mark.asyncio +async def test_experience_set_loader_requires_request_context(): + root = "viking://user/u/memories/experiences" + fs = FakeVikingFS({}) + + with pytest.raises(ValueError, match="requires request_context ctx"): + await ExperienceSetLoader(viking_fs=fs).load(root) + + +@pytest.mark.asyncio +async def test_content_hash_snapshotter_is_deterministic(): + snapshotter = ContentHashPolicySnapshotter() + policy_set = _experience_set() + + first = await snapshotter.snapshot(policy_set) + second = await snapshotter.snapshot(policy_set) + + assert first == second + assert first.startswith("policy-snapshot:") + + +@pytest.mark.asyncio +async def test_dry_run_policy_updater_does_not_mutate_policy_set(): + policy_set = _experience_set() + plan = PolicyUpdatePlan(metadata={"hello": "world"}) + + result = await DryRunPolicyUpdater().apply(plan, policy_set) + + assert result.updated_policy_set is policy_set + assert result.written_uris == [] + assert result.deleted_uris == [] + assert result.metadata["dry_run"] is True + assert result.metadata["simulated"] is True + assert result.metadata["plan"] == {"hello": "world"} + + +@pytest.mark.asyncio +async def test_dry_run_policy_updater_simulates_patch_plan_items(): + policy_set = _experience_set() + gradient = _patch_gradient(uri=policy_set.policies[0].uri, before="content", after="new content") + plan = _plan_from_gradient(gradient) + + result = await DryRunPolicyUpdater().apply(plan, policy_set) + + assert result.updated_policy_set is not policy_set + assert result.updated_policy_set.policies[0].content == "new content" + assert result.updated_policy_set.policies[0].version == 2 + assert result.written_uris == [] + assert result.metadata["dry_run"] is True + assert result.metadata["simulated"] is True + + +@pytest.mark.asyncio +async def test_dry_run_policy_updater_simulates_delete_plan_items(): + policy_set = _experience_set() + plan = _delete_plan(uri=policy_set.policies[0].uri) + + result = await DryRunPolicyUpdater().apply(plan, policy_set) + + assert result.updated_policy_set is not policy_set + assert result.updated_policy_set.policies == [] + assert result.written_uris == [] + assert result.deleted_uris == [] + assert result.metadata["dry_run"] is True + assert result.metadata["simulated"] is True + + +@pytest.mark.asyncio +async def test_memory_file_policy_updater_writes_experience_files(): + policy_set = _experience_set() + fs = FakeVikingFS({}) + gradient = _patch_gradient(uri=policy_set.policies[0].uri, before="content", after="new content") + plan = _plan_from_gradient(gradient) + + result = await MemoryFilePolicyUpdater(viking_fs=fs).apply( + plan, + policy_set, + fake_request_context(), + ) + + assert result.errors == [] + assert result.written_uris == [policy_set.policies[0].uri] + written = fs.files[policy_set.policies[0].uri] + assert written.startswith("new content") + assert '"memory_type": "experiences"' in written + assert '"experience_name": "booking_duplicate_handling"' in written + assert '"version": 2' in written + + +@pytest.mark.asyncio +async def test_memory_file_policy_updater_vectorizes_written_experience_files(): + policy_set = _experience_set() + fs = FakeVikingFS({}) + vikingdb = FakeVikingDB() + gradient = _patch_gradient(uri=policy_set.policies[0].uri, before="content", after="new content") + plan = _plan_from_gradient(gradient) + + from openviking.server.identity import RequestContext, Role + from openviking_cli.session.user_id import UserIdentifier + + result = await MemoryFilePolicyUpdater(viking_fs=fs, vikingdb=vikingdb).apply( + plan, + policy_set, + RequestContext(user=UserIdentifier("default", "u"), role=Role.USER), + ) + + assert result.errors == [] + assert result.written_uris == [policy_set.policies[0].uri] + assert len(vikingdb.embedding_messages) == 1 + embedding_msg = vikingdb.embedding_messages[0] + assert embedding_msg.context_data["uri"] == policy_set.policies[0].uri + assert embedding_msg.context_data["context_type"] == "memory" + assert "new content" in embedding_msg.message + + +@pytest.mark.asyncio +async def test_memory_file_policy_updater_writes_v2_compatible_source_trajectory_links(): + policy_set = _experience_set() + exp_uri = policy_set.policies[0].uri + traj_uri = "viking://user/u/memories/trajectories/booking_duplicate.md" + fs = FakeVikingFS( + { + traj_uri: MemoryFileUtils.write( + MemoryFile( + uri=traj_uri, + content="trajectory content", + memory_type="trajectories", + extra_fields={ + "memory_type": "trajectories", + "trajectory_name": "booking_duplicate", + }, + ) + ) + } + ) + gradient = _patch_gradient( + uri=exp_uri, + before="content", + after="new content", + links=[ + StoredLink( + from_uri=exp_uri, + to_uri=traj_uri, + link_type="derived_from", + weight=1.0, + ) + ], + ) + plan = _plan_from_gradient(gradient) + + result = await MemoryFilePolicyUpdater(viking_fs=fs).apply(plan, policy_set) + + assert result.errors == [] + exp_mf = MemoryFileUtils.read(fs.files[exp_uri], uri=exp_uri) + assert any( + link.get("from_uri") == exp_uri + and link.get("to_uri") == traj_uri + and link.get("link_type") == "derived_from" + and link.get("match_text") is None + and link.get("description") == "" + for link in exp_mf.links + ) + + traj_mf = MemoryFileUtils.read(fs.files[traj_uri], uri=traj_uri) + assert any( + link.get("from_uri") == exp_uri + and link.get("to_uri") == traj_uri + and link.get("link_type") == "derived_from" + and link.get("match_text") is None + and link.get("description") == "" + for link in traj_mf.backlinks + ) + + +@pytest.mark.asyncio +async def test_memory_file_policy_updater_deletes_experience_files(): + policy_set = _experience_set() + uri = policy_set.policies[0].uri + fs = FakeVikingFS({uri: "content"}) + plan = _delete_plan(uri=uri) + lock_handle = object() + + result = await MemoryFilePolicyUpdater(viking_fs=fs).apply( + plan, + policy_set, + transaction_handle=lock_handle, + ) + + assert result.errors == [] + assert result.written_uris == [] + assert result.deleted_uris == [uri] + assert result.updated_policy_set.policies == [] + assert uri not in fs.files + assert fs.rm_lock_handles == [lock_handle] + + +@pytest.mark.asyncio +async def test_memory_file_policy_updater_detects_base_content_mismatch(): + policy_set = _experience_set() + fs = FakeVikingFS({}) + gradient = _patch_gradient( + uri=policy_set.policies[0].uri, + before="stale content", + after="new content", + ) + plan = _plan_from_gradient(gradient) + + result = await MemoryFilePolicyUpdater(viking_fs=fs).apply(plan, policy_set) + + assert result.written_uris == [] + assert result.errors == [ + "base content mismatch for booking_duplicate_handling: expected gradient before_content" + ] + assert policy_set.policies[0].uri not in fs.files + + +@pytest.mark.asyncio +async def test_patch_merge_policy_optimizer_runs_patch_merge_extract_loop(monkeypatch): + from openviking.session.memory.dataclass import ( + MemoryFile, + ResolvedOperation, + ResolvedOperations, + ) + + policy_set = _experience_set() + gradient = _patch_gradient( + uri=policy_set.policies[0].uri, + before="stale content", + after="merged content", + ) + captured = {} + + class FakeExtractLoop: + def __init__(self, **kwargs): + captured.update(kwargs) + + async def run(self): + provider = captured["context_provider"] + captured["prefetch_messages"] = await provider.prefetch() + return ( + ResolvedOperations( + upsert_operations=[ + ResolvedOperation( + old_memory_file_content=MemoryFile( + uri=policy_set.policies[0].uri, + content="content", + memory_type="experiences", + extra_fields={ + "experience_name": "booking_duplicate_handling", + "version": 1, + }, + ), + memory_fields={ + "experience_name": "booking_duplicate_handling", + "content": "merged content", + }, + memory_type="experiences", + uris=[policy_set.policies[0].uri], + ) + ], + delete_file_contents=[], + errors=[], + ), + [], + ) + + monkeypatch.setattr("openviking.session.train.components.policy_optimizer.ExtractLoop", FakeExtractLoop) + + plan = await PatchMergePolicyOptimizer(viking_fs=FakeVikingFS({}), vlm=object()).plan( + [gradient], + policy_set, + PatchMergePolicyOptimizerContext(request_context=fake_request_context()), + ) + + assert plan.metadata["optimizer"] == "patch_merge" + assert plan.items[0].kind == "upsert" + assert plan.items[0].target_uri == policy_set.policies[0].uri + assert plan.items[0].before_content == "content" + assert plan.items[0].after_content == "merged content" + assert [link.to_uri for link in plan.items[0].links] == [ + "viking://user/u/memories/trajectories/traj1.md" + ] + assert captured["context_provider"].__class__.__name__ == "PatchMergeContextProvider" + assert captured["context_provider"].get_tools() == [] + assert "Patch 1" in captured["prefetch_messages"][-1]["content"] + assert " content:" in captured["prefetch_messages"][-1]["content"] + assert "-stale content" in captured["prefetch_messages"][-1]["content"] + assert "+merged content" in captured["prefetch_messages"][-1]["content"] + + +@pytest.mark.asyncio +async def test_patch_merge_policy_optimizer_merges_all_patch_gradients_once(monkeypatch): + from openviking.session.memory.dataclass import ( + ResolvedOperation, + ResolvedOperations, + ) + + policy_set = _experience_set() + root = policy_set.root_uri + gradients = [ + _patch_gradient( + name="重复预订处理", + uri=f"{root}/重复预订处理.md", + before=None, + after="核对订单后只取消重复订单", + base_version=None, + rationale="r1", + links=[ + StoredLink( + from_uri=f"{root}/重复预订处理.md", + to_uri="viking://user/u/memories/trajectories/traj1.md", + link_type="derived_from", + weight=1.0, + ) + ], + confidence=0.8, + ), + _patch_gradient( + name="处理酒店重复预订", + uri=f"{root}/处理酒店重复预订.md", + before=None, + after="识别有效订单并取消重复订单", + base_version=None, + rationale="r2", + links=[ + StoredLink( + from_uri=f"{root}/处理酒店重复预订.md", + to_uri="viking://user/u/memories/trajectories/traj2.md", + link_type="derived_from", + weight=1.0, + ) + ], + confidence=0.9, + ), + ] + captured = {"constructed": 0} + + class FakeExtractLoop: + def __init__(self, **kwargs): + captured["constructed"] += 1 + captured.update(kwargs) + + async def run(self): + provider = captured["context_provider"] + captured["prefetch_messages"] = await provider.prefetch() + return ( + ResolvedOperations( + upsert_operations=[ + ResolvedOperation( + old_memory_file_content=None, + memory_fields={ + "experience_name": "重复预订处理", + "content": "合并后的重复预订处理经验", + }, + memory_type="experiences", + uris=[f"{root}/重复预订处理.md"], + ) + ], + delete_file_contents=[], + errors=[], + ), + [], + ) + + monkeypatch.setattr("openviking.session.train.components.policy_optimizer.ExtractLoop", FakeExtractLoop) + + plan = await PatchMergePolicyOptimizer(viking_fs=FakeVikingFS({}), vlm=object()).plan( + gradients, + policy_set, + PatchMergePolicyOptimizerContext(request_context=fake_request_context()), + ) + + assert captured["constructed"] == 1 + provider = captured["context_provider"] + assert provider.required_file_uris == [ + f"{root}/重复预订处理.md", + f"{root}/处理酒店重复预订.md", + ] + assert len(provider.patches) == 2 + assert captured["prefetch_messages"][-1]["content"].count("\nPatch ") == 2 + assert plan.metadata["optimizer"] == "patch_merge" + assert plan.metadata["patch_gradient_count"] == 2 + assert len(plan.items) == 1 + assert plan.items[0].target_name == "重复预订处理" + assert [link.to_uri for link in plan.items[0].links] == [ + "viking://user/u/memories/trajectories/traj1.md", + "viking://user/u/memories/trajectories/traj2.md", + ] + assert {link.from_uri for link in plan.items[0].links} == {f"{root}/重复预订处理.md"} + + +@pytest.mark.asyncio +async def test_patch_merge_policy_optimizer_keeps_distinct_output_source_links_scoped(monkeypatch): + from openviking.session.memory.dataclass import ( + ResolvedOperation, + ResolvedOperations, + ) + + policy_set = ExperienceSet(root_uri="viking://user/u/memories/experiences", policies=[]) + root = policy_set.root_uri + gradients = [ + _patch_gradient( + name="取消资格核验", + uri=f"{root}/取消资格核验.md", + before=None, + after="取消前核验资格", + base_version=None, + links=[ + StoredLink( + from_uri=f"{root}/取消资格核验.md", + to_uri="viking://user/u/memories/trajectories/traj_cancel.md", + link_type="derived_from", + weight=1.0, + ) + ], + ), + _patch_gradient( + name="退款总额传达", + uri=f"{root}/退款总额传达.md", + before=None, + after="多笔退款后传达总额", + base_version=None, + links=[ + StoredLink( + from_uri=f"{root}/退款总额传达.md", + to_uri="viking://user/u/memories/trajectories/traj_refund.md", + link_type="derived_from", + weight=1.0, + ) + ], + ), + ] + + class FakeExtractLoop: + def __init__(self, **kwargs): + pass + + async def run(self): + return ( + ResolvedOperations( + upsert_operations=[ + ResolvedOperation( + old_memory_file_content=None, + memory_fields={ + "experience_name": "取消资格核验", + "content": "取消前核验资格", + }, + memory_type="experiences", + uris=[f"{root}/取消资格核验.md"], + ), + ResolvedOperation( + old_memory_file_content=None, + memory_fields={ + "experience_name": "退款总额传达", + "content": "多笔退款后传达总额", + }, + memory_type="experiences", + uris=[f"{root}/退款总额传达.md"], + ), + ], + delete_file_contents=[], + errors=[], + ), + [], + ) + + monkeypatch.setattr("openviking.session.train.components.policy_optimizer.ExtractLoop", FakeExtractLoop) + + plan = await PatchMergePolicyOptimizer(viking_fs=FakeVikingFS({}), vlm=object()).plan( + gradients, + policy_set, + PatchMergePolicyOptimizerContext(request_context=fake_request_context()), + ) + + links_by_name = {item.target_name: {link.to_uri for link in item.links} for item in plan.items} + assert links_by_name == { + "取消资格核验": {"viking://user/u/memories/trajectories/traj_cancel.md"}, + "退款总额传达": {"viking://user/u/memories/trajectories/traj_refund.md"}, + } + + +@pytest.mark.asyncio +async def test_patch_merge_policy_optimizer_single_canonical_output_inherits_all_source_links(monkeypatch): + from openviking.session.memory.dataclass import ( + ResolvedOperation, + ResolvedOperations, + ) + + policy_set = ExperienceSet(root_uri="viking://user/u/memories/experiences", policies=[]) + root = policy_set.root_uri + gradients = [ + _patch_gradient( + name="重复预订处理", + uri=f"{root}/重复预订处理.md", + before=None, + after="核对订单后只取消重复订单", + base_version=None, + links=[ + StoredLink( + from_uri=f"{root}/重复预订处理.md", + to_uri="viking://user/u/memories/trajectories/traj1.md", + link_type="derived_from", + weight=1.0, + ) + ], + ), + _patch_gradient( + name="处理酒店重复预订", + uri=f"{root}/处理酒店重复预订.md", + before=None, + after="识别有效订单并取消重复订单", + base_version=None, + links=[ + StoredLink( + from_uri=f"{root}/处理酒店重复预订.md", + to_uri="viking://user/u/memories/trajectories/traj2.md", + link_type="derived_from", + weight=1.0, + ) + ], + ), + ] + + class FakeExtractLoop: + def __init__(self, **kwargs): + pass + + async def run(self): + return ( + ResolvedOperations( + upsert_operations=[ + ResolvedOperation( + old_memory_file_content=None, + memory_fields={ + "experience_name": "重复预订处理", + "content": "合并后的重复预订处理经验", + }, + memory_type="experiences", + uris=[f"{root}/重复预订处理.md"], + ) + ], + delete_file_contents=[], + errors=[], + ), + [], + ) + + monkeypatch.setattr("openviking.session.train.components.policy_optimizer.ExtractLoop", FakeExtractLoop) + + plan = await PatchMergePolicyOptimizer(viking_fs=FakeVikingFS({}), vlm=object()).plan( + gradients, + policy_set, + PatchMergePolicyOptimizerContext(request_context=fake_request_context()), + ) + + assert len(plan.items) == 1 + assert {link.to_uri for link in plan.items[0].links} == { + "viking://user/u/memories/trajectories/traj1.md", + "viking://user/u/memories/trajectories/traj2.md", + } + + +@pytest.mark.asyncio +async def test_patch_merge_policy_optimizer_runs_llm_for_single_patch(monkeypatch): + from openviking.session.memory.dataclass import ( + MemoryFile, + ResolvedOperation, + ResolvedOperations, + ) + + policy_set = _experience_set() + gradient = _patch_gradient( + uri=policy_set.policies[0].uri, + before="content", + after="merged update", + ) + captured = {"constructed": False} + + class FakeExtractLoop: + def __init__(self, **kwargs): + captured["constructed"] = True + captured.update(kwargs) + + async def run(self): + return ( + ResolvedOperations( + upsert_operations=[ + ResolvedOperation( + old_memory_file_content=MemoryFile( + uri=policy_set.policies[0].uri, + content="content", + memory_type="experiences", + extra_fields={ + "experience_name": "booking_duplicate_handling", + "version": 1, + }, + ), + memory_fields={ + "experience_name": "booking_duplicate_handling", + "content": "merged update", + }, + memory_type="experiences", + uris=[policy_set.policies[0].uri], + ) + ], + delete_file_contents=[], + errors=[], + ), + [], + ) + + monkeypatch.setattr("openviking.session.train.components.policy_optimizer.ExtractLoop", FakeExtractLoop) + + plan = await PatchMergePolicyOptimizer(viking_fs=FakeVikingFS({}), vlm=object()).plan( + [gradient], + policy_set, + PatchMergePolicyOptimizerContext(request_context=fake_request_context()), + ) + + assert captured["constructed"] is True + assert plan.metadata["patch_gradient_count"] == 1 + assert plan.items[0].after_content == "merged update" diff --git a/tests/session/train/test_train_framework.py b/tests/session/train/test_train_framework.py new file mode 100644 index 0000000000..82685e1967 --- /dev/null +++ b/tests/session/train/test_train_framework.py @@ -0,0 +1,1929 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 + +from __future__ import annotations + +import asyncio +import json +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import pytest +from test_fakes import InMemoryAGFS, fake_request_context + +from openviking.message import Message, TextPart +from openviking.session.memory.dataclass import StoredLink +from openviking.session.train import ( + Case, + Experience, + ExperienceSet, + ListCaseLoader, + NoopPipelineLifecycleHook, + OfflinePolicyOptimizationPipeline, + PipelineContext, + PipelineHookDecision, + PipelineReportHook, + PolicyApplyResult, + PolicyUpdatePlan, + Rollout, + RolloutAnalysis, + Rubric, + RubricCriterion, + RubricEvaluation, + Trajectory, +) +from openviking.session.train.components.reporter import ConsolePipelineReporter +from openviking.storage.transaction import init_lock_manager, reset_lock_manager + + +@pytest.fixture(autouse=True) +def _train_lock_manager(): + reset_lock_manager() + init_lock_manager(InMemoryAGFS(), redo_recovery_enabled=False) + yield + reset_lock_manager() + + +def _case() -> Case: + return Case( + name="duplicate_booking", + task_signature="booking_duplicate", + input={"user_request": "cancel the duplicate booking"}, + rubric=Rubric( + name="booking_duplicate_rubric", + description="Cancel only the verified duplicate booking.", + criteria=[ + RubricCriterion( + name="verify_duplicate", + description="Verify duplicate status before cancellation.", + required=True, + weight=1.0, + ) + ], + ), + ) + + +class DummyVikingFS: + def __init__(self): + self.reloads = 0 + self.version = 1 + + async def ls(self, uri: str, output: str = "original", ctx=None): + del output, ctx + return [ + { + "name": "booking_duplicate_handling.md", + "uri": f"{uri.rstrip('/')}/booking_duplicate_handling.md", + "isDir": False, + } + ] + + async def read_file(self, uri: str, ctx=None): + del uri, ctx + self.reloads += 1 + return ( + "## Situation\n- Duplicate booking handling\n\n" + "" + ) + + def _uri_to_path(self, uri: str, ctx=None) -> str: + account_id = getattr(ctx, "account_id", "default") + return f"/local/{account_id}/{uri.removeprefix('viking://').strip('/')}" + + +def _policy_set(*, version: int = 1, viking_fs: DummyVikingFS | None = None) -> ExperienceSet: + return ExperienceSet( + root_uri="viking://user/u/memories/experiences", + policies=[ + Experience( + name="booking_duplicate_handling", + uri="viking://user/u/memories/experiences/booking_duplicate_handling.md", + version=version, + status="production", + content="## Situation\n- Duplicate booking handling", + ) + ], + viking_fs=viking_fs or DummyVikingFS(), + request_context=fake_request_context(), + ) + + +@dataclass +class DummyGradient: + target_name: str + target_uri: str | None + base_version: int | None + rationale: str + links: list[StoredLink] + confidence: float + metadata: dict[str, Any] = field(default_factory=dict) + + +class DummySnapshotter: + def __init__(self): + self.calls = 0 + + async def snapshot(self, policy_set: ExperienceSet, context: Any) -> str: + assert policy_set.root_uri + self.calls += 1 + return f"snapshot-{self.calls}" + + +class DummyExecutor: + def __init__(self): + self.calls = 0 + + async def execute(self, cases: list[Case], policy_set: ExperienceSet, context) -> list[Rollout]: + self.calls += 1 + assert context.policy_snapshot_id.startswith("snapshot-") + epoch = int(context.policy_snapshot_id.removeprefix("snapshot-")) - 1 + return [ + Rollout( + case=case, + messages=[ + Message( + id=f"msg-{case.name}", + role="user", + parts=[TextPart(text=str(case.input["user_request"]))], + ) + ], + policy_snapshot_id=context.policy_snapshot_id, + evaluation=RubricEvaluation( + passed=True, + score=float(epoch), + criterion_results=[], + feedback=[], + ), + ) + for case in cases + ] + + +class DummyAnalyzer: + def __init__(self): + self.calls = [] + + async def analyze(self, rollout: Rollout, context: Any) -> RolloutAnalysis: + self.calls.append(rollout) + epoch = int(rollout.policy_snapshot_id.removeprefix("snapshot-")) - 1 + return RolloutAnalysis( + evaluation=RubricEvaluation( + passed=True, + score=float(epoch), + criterion_results=[], + feedback=[], + ), + trajectories=[ + Trajectory( + name=rollout.case.task_signature, + uri=f"viking://user/u/memories/trajectories/{rollout.case.name}.md", + content="trajectory content", + outcome="success", + retrieval_anchor="Stage: final; Capability: duplicate booking handling", + ) + ], + ) + + +class DummyEstimator: + async def estimate( + self, + analysis: RolloutAnalysis, + experience_set: ExperienceSet, + context: Any, + ) -> list[DummyGradient]: + traj = analysis.trajectories[0] + return [ + DummyGradient( + target_name="booking_duplicate_handling", + target_uri=experience_set.policies[0].uri, + base_version=experience_set.policies[0].version, + rationale="trajectory succeeded", + links=[ + StoredLink( + from_uri=experience_set.policies[0].uri, + to_uri=traj.uri, + link_type="derived_from", + weight=1.0, + ) + ], + confidence=0.9, + ) + ] + + +class DummyOptimizer: + async def plan( + self, + gradients: list[DummyGradient], + policy_set: ExperienceSet, + context: Any, + ) -> PolicyUpdatePlan: + return PolicyUpdatePlan(metadata={"gradient_count": len(gradients)}) + + +class DummyUpdater: + last_instance = None + + def __init__(self): + self.transaction_handles = [] + DummyUpdater.last_instance = self + + async def apply( + self, + plan: PolicyUpdatePlan, + policy_set: ExperienceSet, + context: Any, + *, + transaction_handle: Any = None, + ) -> PolicyApplyResult: + self.transaction_handles.append(transaction_handle) + updated = ExperienceSet( + root_uri=policy_set.root_uri, + policies=[ + Experience( + name=p.name, + uri=p.uri, + version=p.version + 1, + status=p.status, + content=p.content, + metadata=dict(p.metadata), + ) + for p in policy_set.policies + ], + metadata=dict(policy_set.metadata), + viking_fs=policy_set.viking_fs, + request_context=policy_set.request_context, + ) + if hasattr(policy_set.viking_fs, "version") and updated.policies: + policy_set.viking_fs.version = updated.policies[0].version + return PolicyApplyResult( + updated_policy_set=updated, + written_uris=[p.uri for p in updated.policies], + ) + + +class RecordingLifecycleHook(NoopPipelineLifecycleHook): + def __init__(self): + self.events: list[tuple[str, Any]] = [] + + def on_epoch_start(self, *, epoch: int, context: Any) -> None: + del context + self.events.append(("epoch_start", epoch)) + + async def on_epoch_end( + self, + *, + epoch_result: Any, + policy_set: Any, + context: Any, + ) -> PipelineHookDecision | None: + del policy_set, context + self.events.append(("epoch_end", epoch_result.epoch)) + return None + + def on_train_rollout_report( + self, + *, + report: dict[str, Any], + context: Any, + ) -> None: + del context + self.events.append(("train_rollout_report", report["epoch"])) + + def on_train_report( + self, + *, + report: dict[str, Any], + context: Any, + ) -> None: + del context + self.events.append(("train_report", report["epoch"])) + + def on_eval_report( + self, + *, + label: str, + report: dict[str, Any], + context: Any, + ) -> None: + del context + self.events.append(("eval_report", label, report["epoch"])) + + +@pytest.mark.asyncio +async def test_default_policy_optimization_pipeline_runs_one_batch(): + snapshotter = DummySnapshotter() + pipeline = OfflinePolicyOptimizationPipeline( + snapshotter=snapshotter, + rollout_executor=DummyExecutor(), + rollout_analyzer=DummyAnalyzer(), + gradient_estimator=DummyEstimator(), + policy_optimizer=DummyOptimizer(), + policy_updater=DummyUpdater(), + ) + + initial_policy_set = _policy_set() + result = await pipeline.train( + case_loader=ListCaseLoader([_case()]), + policy_set=initial_policy_set, + context=PipelineContext(), + ) + + assert len(result.analyses) == 1 + assert len(result.gradients) == 1 + assert result.plan.metadata == {"gradient_count": 1} + assert result.apply_result.updated_policy_set.policies[0].version == 2 + assert result.apply_result.written_uris == [ + "viking://user/u/memories/experiences/booking_duplicate_handling.md" + ] + assert DummyUpdater.last_instance is not None + assert len(DummyUpdater.last_instance.transaction_handles) == 1 + assert DummyUpdater.last_instance.transaction_handles[0] is not None + assert initial_policy_set.viking_fs.reloads == 1 + assert len(result.epochs) == 1 + assert result.epochs[0].epoch == 0 + assert result.epochs[0].policy_snapshot_ids == ["snapshot-1"] + + +@pytest.mark.asyncio +async def test_offline_policy_optimization_pipeline_supports_train_and_eval(): + snapshotter = DummySnapshotter() + pipeline = OfflinePolicyOptimizationPipeline( + snapshotter=snapshotter, + rollout_executor=DummyExecutor(), + rollout_analyzer=DummyAnalyzer(), + gradient_estimator=DummyEstimator(), + policy_optimizer=DummyOptimizer(), + policy_updater=DummyUpdater(), + ) + policy_set = _policy_set() + + before_eval = await pipeline.eval( + case_loader=ListCaseLoader([_case()]), + policy_set=policy_set, + context=PipelineContext(execution_metadata={"epoch": -1}), + ) + result = await pipeline.train( + case_loader=ListCaseLoader([_case()]), + policy_set=policy_set, + context=PipelineContext(max_epochs=2), + ) + after_eval = await pipeline.eval( + case_loader=ListCaseLoader([_case()]), + policy_set=result.apply_result.updated_policy_set, + context=PipelineContext(execution_metadata={"epoch": 2}), + ) + + assert before_eval.epoch == -1 + assert before_eval.metadata["score"] == 0.0 + assert [item.epoch for item in result.epochs] == [0, 1] + assert result.evaluation_passes == [] + assert result.epochs[0].metadata["score"] == 1.0 + assert result.epochs[1].metadata["score"] == 2.0 + assert after_eval.epoch == 2 + assert after_eval.metadata["score"] == 3.0 + assert result.metadata["first_score"] == 1.0 + assert result.metadata["final_score"] == 2.0 + assert result.metadata["score_delta"] == 1.0 + assert result.apply_result.updated_policy_set.policies[0].version == 3 + + +@pytest.mark.asyncio +async def test_train_trials_expands_training_cases_per_epoch(): + class RecordingExecutor(DummyExecutor): + def __init__(self): + super().__init__() + self.train_trials = [] + + async def execute(self, cases, policy_set, context): + self.train_trials.extend(case.input.get("train_trial") for case in cases) + return await super().execute(cases, policy_set, context) + + executor = RecordingExecutor() + pipeline = OfflinePolicyOptimizationPipeline( + snapshotter=DummySnapshotter(), + rollout_executor=executor, + rollout_analyzer=DummyAnalyzer(), + gradient_estimator=DummyEstimator(), + policy_optimizer=DummyOptimizer(), + policy_updater=DummyUpdater(), + ) + + result = await pipeline.train( + case_loader=ListCaseLoader([_case()]), + policy_set=_policy_set(), + context=PipelineContext(train_trials=3), + ) + + assert len(result.analyses) == 3 + assert executor.train_trials == [0, 1, 2] + + +@pytest.mark.asyncio +async def test_training_updates_execution_metadata_epoch_each_epoch(): + pipeline = OfflinePolicyOptimizationPipeline( + snapshotter=DummySnapshotter(), + rollout_executor=DummyExecutor(), + rollout_analyzer=DummyAnalyzer(), + gradient_estimator=DummyEstimator(), + policy_optimizer=DummyOptimizer(), + policy_updater=DummyUpdater(), + ) + context = PipelineContext(max_epochs=2, execution_metadata={"rollout_stage": "eval_train_rollout"}) + + result = await pipeline.train( + case_loader=ListCaseLoader([_case()]), + policy_set=_policy_set(), + context=context, + ) + + assert [item.epoch for item in result.epochs] == [0, 1] + assert context.execution_metadata["epoch"] == 1 + + +@pytest.mark.asyncio +async def test_train_runs_test_eval_after_each_epoch_when_configured(): + hook = RecordingLifecycleHook() + pipeline = OfflinePolicyOptimizationPipeline( + snapshotter=DummySnapshotter(), + rollout_executor=DummyExecutor(), + rollout_analyzer=DummyAnalyzer(), + gradient_estimator=DummyEstimator(), + policy_optimizer=DummyOptimizer(), + policy_updater=DummyUpdater(), + ) + + result = await pipeline.train( + case_loader=ListCaseLoader([_case()]), + policy_set=_policy_set(), + context=PipelineContext( + max_epochs=2, + eval_each_epoch_case_loader=ListCaseLoader([_case()]), + lifecycle_hooks=[PipelineReportHook(), hook], + ), + ) + + assert [item.epoch for item in result.evaluation_passes] == [0, 1] + assert result.metadata["evaluation_pass_count"] == 2 + assert [ + item.metadata.get("rollout_stage") for item in result.evaluation_passes + ] == ["test_rollout", "test_rollout"] + assert [item.metadata.get("eval_split") for item in result.evaluation_passes] == [ + "test", + "test", + ] + assert ("eval_report", "test_rollout", 0) in hook.events + assert ("eval_report", "test_rollout", 1) in hook.events + + +@pytest.mark.asyncio +async def test_train_epoch_eval_uses_configured_split_metadata(): + hook = RecordingLifecycleHook() + pipeline = OfflinePolicyOptimizationPipeline( + snapshotter=DummySnapshotter(), + rollout_executor=DummyExecutor(), + rollout_analyzer=DummyAnalyzer(), + gradient_estimator=DummyEstimator(), + policy_optimizer=DummyOptimizer(), + policy_updater=DummyUpdater(), + ) + + result = await pipeline.train( + case_loader=ListCaseLoader([_case()]), + policy_set=_policy_set(), + context=PipelineContext( + max_epochs=1, + eval_each_epoch_case_loader=ListCaseLoader([_case()]), + execution_metadata={ + "rollout_stage": "eval_train_rollout", + "eval_split": "train", + }, + lifecycle_hooks=[PipelineReportHook(), hook], + ), + ) + + assert len(result.evaluation_passes) == 1 + assert result.evaluation_passes[0].metadata.get("rollout_stage") == "eval_train_rollout" + assert result.evaluation_passes[0].metadata.get("eval_split") == "train" + assert ("eval_report", "eval_train_rollout", 0) in hook.events + + +@pytest.mark.asyncio +async def test_offline_policy_optimization_pipeline_epoch_hook_can_stop_training(): + pipeline = OfflinePolicyOptimizationPipeline( + snapshotter=DummySnapshotter(), + rollout_executor=DummyExecutor(), + rollout_analyzer=DummyAnalyzer(), + gradient_estimator=DummyEstimator(), + policy_optimizer=DummyOptimizer(), + policy_updater=DummyUpdater(), + ) + class StopTrainingHook(NoopPipelineLifecycleHook): + def __init__(self): + self.epochs: list[int] = [] + + async def on_epoch_end( + self, + *, + epoch_result: Any, + policy_set: Any, + context: Any, + ) -> PipelineHookDecision: + del policy_set, context + self.epochs.append(epoch_result.epoch) + return PipelineHookDecision( + stop_training=True, + reason="unit test stop", + metadata={"epoch": epoch_result.epoch}, + ) + + hook = StopTrainingHook() + + result = await pipeline.train( + case_loader=ListCaseLoader([_case()]), + policy_set=_policy_set(), + context=PipelineContext(max_epochs=3, lifecycle_hooks=[hook]), + ) + + assert hook.epochs == [0] + assert [item.epoch for item in result.epochs] == [0] + assert result.metadata["completed_epochs"] == 1 + assert result.metadata["max_epochs"] == 3 + assert result.metadata["stopped_early"] is True + assert result.metadata["stop_reason"] == "unit test stop" + assert result.metadata["stop_metadata"] == {"epoch": 0} + + +@pytest.mark.asyncio +async def test_pipeline_lifecycle_hooks_receive_report_events(): + hook = RecordingLifecycleHook() + pipeline = OfflinePolicyOptimizationPipeline( + snapshotter=DummySnapshotter(), + rollout_executor=DummyExecutor(), + rollout_analyzer=DummyAnalyzer(), + gradient_estimator=DummyEstimator(), + policy_optimizer=DummyOptimizer(), + policy_updater=DummyUpdater(), + ) + + policy_set = _policy_set() + await pipeline.train( + case_loader=ListCaseLoader([_case()]), + policy_set=policy_set, + context=PipelineContext( + lifecycle_hooks=[PipelineReportHook(), hook], + ), + ) + await pipeline.eval( + case_loader=ListCaseLoader([_case()]), + policy_set=policy_set, + context=PipelineContext( + lifecycle_hooks=[PipelineReportHook(), hook], + execution_metadata={"epoch": 1, "rollout_stage": "final_test_rollout"}, + ), + ) + + assert hook.events == [ + ("epoch_start", 0), + ("train_rollout_report", 0), + ("epoch_end", 0), + ("train_report", 0), + ("eval_report", "final_test_rollout", 1), + ] + + +@pytest.mark.asyncio +async def test_policy_optimization_pipeline_trains_from_external_rollouts_without_executor(): + snapshotter = DummySnapshotter() + executor = DummyExecutor() + analyzer = DummyAnalyzer() + pipeline = OfflinePolicyOptimizationPipeline( + snapshotter=snapshotter, + rollout_executor=executor, + rollout_analyzer=analyzer, + gradient_estimator=DummyEstimator(), + policy_optimizer=DummyOptimizer(), + policy_updater=DummyUpdater(), + ) + rollout = Rollout( + case=_case(), + messages=[ + Message( + id="external-rollout-user", + role="user", + parts=[TextPart(text="cancel duplicate booking")], + ) + ], + policy_snapshot_id="snapshot-1", + ) + + result = await pipeline.train_from_rollouts( + rollouts=[rollout], + policy_set=_policy_set(), + context=PipelineContext(), + ) + + assert executor.calls == 0 + assert snapshotter.calls == 0 + assert analyzer.calls == [rollout] + assert len(result.analyses) == 1 + assert len(result.gradients) == 1 + assert result.plan.metadata == {"gradient_count": 1} + assert result.apply_result.updated_policy_set.policies[0].version == 2 + assert result.apply_result.written_uris == [ + "viking://user/u/memories/experiences/booking_duplicate_handling.md" + ] + assert result.metadata["source"] == "external_rollouts" + assert result.metadata["rollout_count"] == 1 + + +@pytest.mark.asyncio +async def test_policy_optimization_pipeline_realtime_rollouts_require_case(): + pipeline = OfflinePolicyOptimizationPipeline( + snapshotter=DummySnapshotter(), + rollout_executor=DummyExecutor(), + rollout_analyzer=DummyAnalyzer(), + gradient_estimator=DummyEstimator(), + policy_optimizer=DummyOptimizer(), + policy_updater=DummyUpdater(), + ) + rollout = Rollout( + case=None, + messages=[ + Message( + id="external-rollout-user", + role="user", + parts=[TextPart(text="cancel duplicate booking")], + ) + ], + policy_snapshot_id="snapshot-1", + ) + + with pytest.raises(ValueError, match="requires Rollout.case"): + await pipeline.train_from_rollouts( + rollouts=[rollout], + policy_set=_policy_set(), + context=PipelineContext(), + ) + + +@pytest.mark.asyncio +async def test_list_case_loader_yields_copy(): + cases = [_case()] + loader = ListCaseLoader(cases) + batches = [batch async for batch in loader.batches(None)] + + assert batches == [cases] + assert batches[0] is not cases + + +@pytest.mark.asyncio +async def test_batch_policy_trainer_trains_from_rollout_batch(): + from openviking.session.train import BatchPolicyTrainer + + trainer = BatchPolicyTrainer( + rollout_analyzer=DummyAnalyzer(), + gradient_estimator=DummyEstimator(), + policy_optimizer=DummyOptimizer(), + policy_updater=DummyUpdater(), + ) + rollouts = [ + Rollout( + case=_case(), + messages=[ + Message( + id="batch-rollout-user", + role="user", + parts=[TextPart(text="cancel duplicate booking")], + ) + ], + policy_snapshot_id="snapshot-1", + ) + ] + + result = await trainer.train_rollouts( + rollouts=rollouts, + policy_set=_policy_set(), + context=PipelineContext(), + ) + + assert len(result.analyses) == 1 + assert len(result.gradients) == 1 + assert result.plan.metadata == {"gradient_count": 1} + assert result.apply_result.updated_policy_set.policies[0].version == 2 + assert result.metadata["source"] == "batch_rollouts" + assert result.metadata["rollout_count"] == 1 + + +@pytest.mark.asyncio +async def test_streaming_policy_trainer_flushes_on_gradient_count(): + from openviking.session.train import ( + PolicyPlanItem, + StreamingPolicyTrainer, + StreamingPolicyTrainerConfig, + ) + + class LinkingOptimizer(DummyOptimizer): + async def plan(self, gradients, policy_set, context): + del policy_set, context + return PolicyUpdatePlan( + items=[ + PolicyPlanItem( + kind="upsert", + memory_type="experiences", + target_name=f"exp_{idx}", + target_uri=f"viking://user/u/memories/experiences/exp_{idx}.md", + before_content=None, + after_content=f"content {idx}", + links=list(gradient.links), + ) + for idx, gradient in enumerate(gradients) + ], + metadata={"gradient_count": len(gradients)}, + ) + + trainer = StreamingPolicyTrainer( + policy_set=_policy_set(), + rollout_analyzer=DummyAnalyzer(), + gradient_estimator=DummyEstimator(), + policy_optimizer=LinkingOptimizer(), + policy_updater=DummyUpdater(), + context=PipelineContext(), + config=StreamingPolicyTrainerConfig( + max_gradients_per_update=2, + max_wait_seconds=60.0, + timer_check_interval_seconds=60.0, + ), + ) + rollout1 = Rollout( + case=_case(), + messages=[Message(id="s1", role="user", parts=[TextPart(text="first")])], + policy_snapshot_id="snapshot-1", + ) + rollout2 = Rollout( + case=_case(), + messages=[Message(id="s2", role="user", parts=[TextPart(text="second")])], + policy_snapshot_id="snapshot-1", + ) + + first, second = await asyncio.gather( + trainer.submit_rollout(rollout1), + trainer.submit_rollout(rollout2), + ) + assert first is not second + assert first.batch_result is second.batch_result + assert len(first.analyses) == 1 + assert len(second.analyses) == 1 + assert [analysis.trajectories[0].uri for analysis in first.analyses] == [ + "viking://user/u/memories/trajectories/duplicate_booking.md" + ] + assert len(first.plan.items) == 2 + assert len(second.plan.items) == 2 + assert second.metadata["flush_reason"] == "count" + assert second.metadata["gradient_count"] == 1 + assert second.metadata["batch_gradient_count"] == 2 + assert second.apply_result.updated_policy_set.policies[0].version == 2 + assert await trainer.get_buffered_gradient_count() == 0 + assert trainer.last_apply_result is second.batch_result.apply_result + + assert await trainer.close() is None + assert trainer.closed is True + + +@pytest.mark.asyncio +async def test_streaming_policy_trainer_scopes_concurrent_submit_results_by_source_trajectory(): + from openviking.session.train import ( + PolicyPlanItem, + StreamingPolicyTrainer, + StreamingPolicyTrainerConfig, + ) + + class CaseAwareAnalyzer: + async def analyze(self, rollout, context): + del context + return RolloutAnalysis( + evaluation=RubricEvaluation( + passed=True, + score=1.0, + criterion_results=[], + feedback=[], + ), + trajectories=[ + Trajectory( + name=rollout.case.name, + uri=f"viking://user/u/memories/trajectories/{rollout.case.name}.md", + content=f"trajectory {rollout.case.name}", + outcome="success", + retrieval_anchor="", + ) + ], + ) + + class NewExpEstimator: + async def estimate(self, analysis, experience_set, context): + del context + traj = analysis.trajectories[0] + target_uri = f"{experience_set.root_uri}/{traj.name}.md" + return [ + DummyGradient( + target_name=traj.name, + target_uri=target_uri, + base_version=None, + rationale="new scoped experience", + links=[ + StoredLink( + from_uri=target_uri, + to_uri=traj.uri, + link_type="derived_from", + weight=1.0, + ) + ], + confidence=0.9, + ) + ] + + class LinkingOptimizer: + async def plan(self, gradients, policy_set, context): + del policy_set, context + return PolicyUpdatePlan( + items=[ + PolicyPlanItem( + kind="upsert", + memory_type="experiences", + target_name=gradient.target_name, + target_uri=gradient.target_uri, + before_content=None, + after_content=f"content {gradient.target_name}", + links=list(gradient.links), + ) + for gradient in gradients + ], + metadata={"gradient_count": len(gradients)}, + ) + + class PassthroughUpdater: + async def apply(self, plan, policy_set, context, *, transaction_handle=None): + del context, transaction_handle + return PolicyApplyResult( + updated_policy_set=policy_set, + written_uris=[item.target_uri for item in plan.items if item.target_uri], + errors=[], + ) + + def make_case(name: str) -> Case: + case = _case() + return Case( + name=name, + task_signature=case.task_signature, + input=case.input, + rubric=case.rubric, + ) + + trainer = StreamingPolicyTrainer( + policy_set=_policy_set(), + rollout_analyzer=CaseAwareAnalyzer(), + gradient_estimator=NewExpEstimator(), + policy_optimizer=LinkingOptimizer(), + policy_updater=PassthroughUpdater(), + context=PipelineContext(), + config=StreamingPolicyTrainerConfig( + max_gradients_per_update=2, + max_wait_seconds=60.0, + timer_check_interval_seconds=60.0, + ), + ) + rollout_a = Rollout( + case=make_case("case_a"), + messages=[Message(id="a", role="user", parts=[TextPart(text="a")])], + policy_snapshot_id="snapshot-1", + ) + rollout_b = Rollout( + case=make_case("case_b"), + messages=[Message(id="b", role="user", parts=[TextPart(text="b")])], + policy_snapshot_id="snapshot-1", + ) + + first, second = await asyncio.gather( + trainer.submit_rollout(rollout_a), + trainer.submit_rollout(rollout_b), + ) + + assert first.batch_result is second.batch_result + assert {item.target_name for item in first.batch_result.plan.items} == {"case_a", "case_b"} + assert [item.target_name for item in first.plan.items] == ["case_a"] + assert [item.target_name for item in second.plan.items] == ["case_b"] + assert first.apply_result.written_uris == [ + "viking://user/u/memories/experiences/case_a.md" + ] + assert second.apply_result.written_uris == [ + "viking://user/u/memories/experiences/case_b.md" + ] + + assert await trainer.close() is None + + +@pytest.mark.asyncio +async def test_streaming_policy_trainer_splits_flush_by_gradient_count(): + from openviking.session.train import StreamingPolicyTrainer, StreamingPolicyTrainerConfig + + class MultiGradientEstimator: + async def estimate(self, analysis, experience_set, context): + del context + traj = analysis.trajectories[0] + return [ + DummyGradient( + target_name="booking_duplicate_handling", + target_uri=experience_set.policies[0].uri, + base_version=experience_set.policies[0].version, + rationale=f"gradient {idx}", + links=[ + StoredLink( + from_uri=experience_set.policies[0].uri, + to_uri=traj.uri, + link_type="derived_from", + weight=1.0, + ) + ], + confidence=0.9, + ) + for idx in range(5) + ] + + class RecordingOptimizer: + def __init__(self): + self.gradient_counts = [] + + async def plan(self, gradients, policy_set, context): + del policy_set, context + self.gradient_counts.append(len(gradients)) + return PolicyUpdatePlan(metadata={"gradient_count": len(gradients)}) + + optimizer = RecordingOptimizer() + trainer = StreamingPolicyTrainer( + policy_set=_policy_set(), + rollout_analyzer=DummyAnalyzer(), + gradient_estimator=MultiGradientEstimator(), + policy_optimizer=optimizer, + policy_updater=DummyUpdater(), + context=PipelineContext(), + config=StreamingPolicyTrainerConfig( + max_gradients_per_update=3, + max_wait_seconds=0.01, + timer_check_interval_seconds=0.01, + ), + ) + rollout = Rollout( + case=_case(), + messages=[Message(id="split", role="user", parts=[TextPart(text="split")])], + policy_snapshot_id="snapshot-1", + ) + + result = await trainer.submit_rollout(rollout) + + assert optimizer.gradient_counts == [3, 2] + assert result.metadata["gradient_count"] == 5 + assert result.metadata["chunk_count"] == 2 + assert result.metadata["chunk_gradient_counts"] == [3, 2] + assert result.plan.metadata["chunk_item_counts"] == [0, 0] + assert result.apply_result.metadata["chunk_count"] == 2 + assert result.apply_result.updated_policy_set.policies[0].version == 3 + + assert await trainer.close() is None + + +@pytest.mark.asyncio +async def test_streaming_policy_trainer_chunks_multiple_target_gradients_by_count(): + """Gradients from different targets share the same chunk pool — split only by count.""" + from openviking.session.train import StreamingPolicyTrainer, StreamingPolicyTrainerConfig + + class MultiTargetEstimator: + async def estimate(self, analysis, experience_set, context): + del context + traj = analysis.trajectories[0] + return [ + DummyGradient( + target_name=f"target_{idx}", + target_uri=f"{experience_set.root_uri}/target_{idx}.md", + base_version=None, + rationale=f"gradient {idx}", + links=[ + StoredLink( + from_uri=f"{experience_set.root_uri}/target_{idx}.md", + to_uri=traj.uri, + link_type="derived_from", + weight=1.0, + ) + ], + confidence=0.9, + ) + for idx in range(5) + ] + + class RecordingOptimizer: + def __init__(self): + self.gradient_counts = [] + + async def plan(self, gradients, policy_set, context): + del policy_set, context + self.gradient_counts.append(len(gradients)) + return PolicyUpdatePlan(metadata={"gradient_count": len(gradients)}) + + optimizer = RecordingOptimizer() + trainer = StreamingPolicyTrainer( + policy_set=_policy_set(), + rollout_analyzer=DummyAnalyzer(), + gradient_estimator=MultiTargetEstimator(), + policy_optimizer=optimizer, + policy_updater=DummyUpdater(), + context=PipelineContext(), + config=StreamingPolicyTrainerConfig( + max_gradients_per_update=3, + max_wait_seconds=0.01, + timer_check_interval_seconds=0.01, + ), + ) + rollout = Rollout( + case=_case(), + messages=[Message(id="split-targets", role="user", parts=[TextPart(text="split")])], + policy_snapshot_id="snapshot-1", + ) + + result = await trainer.submit_rollout(rollout) + + # gradients are chunked purely by count, target boundaries don't affect chunking + assert optimizer.gradient_counts == [3, 2] + assert result.metadata["chunk_gradient_counts"] == [3, 2] + assert "chunk_target_counts" not in result.metadata + assert await trainer.close() is None + + +@pytest.mark.asyncio +async def test_streaming_policy_trainer_mixes_categories_in_chunks(): + """All gradients share the same chunk pool regardless of training_category.""" + from openviking.session.train import StreamingPolicyTrainer, StreamingPolicyTrainerConfig + + class CategorizedEstimator: + async def estimate(self, analysis, experience_set, context): + del analysis, context + return [ + DummyGradient( + target_name=f"target_{idx}", + target_uri=f"{experience_set.root_uri}/target_{idx}.md", + base_version=None, + rationale=f"gradient {idx}", + links=[], + confidence=0.9, + metadata={"training_category": "category_a" if idx < 2 else "category_b"}, + ) + for idx in range(4) + ] + + class RecordingOptimizer: + def __init__(self): + self.gradient_counts = [] + self.categories = [] + + async def plan(self, gradients, policy_set, context): + del policy_set, context + self.gradient_counts.append(len(gradients)) + self.categories.append( + [gradient.metadata["training_category"] for gradient in gradients] + ) + return PolicyUpdatePlan(metadata={"gradient_count": len(gradients)}) + + optimizer = RecordingOptimizer() + trainer = StreamingPolicyTrainer( + policy_set=_policy_set(), + rollout_analyzer=DummyAnalyzer(), + gradient_estimator=CategorizedEstimator(), + policy_optimizer=optimizer, + policy_updater=DummyUpdater(), + context=PipelineContext(), + config=StreamingPolicyTrainerConfig( + max_gradients_per_update=3, + max_wait_seconds=0.01, + timer_check_interval_seconds=0.01, + ), + ) + rollout = Rollout( + case=_case(), + messages=[Message(id="split-categories", role="user", parts=[TextPart(text="split")])], + policy_snapshot_id="snapshot-1", + ) + + result = await trainer.submit_rollout(rollout) + + # categories are no longer kept separate — all gradients chunked purely by count + assert optimizer.gradient_counts == [3, 1] + assert optimizer.categories == [ + ["category_a", "category_a", "category_b"], + ["category_b"], + ] + assert result.metadata["chunk_gradient_counts"] == [3, 1] + assert "chunk_categories" not in result.metadata + assert await trainer.close() is None + + +@pytest.mark.asyncio +async def test_streaming_policy_trainer_flushes_on_timer(): + from openviking.session.train import StreamingPolicyTrainer, StreamingPolicyTrainerConfig + + trainer = StreamingPolicyTrainer( + policy_set=_policy_set(), + rollout_analyzer=DummyAnalyzer(), + gradient_estimator=DummyEstimator(), + policy_optimizer=DummyOptimizer(), + policy_updater=DummyUpdater(), + context=PipelineContext(), + config=StreamingPolicyTrainerConfig( + max_gradients_per_update=10, + max_wait_seconds=0.01, + timer_check_interval_seconds=0.01, + ), + ) + rollout = Rollout( + case=_case(), + messages=[Message(id="timer", role="user", parts=[TextPart(text="timer")])], + policy_snapshot_id="snapshot-1", + ) + + result = await trainer.submit_rollout(rollout) + assert result is not None + assert result.metadata["flush_reason"] == "time" + assert result.metadata["gradient_count"] == 1 + assert trainer.last_apply_result is not None + assert trainer.last_apply_result.updated_policy_set.policies[0].version == 2 + assert await trainer.get_buffered_gradient_count() == 0 + + assert await trainer.close() is None + assert trainer.closed is True + + +@pytest.mark.asyncio +async def test_streaming_policy_trainer_close_flushes_buffer_and_rejects_submit(): + from openviking.session.train import StreamingPolicyTrainer, StreamingPolicyTrainerConfig + + trainer = StreamingPolicyTrainer( + policy_set=_policy_set(), + rollout_analyzer=DummyAnalyzer(), + gradient_estimator=DummyEstimator(), + policy_optimizer=DummyOptimizer(), + policy_updater=DummyUpdater(), + context=PipelineContext(), + config=StreamingPolicyTrainerConfig( + max_gradients_per_update=10, + max_wait_seconds=0.01, + timer_check_interval_seconds=0.01, + ), + ) + rollout = Rollout( + case=_case(), + messages=[Message(id="close", role="user", parts=[TextPart(text="close")])], + policy_snapshot_id="snapshot-1", + ) + + submit_task = asyncio.create_task(trainer.submit_rollout(rollout)) + await asyncio.sleep(0) + assert await trainer.get_buffered_gradient_count() == 1 + + result = await trainer.close() + + assert result is not None + assert result.metadata["flush_reason"] == "close" + assert result.metadata["gradient_count"] == 1 + submit_result = await submit_task + assert submit_result.batch_result is result + assert trainer.closed is True + assert await trainer.get_buffered_gradient_count() == 0 + assert trainer.last_apply_result is result.apply_result + assert await trainer.close() is None + + with pytest.raises(RuntimeError, match="closed"): + await trainer.submit_rollout(rollout) + + +@pytest.mark.asyncio +async def test_get_streaming_policy_trainer_returns_process_global_instance(): + from openviking.session.train import ( + StreamingPolicyTrainerConfig, + get_streaming_policy_trainer, + make_streaming_policy_trainer_key, + ) + + policy_set = _policy_set() + key = make_streaming_policy_trainer_key( + policy_root_uri=policy_set.root_uri, + request_context=policy_set.request_context, + ) + + first = await get_streaming_policy_trainer( + key=key, + policy_set=policy_set, + rollout_analyzer=DummyAnalyzer(), + gradient_estimator=DummyEstimator(), + policy_optimizer=DummyOptimizer(), + policy_updater=DummyUpdater(), + context=PipelineContext(), + config=StreamingPolicyTrainerConfig(max_gradients_per_update=3), + ) + second = await get_streaming_policy_trainer( + key=key, + policy_set=policy_set, + rollout_analyzer=DummyAnalyzer(), + gradient_estimator=DummyEstimator(), + policy_optimizer=DummyOptimizer(), + policy_updater=DummyUpdater(), + context=PipelineContext(), + config=StreamingPolicyTrainerConfig(max_gradients_per_update=9), + ) + + assert second is first + assert first.config.max_gradients_per_update == 3 + +class FakeSessionCommitClient: + def __init__(self): + self.created_sessions = [] + self.messages = {} + self.committed_sessions = [] + self.task_poll_counts = {} + + async def create_session(self, *, session_id, memory_policy=None): + self.created_sessions.append((session_id, memory_policy)) + + async def batch_add_messages(self, session_id, messages): + self.messages.setdefault(session_id, []).extend(messages) + + async def commit_session(self, session_id, telemetry=False, *, keep_recent_count=0): + self.committed_sessions.append((session_id, keep_recent_count, telemetry)) + return { + "task_id": f"task-{session_id}", + "archive_uri": f"viking://user/default/sessions/{session_id}/history/archive_001", + "trace_id": f"trace-{session_id}", + } + + async def get_task(self, task_id): + self.task_poll_counts[task_id] = self.task_poll_counts.get(task_id, 0) + 1 + return {"task_id": task_id, "status": "completed", "result": {}} + + +@pytest.mark.asyncio +async def test_session_commit_policy_trainer_records_commit_trace_id(): + from openviking.session.train import SessionCommitPolicyTrainer + + client = FakeSessionCommitClient() + trainer = SessionCommitPolicyTrainer( + client=client, + run_id="run1", + keep_recent_count=2, + poll_interval_seconds=0.01, + ) + case = _case() + rollout = Rollout( + case=case, + messages=[Message(id="m1", role="user", parts=[TextPart(text="hello")])], + policy_snapshot_id="snapshot-1", + evaluation=RubricEvaluation(passed=True, score=1.0, criterion_results=[], feedback=[]), + metadata={"data_split": "unit", "task_no": 7, "execution_metadata": {"epoch": 3}}, + ) + + result = await trainer.train_rollouts([rollout], _policy_set()) + + commit_result = result.apply_result.metadata["commit_results"][0] + assert commit_result["task_id"] == f"task-{commit_result['session_id']}" + assert commit_result["archive_uri"].endswith("/history/archive_001") + assert commit_result["trace_id"] == f"trace-{commit_result['session_id']}" + assert commit_result["task_status"] == "completed" + assert client.committed_sessions == [(commit_result["session_id"], 2, True)] + assert client.created_sessions == [ + ( + commit_result["session_id"], + { + "memory_types": ["cases", "trajectories", "experiences"], + "working_memory": {"enabled": False}, + }, + ) + ] + + +@pytest.mark.asyncio +async def test_session_commit_policy_trainer_splits_large_message_batches(): + from openviking.session.train import SessionCommitPolicyTrainer + + client = FakeSessionCommitClient() + batch_sizes = [] + + async def batch_add_messages(session_id, messages): + batch_sizes.append(len(messages)) + client.messages.setdefault(session_id, []).extend(messages) + + client.batch_add_messages = batch_add_messages + trainer = SessionCommitPolicyTrainer( + client=client, + run_id="run1", + poll_interval_seconds=0.01, + ) + rollout = Rollout( + case=_case(), + messages=[ + Message(id=f"m{i}", role="user", parts=[TextPart(text=f"hello {i}")]) + for i in range(250) + ], + policy_snapshot_id="snapshot-1", + evaluation=RubricEvaluation(passed=True, score=1.0, criterion_results=[], feedback=[]), + metadata={"data_split": "unit", "task_no": 7, "execution_metadata": {"epoch": 3}}, + ) + + result = await trainer.train_rollouts([rollout], _policy_set()) + + commit_result = result.apply_result.metadata["commit_results"][0] + assert commit_result["error"] is None + assert batch_sizes == [100, 100, 52] + assert len(client.messages[commit_result["session_id"]]) == 252 + + +@pytest.mark.asyncio +async def test_session_commit_policy_trainer_streams_commit_trace_events(tmp_path): + from openviking.session.train import JsonlEventRecorder, SessionCommitPolicyTrainer + + client = FakeSessionCommitClient() + events_path = tmp_path / "events.jsonl" + trainer = SessionCommitPolicyTrainer( + client=client, + run_id="run1", + poll_interval_seconds=0.01, + event_recorder=JsonlEventRecorder( + events_path, + default_fields={"dataset": "unit", "domain": "booking", "run_id": "run1"}, + ), + ) + rollout = Rollout( + case=_case(), + messages=[Message(id="m1", role="user", parts=[TextPart(text="hello")])], + policy_snapshot_id="snapshot-1", + evaluation=RubricEvaluation(passed=True, score=1.0, criterion_results=[], feedback=[]), + metadata={"data_split": "unit", "task_no": 7, "execution_metadata": {"epoch": 3}}, + ) + + result = await trainer.train_rollouts([rollout], _policy_set()) + + lines = [json.loads(line) for line in events_path.read_text().splitlines()] + assert [line["event"] for line in lines] == ["train_commit_submitted", "train_commit_done"] + commit_result = result.apply_result.metadata["commit_results"][0] + assert lines[0]["trace_id"] == commit_result["trace_id"] + assert lines[1]["trace_id"] == commit_result["trace_id"] + assert lines[1]["task_status"] == "completed" + assert lines[1]["session_id"] == commit_result["session_id"] + assert lines[1]["task_no"] == 7 + assert lines[1]["dataset"] == "unit" + + +@pytest.mark.asyncio +async def test_jsonl_pipeline_event_hook_omits_full_commit_results(tmp_path): + from openviking.session.train import JsonlEventRecorder, JsonlPipelineEventHook + + events_path = tmp_path / "pipeline.jsonl" + hook = JsonlPipelineEventHook(JsonlEventRecorder(events_path)) + await hook.on_train_report( + report={ + "epoch": 0, + "committed_rollout_count": 1, + "errors": [], + "commit_results": [ + {"trace_id": "trace-1", "task_id": "task-1", "telemetry_id": "telemetry-1"} + ], + }, + context=PipelineContext(execution_metadata={"epoch": 0, "training": True}), + ) + + data = json.loads(events_path.read_text()) + assert data["event"] == "train_result" + assert data["commit_trace_ids"] == ["trace-1"] + assert data["commit_task_ids"] == ["task-1"] + assert "commit_results" not in data + + +@pytest.mark.asyncio +async def test_rollout_artifact_recorder_writes_train_rollouts_before_commit(tmp_path): + from openviking.session.train import RolloutArtifactRecorder + + recorder = RolloutArtifactRecorder(run_dir=tmp_path) + case = Case( + name="tau2_airline_train_7", + task_signature="tau2:airline:train:7", + input={ + "data_split": "airline_train", + "task_no": 7, + "task_id": "task-7", + "user_request": "change my seat", + }, + rubric=Rubric(name="reward", description="reward", criteria=[]), + ) + rollout = Rollout( + case=case, + messages=[Message(id="m1", role="user", parts=[TextPart(text="hello")])], + policy_snapshot_id="snapshot-1", + evaluation=RubricEvaluation(passed=False, score=0.0, criterion_results=[], feedback=[]), + metadata={ + "memory": "remember airline seat-change rules", + "task_case_experience_skill": "# task_case_experience\nlinked exp content", + }, + ) + + recorder.on_train_rollout_end( + epoch=0, + rollouts=[rollout], + snapshot_id="snapshot-1", + policy_set=None, + context=None, + ) + + rollout_dir = ( + tmp_path + / "rollouts" + / "airline_train_task_7_task-7" + / "epoch_0" + / "1.train_rollout" + / "trial_0" + ) + assert (rollout_dir / "status.json").exists() + assert (rollout_dir / "rollout.json").exists() + assert (rollout_dir / "evaluation.json").exists() + assert (rollout_dir / "prompt_for_llm.md").exists() + assert (rollout_dir / "memory_context.md").read_text() == "remember airline seat-change rules" + skill_path = rollout_dir / "task_case_experience_skill.md" + assert skill_path.read_text() == "# task_case_experience\nlinked exp content" + assert (rollout_dir / "commit_messages.json").exists() + assert not (rollout_dir / "commit_result.json").exists() + status = json.loads((rollout_dir / "status.json").read_text()) + assert status["artifact_state"] == "rollout_done" + assert status["has_task_case_experience_skill"] is True + assert status["task_case_experience_skill_path"] == "task_case_experience_skill.md" + index = json.loads((tmp_path / "rollouts_index.json").read_text()) + assert index["case_groups"][0]["rollouts"][0]["artifact_state"] == "rollout_done" + + +def test_rollout_artifact_recorder_separates_epoch_eval_dirs(tmp_path): + from openviking.session.train import RolloutArtifactRecorder + + recorder = RolloutArtifactRecorder(run_dir=tmp_path) + case = Case( + name="tau2_airline_test_0_t0", + task_signature="tau2:airline:test:2:trial:0", + input={ + "data_split": "airline_test", + "task_no": 0, + "task_id": "2", + "eval_trial": 0, + "original_case_name": "tau2_airline_test_0", + }, + rubric=Rubric(name="reward", description="reward", criteria=[]), + ) + + for epoch in (0, 1): + rollout = Rollout( + case=case, + messages=[Message(id=f"m{epoch}", role="user", parts=[TextPart(text="hello")])], + policy_snapshot_id=f"snapshot-{epoch}", + evaluation=RubricEvaluation( + passed=epoch == 1, + score=float(epoch == 1), + criterion_results=[], + feedback=[], + ), + ) + recorder.record_eval( + label="test_rollout", + epoch=epoch, + analyses=[ + RolloutAnalysis( + evaluation=rollout.evaluation, + trajectories=[], + metadata={"rollout": rollout}, + ) + ], + ) + + group_dir = tmp_path / "rollouts" / "airline_test_task_0_2" + assert (group_dir / "epoch_0" / "4.test_rollout" / "trial_0" / "status.json").exists() + assert (group_dir / "epoch_1" / "4.test_rollout" / "trial_0" / "status.json").exists() + assert not (group_dir / "4.test_rollout" / "trial_0").exists() + + index = recorder.finalize().to_dict() + rollout_stages = [item["stage"] for item in index["case_groups"][0]["rollouts"]] + assert rollout_stages == ["epoch_0/4.test_rollout", "epoch_1/4.test_rollout"] + + +def test_rollout_artifact_recorder_uses_stage_name_dirs(tmp_path): + from openviking.session.train import RolloutArtifactRecorder + from openviking.session.train.context import ExecutionContext + + recorder = RolloutArtifactRecorder(run_dir=tmp_path) + case = Case( + name="tau2_airline_train_7", + task_signature="tau2:airline:train:7", + input={ + "data_split": "airline_train", + "task_no": 7, + "task_id": "task-7", + }, + rubric=Rubric(name="reward", description="reward", criteria=[]), + ) + rollout = Rollout( + case=case, + messages=[Message(id="m1", role="user", parts=[TextPart(text="hello")])], + policy_snapshot_id="snapshot-1", + evaluation=RubricEvaluation(passed=True, score=1.0, criterion_results=[], feedback=[]), + ) + + recorder.record_rollout_completion( + rollout=rollout, + index=0, + context=ExecutionContext( + policy_snapshot_id="snapshot-1", + metadata={"epoch": 0, "training": True, "rollout_stage": "eval_train_rollout"}, + ), + ) + + group_dir = tmp_path / "rollouts" / "airline_train_task_7_task-7" + assert (group_dir / "epoch_0" / "3.eval_train_rollout" / "trial_0" / "status.json").exists() + assert not (group_dir / "epoch_0" / "2.train" / "trial_0").exists() + + index = recorder.finalize().to_dict() + rollout_index = index["case_groups"][0]["rollouts"][0] + assert rollout_index["stage"] == "epoch_0/3.eval_train_rollout" + assert rollout_index["path"].endswith("epoch_0/3.eval_train_rollout/trial_0") + + +def test_rollout_artifact_recorder_keeps_baseline_and_final_eval_dirs(tmp_path): + from openviking.session.train import RolloutArtifactRecorder + + recorder = RolloutArtifactRecorder(run_dir=tmp_path) + case = Case( + name="tau2_airline_test_0_t0", + task_signature="tau2:airline:test:2:trial:0", + input={ + "data_split": "airline_test", + "task_no": 0, + "task_id": "2", + "eval_trial": 0, + }, + rubric=Rubric(name="reward", description="reward", criteria=[]), + ) + + rollout = Rollout( + case=case, + messages=[Message(id="m", role="user", parts=[TextPart(text="hello")])], + policy_snapshot_id="snapshot", + evaluation=RubricEvaluation(passed=True, score=1.0, criterion_results=[], feedback=[]), + ) + analysis = RolloutAnalysis( + evaluation=rollout.evaluation, + trajectories=[], + metadata={"rollout": rollout}, + ) + + recorder.record_eval(label="baseline_test_rollout", epoch=-1, analyses=[analysis]) + recorder.record_eval(label="final_test_rollout", epoch=2, analyses=[analysis]) + + group_dir = tmp_path / "rollouts" / "airline_test_task_0_2" + assert (group_dir / "epoch_-1" / "0.baseline_test_rollout" / "trial_0" / "status.json").exists() + assert (group_dir / "epoch_2" / "5.final_test_rollout" / "trial_0" / "status.json").exists() + + +def test_console_reporter_highlights_accuracy_and_prints_epoch_summary(capsys): + reporter = ConsolePipelineReporter(use_rich=False) + context = PipelineContext(eval_each_epoch_case_loader=object()) + + reporter.on_train_rollout_report( + report={ + "epoch": 1, + "case_count": 30, + "accuracy": 0.6, + "passed_count": 18, + "average_reward": 0.6, + }, + context=context, + ) + reporter.on_train_report( + report={ + "epoch": 1, + "committed_rollout_count": 30, + "errors": [], + "train_rollout": { + "epoch": 1, + "case_count": 30, + "accuracy": 0.6, + "passed_count": 18, + "average_reward": 0.6, + }, + }, + context=context, + ) + reporter.on_eval_report( + label="test_rollout", + report={ + "epoch": 1, + "rollout_stage": "test_rollout", + "split": "test", + "trial_count": 8, + "case_count": 160, + "total_rollout_count": 160, + "case_count_per_trial": 20, + "accuracy_mean": 0.58125, + "accuracy_std": 0.055551, + "average_reward_mean": 0.58125, + "average_reward_std": 0.055551, + }, + context=context, + ) + + output = capsys.readouterr().out + + assert "accuracy=\x1b[1;33m60.00%\x1b[0m" in output + assert "epoch 1 summary" in output + assert "TRAIN accuracy: \x1b[0m\x1b[1;33m60.00%\x1b[0m" in output + assert "EVAL accuracy: \x1b[0m\x1b[1;33m58.13%\x1b[0m" in output + assert output.count("------------------------------------------------------------") >= 3 + + +def test_rollout_artifact_event_recorder_enriches_commit_result(tmp_path): + from openviking.session.train import RolloutArtifactEventRecorder, RolloutArtifactRecorder + + recorder = RolloutArtifactRecorder(run_dir=tmp_path) + case = Case( + name="tau2_airline_train_7", + task_signature="tau2:airline:train:7", + input={ + "data_split": "airline_train", + "task_no": 7, + "task_id": "task-7", + "user_request": "change my seat", + }, + rubric=Rubric(name="reward", description="reward", criteria=[]), + ) + rollout = Rollout( + case=case, + messages=[Message(id="m1", role="user", parts=[TextPart(text="hello")])], + policy_snapshot_id="snapshot-1", + evaluation=RubricEvaluation(passed=True, score=1.0, criterion_results=[], feedback=[]), + ) + recorder.record_train_rollouts(epoch=0, rollouts=[rollout]) + event_recorder = RolloutArtifactEventRecorder(recorder) + + event_recorder.record( + "train_commit_submitted", + index=0, + epoch=0, + split="airline_train", + task_no=7, + case_task_id="task-7", + case_name="tau2_airline_train_7", + session_id="session-1", + stage="commit_session", + task_id="commit-task-1", + archive_uri="viking://archive", + trace_id="trace-1", + telemetry_id="telemetry-1", + task_status=None, + score=1.0, + error=None, + ) + + rollout_dir = ( + tmp_path + / "rollouts" + / "airline_train_task_7_task-7" + / "epoch_0" + / "1.train_rollout" + / "trial_0" + ) + commit_dir = ( + tmp_path + / "rollouts" + / "airline_train_task_7_task-7" + / "epoch_0" + / "2.train" + / "trial_0" + ) + assert not (rollout_dir / "commit_result.json").exists() + commit_result = json.loads((commit_dir / "commit_result.json").read_text()) + status = json.loads((rollout_dir / "status.json").read_text()) + index = json.loads((tmp_path / "rollouts_index.json").read_text()) + assert commit_result["artifact_state"] == "commit_submitted" + assert commit_result["session_id"] == "session-1" + assert status["artifact_state"] == "commit_submitted" + assert status["archive_uri"] == "viking://archive" + assert status["commit_path"] == str(commit_dir) + rollout_index = index["case_groups"][0]["rollouts"][0] + assert rollout_index["artifact_state"] == "commit_submitted" + assert rollout_index["path"] == str(rollout_dir) + assert rollout_index["commit_path"] == str(commit_dir) + + +@pytest.mark.asyncio +async def test_rollout_artifact_recorder_writes_epoch_commit_artifacts_under_commit_dir(tmp_path): + from openviking.session.train import RolloutArtifactRecorder + + class CommitArtifactClient: + async def read(self, uri): + assert uri == "viking://archive/memory_diff.json" + return json.dumps({ + "operations": { + "adds": [ + {"uri": "viking://memory/new.md", "after": "# New\nbody"} + ], + "updates": [ + { + "uri": "viking://memory/old.md", + "before": "# Old\nbody", + "after": "# Old\nnew body", + } + ], + "deletes": [], + }, + "summary": {"total_adds": 1, "total_updates": 1, "total_deletes": 0}, + }) + + recorder = RolloutArtifactRecorder(run_dir=tmp_path, client=CommitArtifactClient()) + case = Case( + name="tau2_airline_train_7", + task_signature="tau2:airline:train:7", + input={ + "data_split": "airline_train", + "task_no": 7, + "task_id": "task-7", + "user_request": "change my seat", + }, + rubric=Rubric(name="reward", description="reward", criteria=[]), + ) + rollout = Rollout( + case=case, + messages=[Message(id="m1", role="user", parts=[TextPart(text="hello")])], + policy_snapshot_id="snapshot-1", + evaluation=RubricEvaluation(passed=True, score=1.0, criterion_results=[], feedback=[]), + ) + analysis = RolloutAnalysis( + evaluation=rollout.evaluation, + trajectories=[], + metadata={"rollout": rollout}, + ) + + await recorder.record_train_epoch( + epoch=0, + analyses=[analysis], + commit_results=[ + { + "index": 0, + "archive_uri": "viking://archive", + "task_status": "completed", + "error": None, + } + ], + ) + + train_dir = ( + tmp_path + / "rollouts" + / "airline_train_task_7_task-7" + / "epoch_0" + / "1.train_rollout" + / "trial_0" + ) + commit_dir = ( + tmp_path + / "rollouts" + / "airline_train_task_7_task-7" + / "epoch_0" + / "2.train" + / "trial_0" + ) + assert (train_dir / "status.json").exists() + assert not (train_dir / "commit_result.json").exists() + assert not (train_dir / "memory_diff.json").exists() + assert (commit_dir / "commit_result.json").exists() + memory_diff_json = json.loads((commit_dir / "memory_diff.json").read_text()) + assert memory_diff_json["summary"]["total_adds"] == 1 + memory_diff_md = (commit_dir / "memory_diff.md").read_text() + assert "--- /dev/null" in memory_diff_md + assert "+++ viking://memory/new.md" in memory_diff_md + assert "--- viking://memory/old.md" in memory_diff_md + assert "+new body" in memory_diff_md + + status = json.loads((train_dir / "status.json").read_text()) + assert status["artifact_state"] == "memory_diff_done" + assert status["commit_path"] == str(commit_dir) + assert status["memory_diff_path"] == str(commit_dir / "memory_diff.json") + assert status["memory_diff_markdown_path"] == str(commit_dir / "memory_diff.md") + + +class DelayedSessionCommitClient(FakeSessionCommitClient): + def __init__(self, *, pending_polls: int): + super().__init__() + self.pending_polls = pending_polls + + async def get_task(self, task_id): + self.task_poll_counts[task_id] = self.task_poll_counts.get(task_id, 0) + 1 + if self.task_poll_counts[task_id] <= self.pending_polls: + return {"task_id": task_id, "status": "running", "result": {}} + return {"task_id": task_id, "status": "completed", "result": {}} + + +@pytest.mark.asyncio +async def test_session_commit_policy_trainer_waits_without_default_timeout(): + from openviking.session.train import SessionCommitPolicyTrainer + + client = DelayedSessionCommitClient(pending_polls=3) + trainer = SessionCommitPolicyTrainer( + client=client, + run_id="run1", + poll_interval_seconds=0.01, + ) + rollout = Rollout( + case=_case(), + messages=[Message(id="m1", role="user", parts=[TextPart(text="hello")])], + policy_snapshot_id="snapshot-1", + evaluation=RubricEvaluation(passed=True, score=1.0, criterion_results=[], feedback=[]), + metadata={"data_split": "unit", "task_no": 7, "execution_metadata": {"epoch": 3}}, + ) + + result = await asyncio.wait_for(trainer.train_rollouts([rollout], _policy_set()), timeout=1.0) + + commit_result = result.apply_result.metadata["commit_results"][0] + assert commit_result["task_status"] == "completed" + assert commit_result["error"] is None + assert client.task_poll_counts[commit_result["task_id"]] == 4 + + +@pytest.mark.asyncio +async def test_session_commit_policy_trainer_can_still_use_explicit_timeout(): + from openviking.session.train import SessionCommitPolicyTrainer + + client = DelayedSessionCommitClient(pending_polls=100) + trainer = SessionCommitPolicyTrainer( + client=client, + run_id="run1", + poll_interval_seconds=0.01, + timeout_seconds=0.02, + ) + rollout = Rollout( + case=_case(), + messages=[Message(id="m1", role="user", parts=[TextPart(text="hello")])], + policy_snapshot_id="snapshot-1", + evaluation=RubricEvaluation(passed=True, score=1.0, criterion_results=[], feedback=[]), + metadata={"data_split": "unit", "task_no": 7, "execution_metadata": {"epoch": 3}}, + ) + + result = await trainer.train_rollouts([rollout], _policy_set()) + + commit_result = result.apply_result.metadata["commit_results"][0] + assert commit_result["task_status"] == "timeout" + assert commit_result["error"] == "commit task timeout" + + +def test_tau2_case_loader_selects_exact_task_indices(tmp_path: Path, monkeypatch): + from benchmark.tau2.train import case_loader as tau2_case_loader + + domain_dir = tmp_path / "domains" / "airline" + domain_dir.mkdir(parents=True) + (domain_dir / "split_tasks.json").write_text( + json.dumps({"train": ["a", "b", "c"], "test": ["x", "y"]}), + encoding="utf-8", + ) + + class Task: + def __init__(self, task_id: str) -> None: + self.id = task_id + self.evaluation_criteria = f"criteria-{task_id}" + self.user_scenario = f"scenario-{task_id}" + + monkeypatch.setattr(tau2_case_loader, "_load_tau2_task", lambda _domain, task_id: Task(task_id)) + + loader = tau2_case_loader.Tau2CaseLoader( + domain="airline", + split="train", + data_root=str(tmp_path), + task_indices=[1], + ) + + cases = loader.load_cases() + + assert [case.input["task_id"] for case in cases] == ["b"] + assert [case.input["task_no"] for case in cases] == [1] + assert cases[0].name == "tau2_airline_train_1" + + +def test_tau2_service_filter_parses_task_indices(): + from benchmark.tau2.train.service_app import _task_indices_from_filters + + assert _task_indices_from_filters({}) is None + assert _task_indices_from_filters({"task_indices": [0, "3"]}) == [0, 3] + with pytest.raises(ValueError, match="task_indices filter must be a list"): + _task_indices_from_filters({"task_indices": 2}) + with pytest.raises(ValueError, match="task index must be >= 0"): + _task_indices_from_filters({"task_indices": [-1]}) diff --git a/tests/session/train/test_trajectory_analyzer_component.py b/tests/session/train/test_trajectory_analyzer_component.py new file mode 100644 index 0000000000..8d4d50c911 --- /dev/null +++ b/tests/session/train/test_trajectory_analyzer_component.py @@ -0,0 +1,196 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from openviking.message import Message, TextPart +from openviking.session.memory.dataclass import ResolvedOperation, ResolvedOperations +from openviking.session.train import ( + Case, + CriterionResult, + Rollout, + Rubric, + RubricEvaluation, +) +from openviking.session.train.components.trajectory_analyzer import ( + TrajectoryAnalyzerContext, + TrajectoryRolloutAnalyzer, +) + + +class FakeExtractLoop: + created = [] + + def __init__(self, **kwargs): + self.kwargs = kwargs + self._transaction_handle = None + FakeExtractLoop.created.append(self) + + async def run(self): + return ( + ResolvedOperations( + upsert_operations=[ + ResolvedOperation( + old_memory_file_content=None, + memory_fields={ + "trajectory_name": "task", + "outcome": "success", + "retrieval_anchor": "Stage: final", + "content": "# task\nbody", + }, + memory_type="trajectories", + uris=["viking://user/u/memories/trajectories/task_20260607120000.md"], + page_id=100, + ) + ], + delete_file_contents=[], + errors=[], + resolved_links=[], + ), + [], + ) + + +class FakeVikingFS: + agfs = None + + def __init__(self): + self.files = {} + self.writes = [] + + async def read_file(self, uri, ctx=None): + return self.files[uri] + + async def write_file(self, uri, content, ctx=None): + self.files[uri] = content + self.writes.append((uri, content, ctx)) + + +class FakeRolloutEvaluator: + def __init__(self): + self.calls = [] + + async def evaluate(self, rollout, context): + self.calls.append((rollout, context)) + return RubricEvaluation( + passed=False, + score=0.25, + criterion_results=[ + CriterionResult( + criterion_name="tau2_reward", + passed=False, + score=0.0, + feedback=["reward was zero"], + evidence=["missing confirmation"], + ) + ], + feedback=["task failed"], + metadata={"source": "fake"}, + ) + + +def _rollout() -> Rollout: + return Rollout( + case=Case( + name="case", + task_signature="task", + input={}, + rubric=Rubric(name="r", description="d", criteria=[]), + ), + messages=[ + Message( + id="m", + role="user", + parts=[TextPart(text="hello")], + created_at="2026-06-07T12:00:00", + ) + ], + policy_snapshot_id="snapshot", + ) + + +@pytest.mark.asyncio +async def test_trajectory_rollout_analyzer_extracts_and_persists_trajectory(monkeypatch): + from openviking.session.train.components import trajectory_analyzer as module + + FakeExtractLoop.created.clear() + fs = FakeVikingFS() + monkeypatch.setattr(module, "ExtractLoop", FakeExtractLoop) + monkeypatch.setattr(module, "get_viking_fs", lambda: fs) + + analyzer = TrajectoryRolloutAnalyzer(viking_fs=fs, vlm=SimpleNamespace(model="fake")) + context = TrajectoryAnalyzerContext( + request_context=SimpleNamespace( + user=SimpleNamespace(account_id="default", user_id="u"), + account_id="default", + ) + ) + + analysis = await analyzer.analyze(_rollout(), context) + + assert FakeExtractLoop.created + created_loop = FakeExtractLoop.created[0] + assert created_loop._transaction_handle is None + provider = created_loop.kwargs["context_provider"] + assert provider._transaction_handle is None + assert [schema.memory_type for schema in provider.get_memory_schemas(context.request_context)] == [ + "trajectories" + ] + assert len(fs.writes) == 1 + assert fs.writes[0][0] == "viking://user/u/memories/trajectories/task_20260607120000.md" + assert '"case_name": "case"' in fs.writes[0][1] + assert len(analysis.trajectories) == 1 + traj = analysis.trajectories[0] + assert traj.name == "task" + assert traj.outcome == "success" + assert traj.retrieval_anchor == "Stage: final" + assert traj.metadata["case_name"] == "case" + assert analysis.evaluation.passed is True + assert analysis.metadata["policy_snapshot_id"] == "snapshot" + + +@pytest.mark.asyncio +async def test_trajectory_rollout_analyzer_evaluates_before_extracting_trajectory(monkeypatch): + from openviking.session.train.components import trajectory_analyzer as module + + FakeExtractLoop.created.clear() + fs = FakeVikingFS() + evaluator = FakeRolloutEvaluator() + evaluator_context = {"benchmark": "tau2"} + monkeypatch.setattr(module, "ExtractLoop", FakeExtractLoop) + monkeypatch.setattr(module, "get_viking_fs", lambda: fs) + + analyzer = TrajectoryRolloutAnalyzer( + viking_fs=fs, + vlm=SimpleNamespace(model="fake"), + evaluator=evaluator, + ) + context = TrajectoryAnalyzerContext( + request_context=SimpleNamespace( + user=SimpleNamespace(account_id="default", user_id="u"), + account_id="default", + ), + evaluator_context=evaluator_context, + ) + + rollout = _rollout() + analysis = await analyzer.analyze(rollout, context) + + assert evaluator.calls == [(rollout, evaluator_context)] + assert analysis.evaluation.score == 0.25 + assert analysis.evaluation.metadata == {"source": "fake"} + created_loop = FakeExtractLoop.created[0] + provider = created_loop.kwargs["context_provider"] + assert len(provider.messages) == 2 + assert provider.messages[0] is rollout.messages[0] + feedback_message = provider.messages[1] + assert feedback_message.role == "user" + assert "[Rollout Evaluation]" in feedback_message.content + assert "score: 0.25" in feedback_message.content + assert "task failed" in feedback_message.content + assert "missing confirmation" in feedback_message.content + assert analysis.metadata["extraction_message_count"] == 2 diff --git a/tests/test_config_loader.py b/tests/test_config_loader.py index 2ade6a1820..05ae0cc2f0 100644 --- a/tests/test_config_loader.py +++ b/tests/test_config_loader.py @@ -186,7 +186,7 @@ def test_openviking_config_rejects_unknown_memory_field(monkeypatch): OpenVikingConfigSingleton.reset_instance() -def test_openviking_config_rejects_memory_v1(monkeypatch): +def test_openviking_config_ignores_deprecated_memory_version(monkeypatch): monkeypatch.setenv(OPENVIKING_CONFIG_ENV, "/tmp/codex-no-config.json") from openviking_cli.utils.config.open_viking_config import ( @@ -194,8 +194,12 @@ def test_openviking_config_rejects_memory_v1(monkeypatch): OpenVikingConfigSingleton, ) - with pytest.raises(ValueError, match="legacy memory v1 has been removed"): - OpenVikingConfig.from_dict({"memory": {"version": "v1"}}) + config = OpenVikingConfig.from_dict({}) + assert config.memory.version == "v3" + + for configured_version in ("v1", "v2", "v3", "unsupported"): + config = OpenVikingConfig.from_dict({"memory": {"version": configured_version}}) + assert config.memory.version == "v3" OpenVikingConfigSingleton.reset_instance() @@ -209,11 +213,15 @@ def test_openviking_config_ignores_deprecated_agent_memory_enabled(monkeypatch): ) legacy_config = OpenVikingConfig.from_dict({"memory": {"agent_memory_enabled": False}}) + working_memory_config = OpenVikingConfig.from_dict( + {"memory": {"working_memory_enabled": False}} + ) experimental_config = OpenVikingConfig.from_dict( {"memory": {"experimental_memory_switch": True}} ) assert not hasattr(legacy_config.memory, "agent_memory_enabled") + assert not hasattr(working_memory_config.memory, "working_memory_enabled") assert experimental_config.memory.experimental_memory_switch is True OpenVikingConfigSingleton.reset_instance() diff --git a/tests/test_prompt_manager.py b/tests/test_prompt_manager.py index 33b169c2e3..9557537f61 100644 --- a/tests/test_prompt_manager.py +++ b/tests/test_prompt_manager.py @@ -69,7 +69,7 @@ def teardown_function() -> None: OpenVikingConfigSingleton.reset_instance() -def test_profile_memory_template_keeps_profile_minimal_and_migrates_preferences(): +def test_profile_memory_template_includes_stable_identity_work_style_and_preferences(): template_path = PromptManager._get_bundled_templates_dir() / "memory" / "profile.yaml" schema = yaml.safe_load(template_path.read_text(encoding="utf-8")) text = "\n".join( @@ -79,19 +79,16 @@ def test_profile_memory_template_keeps_profile_minimal_and_migrates_preferences( ] ) - assert "identity summary" in text - assert "5-8" in text - assert "Complete but minimal" in text - assert "Rewrite the full profile" in text - assert "Do not append" in text - assert "migrate" in text - assert "preferences" in text - assert "Do not keep concrete preference examples" in text - assert "patch" in text - assert "rewrite the whole profile" in text + assert '"who the user is"' in text + assert "identity, work style, and preferences" in text + assert "profession, experience level, technical background" in text + assert "communication style, work habits" in text + assert "Do NOT include transient conversation content" in text + assert "Each item: self-contained" in text + assert "Only record objective statuses" in text -def test_preferences_memory_template_limits_topics_and_splits_when_too_large(): +def test_preferences_memory_template_keeps_topic_specific_preferences(): template_path = PromptManager._get_bundled_templates_dir() / "memory" / "preferences.yaml" schema = yaml.safe_load(template_path.read_text(encoding="utf-8")) text = "\n".join( @@ -102,14 +99,12 @@ def test_preferences_memory_template_limits_topics_and_splits_when_too_large(): ] ) - assert "Complete but minimal" in text - assert "3-8" in text - assert "800" in text - assert "split" in text - assert "semantic subtopics" in text - assert "evidenced by" in text - assert "as of" in text - assert "not become a second profile" in text + assert "what the user likes/dislikes or is accustomed to" in text + assert "specific preferences" in text + assert "specific topic" in text + assert "code style, communication style, tools, workflow" in text + assert "Store different topics as separate memory files" in text + assert "do not mix unrelated preferences" in text def test_prompt_manager_prefers_env_templates_dir_over_config(tmp_path, monkeypatch): diff --git a/tests/unit/isolated/test_tau2_service_app_progress.py b/tests/unit/isolated/test_tau2_service_app_progress.py new file mode 100644 index 0000000000..7721393200 --- /dev/null +++ b/tests/unit/isolated/test_tau2_service_app_progress.py @@ -0,0 +1,77 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 + +from __future__ import annotations + + +def test_tau2_service_configures_tau2_loguru_at_warning(monkeypatch): + import benchmark.tau2.train.service_app as service_app + + calls = [] + + class FakeLoguruLogger: + def remove(self): + calls.append(("remove",)) + + def add(self, sink, *, level): + del sink + calls.append(("add", level)) + + monkeypatch.setitem( + __import__("sys").modules, + "loguru", + type("FakeLoguruModule", (), {"logger": FakeLoguruLogger()})(), + ) + + service_app.configure_tau2_service_logging() + + assert calls == [("remove",), ("add", "WARNING")] + + +def test_tau2_service_disables_rollout_progress_by_default(monkeypatch): + import benchmark.tau2.train.service_app as service_app + + calls = [] + + def fake_create_dataset_service_app(**kwargs): + calls.append(kwargs) + return kwargs + + class FakeExecutor: + pass + + def fake_make_tau2_rollout_executor(**kwargs): + calls.append({"factory": kwargs}) + return FakeExecutor() + + monkeypatch.setattr(service_app, "create_dataset_service_app", fake_create_dataset_service_app) + monkeypatch.setattr(service_app, "make_tau2_rollout_executor", fake_make_tau2_rollout_executor) + + app = service_app.create_app(rollout_backend="native") + executor = app["make_rollout_executor"]({"rollout_backend": "vikingbot", "max_iterations": 5}) + + assert isinstance(executor, FakeExecutor) + assert calls[-1]["factory"]["backend"] == "vikingbot" + assert calls[-1]["factory"]["options"]["max_iterations"] == 5 + assert calls[-1]["factory"]["options"]["show_progress"] is False + + +def test_tau2_service_keeps_explicit_rollout_progress_override(monkeypatch): + import benchmark.tau2.train.service_app as service_app + + calls = [] + + def fake_create_dataset_service_app(**kwargs): + return kwargs + + def fake_make_tau2_rollout_executor(**kwargs): + calls.append(kwargs) + return object() + + monkeypatch.setattr(service_app, "create_dataset_service_app", fake_create_dataset_service_app) + monkeypatch.setattr(service_app, "make_tau2_rollout_executor", fake_make_tau2_rollout_executor) + + app = service_app.create_app(rollout_backend="native") + app["make_rollout_executor"]({"rollout_backend": "native", "show_progress": True}) + + assert calls[-1]["options"]["show_progress"] is True diff --git a/tests/unit/session/memory/test_embedding_template.py b/tests/unit/session/memory/test_embedding_template.py index 1912ef80fc..b16c5f3114 100644 --- a/tests/unit/session/memory/test_embedding_template.py +++ b/tests/unit/session/memory/test_embedding_template.py @@ -67,6 +67,29 @@ def test_content_template_can_access_content_extra_fields_and_extract_context(se assert rendered.startswith("Road Trip | Body content | 2026") + def test_content_template_render_failure_is_logged_before_plain_content_fallback(self): + memory_file = MemoryFile( + content="Body content", + extra_fields={"title": "Road Trip", "ranges": "0-1"}, + ) + extract_context = SimpleNamespace( + get_year=Mock(side_effect=RuntimeError("template render failed")) + ) + + with patch( + "openviking.session.memory.utils.memory_file_utils.logger.exception" + ) as mock_logger_exception: + rendered = MemoryFileUtils.write( + memory_file, + content_template="{{ title }} | {{ content }} | {{ extract_context.get_year(ranges) }}", + extract_context=extract_context, + ) + + assert rendered.startswith("Body content") + mock_logger_exception.assert_called_once_with( + "Failed to render memory content template; using plain content fallback" + ) + class TestEmbeddingTextConstruction: @pytest.mark.asyncio diff --git a/tests/unit/session/memory/test_extract_loop_match_text.py b/tests/unit/session/memory/test_extract_loop_match_text.py index c88264cf88..03db4d216e 100644 --- a/tests/unit/session/memory/test_extract_loop_match_text.py +++ b/tests/unit/session/memory/test_extract_loop_match_text.py @@ -180,7 +180,6 @@ async def test_run_always_includes_page_id_rules_when_links_disabled(self): context_provider=context_provider, isolation_handler=isolation_handler, ) - loop._mark_cache_breakpoint = AsyncMock() loop._call_llm = AsyncMock( return_value=( [], @@ -259,7 +258,6 @@ async def test_run_includes_link_page_id_rule_when_links_enabled(self): context_provider=context_provider, isolation_handler=isolation_handler, ) - loop._mark_cache_breakpoint = AsyncMock() loop._call_llm = AsyncMock( return_value=( [], @@ -332,7 +330,6 @@ async def test_run_logs_final_operations_after_old_memory_file_is_hydrated(self) context_provider=context_provider, isolation_handler=isolation_handler, ) - loop._mark_cache_breakpoint = AsyncMock() loop._call_llm = AsyncMock( return_value=( [], diff --git a/tests/unit/session/test_memory_policy.py b/tests/unit/session/test_memory_policy.py index 2576d22540..dbf8c98396 100644 --- a/tests/unit/session/test_memory_policy.py +++ b/tests/unit/session/test_memory_policy.py @@ -16,6 +16,7 @@ def test_memory_policy_defaults_to_self_and_peer(): assert policy.self_enabled is True assert policy.peer_enabled is True assert policy.memory_types is None + assert policy.working_memory_enabled is True def test_memory_policy_can_disable_peer_memory(): @@ -44,6 +45,27 @@ def test_memory_policy_uses_top_level_memory_types(): } +def test_memory_policy_can_disable_working_memory(): + policy = MemoryPolicy.from_dict({"working_memory": {"enabled": False}}) + + assert policy.self_enabled is True + assert policy.peer_enabled is True + assert policy.working_memory_enabled is False + assert policy.to_dict() == { + "self": {"enabled": True}, + "peer": {"enabled": True}, + "working_memory": {"enabled": False}, + } + + +def test_memory_policy_rejects_invalid_working_memory_shape(): + with pytest.raises(InvalidArgumentError, match="working_memory must be an object"): + MemoryPolicy.from_dict({"working_memory": False}) + + with pytest.raises(InvalidArgumentError, match="working_memory supports only: enabled"): + MemoryPolicy.from_dict({"working_memory": {"enabled": True, "mode": "summary"}}) + + def test_memory_policy_rejects_invalid_memory_types(): policy = MemoryPolicy.from_dict({"memory_types": ["profile", "missing"]}) diff --git a/tests/unit/session_train/test_batch_runner_indices.py b/tests/unit/session_train/test_batch_runner_indices.py new file mode 100644 index 0000000000..ec71e96835 --- /dev/null +++ b/tests/unit/session_train/test_batch_runner_indices.py @@ -0,0 +1,96 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 + +from __future__ import annotations + +from openviking.session.train.batch_runner import ( + BatchTrainEvalConfig, + _baseline_cache_key, + _case_loader, +) + + +def test_case_loader_uses_multi_value_train_and_eval_index_filters(): + config = BatchTrainEvalConfig( + dataset="tau2", + domain="airline", + train_index=[1, 5, 5, 6], + eval_index="10,14,18", + benchmark_service_url="http://127.0.0.1:1944", + ) + + train_loader = _case_loader( + config, + split="train", + sample_index=config.train_index, + ) + eval_loader = _case_loader( + config, + split="train", + sample_index=config.eval_index, + ) + + assert config.train_index == [1, 5, 6] + assert config.eval_index == [10, 14, 18] + assert train_loader.filters == {"task_indices": [1, 5, 6]} + assert eval_loader.filters == {"task_indices": [10, 14, 18]} + + +def test_baseline_cache_key_depends_on_multi_eval_index(): + base = BatchTrainEvalConfig( + dataset="tau2", + domain="airline", + eval_index=[1, 5, 6], + trials=8, + benchmark_service_url="http://127.0.0.1:1944", + ) + + assert _baseline_cache_key(base) == _baseline_cache_key( + BatchTrainEvalConfig( + dataset="tau2", + domain="airline", + eval_index="1,5,6", + trials=8, + benchmark_service_url="http://127.0.0.1:1944", + ) + ) + assert _baseline_cache_key(base) != _baseline_cache_key( + BatchTrainEvalConfig( + dataset="tau2", + domain="airline", + eval_index=[1, 5], + trials=8, + benchmark_service_url="http://127.0.0.1:1944", + ) + ) + + +def test_empty_index_filter_is_invalid(): + import pytest + + with pytest.raises(ValueError, match="train_index must not be empty"): + BatchTrainEvalConfig( + dataset="tau2", + domain="airline", + train_index="", + benchmark_service_url="http://127.0.0.1:1944", + ) + + +def test_train_trials_defaults_to_one_and_validates_positive(): + import pytest + + config = BatchTrainEvalConfig( + dataset="tau2", + domain="airline", + benchmark_service_url="http://127.0.0.1:1944", + ) + + assert config.train_trials == 1 + with pytest.raises(ValueError, match="train_trials must be > 0"): + BatchTrainEvalConfig( + dataset="tau2", + domain="airline", + train_trials=0, + benchmark_service_url="http://127.0.0.1:1944", + ) diff --git a/tests/unit/test_locomo_peer_wiring.py b/tests/unit/test_locomo_peer_wiring.py index 3f1f4a22ed..bf85727c61 100644 --- a/tests/unit/test_locomo_peer_wiring.py +++ b/tests/unit/test_locomo_peer_wiring.py @@ -10,6 +10,9 @@ def _load_module(module_name: str, relative_path: str): repo_root = Path(__file__).resolve().parents[2] module_path = repo_root / relative_path + module_dir = str(module_path.parent) + if module_dir not in sys.path: + sys.path.insert(0, module_dir) spec = importlib.util.spec_from_file_location(module_name, module_path) module = importlib.util.module_from_spec(spec) assert spec is not None and spec.loader is not None @@ -46,15 +49,15 @@ def _sample_payload(): } -def test_build_memory_policy_writes_peer_only(): - assert IMPORT_TO_OV.build_memory_policy(False) == { - "self": {"enabled": False}, - "peer": {"enabled": True}, - } - assert IMPORT_TO_OV.build_memory_policy(True) == { +def test_build_memory_policy_writes_peer_only_user_memories(): + expected = { "self": {"enabled": False}, "peer": {"enabled": True}, + "working_memory": {"enabled": False}, + "memory_types": ["entities", "events", "preferences", "profile"], } + assert IMPORT_TO_OV.build_memory_policy(False) == expected + assert IMPORT_TO_OV.build_memory_policy(True) == expected def test_build_session_messages_non_group_uses_sample_peer_and_prefixes_speaker(): @@ -82,6 +85,10 @@ async def create_session(self, memory_policy=None): calls.append(("create_session", memory_policy)) return {"session_id": "sess-1"} + async def get_session(self, session_id): + calls.append(("get_session", session_id)) + return {"commit_count": 0} + async def add_message( self, session_id=None, role=None, parts=None, created_at=None, peer_id=None ): @@ -123,6 +130,70 @@ async def close(self): assert add_calls[0][1]["peer_id"] == "conv-26" +@pytest.mark.asyncio +async def test_process_single_session_marks_ingest_record_with_canonical_sample_id( + monkeypatch, tmp_path +): + ingest_record = {} + success_csv = tmp_path / "import_success.csv" + record_path = tmp_path / ".ingest_record.json" + + async def fake_viking_ingest(*_args, **_kwargs): + return { + "token_usage": { + "embedding": 1, + "llm_input": 2, + "cache": 0, + "reasoning": 0, + "llm_output": 3, + "total": 6, + }, + "task_id": "task-1", + "trace_id": "trace-1", + } + + monkeypatch.setattr(IMPORT_TO_OV, "viking_ingest", fake_viking_ingest) + original_save_ingest_record = IMPORT_TO_OV.save_ingest_record + monkeypatch.setattr( + IMPORT_TO_OV, + "save_ingest_record", + lambda record: original_save_ingest_record(record, str(record_path)), + ) + + result = await IMPORT_TO_OV.process_single_session( + messages=[{"role": "user", "text": "Alice: Hi Bob", "peer_id": "conv-26"}], + sample_id="conv-26", + display_id="sample_0", + session_key="session_1", + meta={"sample_id": "conv-26", "session_key": "session_1"}, + run_time="2026-06-30T00:00:00", + ingest_record=ingest_record, + args=SimpleNamespace( + api_key="user-key", + auth_mode="api_key", + user="", + account="", + openviking_url="http://localhost:1933", + group_chat=False, + success_csv=str(success_csv), + error_log=str(tmp_path / "import_errors.log"), + _show_progress=True, + ), + ) + + assert result["status"] == "success" + assert result["sample_id"] == "conv-26" + assert result["display_id"] == "sample_0" + assert "viking:conv-26:session_1" in ingest_record + assert "viking:sample_0:session_1" not in ingest_record + assert IMPORT_TO_OV.is_already_ingested("conv-26", "session_1", ingest_record) + assert not IMPORT_TO_OV.is_already_ingested("sample_0", "session_1", ingest_record) + + persisted = json.loads(record_path.read_text(encoding="utf-8")) + assert "viking:conv-26:session_1" in persisted + assert "viking:sample_0:session_1" not in persisted + + def test_load_locomo_qa_keeps_internal_and_original_sample_ids(tmp_path): input_path = tmp_path / "locomo.json" input_path.write_text(json.dumps([_sample_payload()]), encoding="utf-8") @@ -138,7 +209,7 @@ def test_load_locomo_qa_keeps_internal_and_original_sample_ids(tmp_path): def test_run_vikingbot_chat_non_group_builds_sender_without_memory_peers(monkeypatch): calls = [] - def fake_run(cmd, capture_output, text, timeout=None, check=False): + def fake_run(cmd, capture_output, text, timeout=None, check=False, env=None): calls.append(cmd) return SimpleNamespace( stdout=json.dumps( @@ -155,13 +226,15 @@ def fake_run(cmd, capture_output, text, timeout=None, check=False): monkeypatch.setattr(RUN_EVAL.subprocess, "run", fake_run) - response, token_usage, _time_cost, iteration, tools_used_names = RUN_EVAL.run_vikingbot_chat( - question="Who said hello?", - question_time="2023-05-08", - sender_peer_id="conv-26", - question_id="sample_0_qa0", - config="/tmp/ov.conf", - memory_peer_ids=None, + response, token_usage, _time_cost, iteration, tools_used_names, _log_file = ( + RUN_EVAL.run_vikingbot_chat( + question="Who said hello?", + question_time="2023-05-08", + sender_peer_id="conv-26", + question_id="sample_0_qa0", + config="/tmp/ov.conf", + memory_peer_ids=None, + ) ) assert response == "ok" diff --git a/tests/unit/test_locomo_progress_utils.py b/tests/unit/test_locomo_progress_utils.py new file mode 100644 index 0000000000..60628689dd --- /dev/null +++ b/tests/unit/test_locomo_progress_utils.py @@ -0,0 +1,82 @@ +"""Tests for LoCoMo benchmark progress utilities.""" + +from rich.progress import ProgressColumn + +from benchmark.locomo.vikingbot.progress_utils import ( + AsyncProgressTracker, + ProgressSummaryColumn, + ThreadSafeProgressTracker, + ThreeStateBarColumn, + make_three_state_progress, +) + + +def test_three_state_bar_column_initializes_progress_column_state(): + column = ThreeStateBarColumn() + + assert isinstance(column, ProgressColumn) + assert hasattr(column, "_renderable_cache") + assert hasattr(column, "_update_time") + + +def test_three_state_progress_renders_without_missing_cache_error(): + progress, task_id = make_three_state_progress(description="Test", transient=True) + progress.update(task_id, total=4, completed=2, running=1, succeeded=1, failed=1) + + renderables = list(progress.get_renderables()) + + assert renderables + + +def test_progress_summary_hides_ok_and_zero_failed_counts(): + progress, task_id = make_three_state_progress(description="Test", transient=True) + progress.update(task_id, total=25, completed=23, running=0, succeeded=23, failed=0) + task = progress.tasks[0] + + summary = ProgressSummaryColumn().render(task).plain + + assert summary == "(23/25)" + assert "ok" not in summary + assert "failed" not in summary + + +def test_progress_summary_shows_failed_only_when_non_zero(): + progress, task_id = make_three_state_progress(description="Test", transient=True) + progress.update(task_id, total=25, completed=24, running=0, succeeded=23, failed=1) + task = progress.tasks[0] + + summary = ProgressSummaryColumn().render(task).plain + + assert summary == "(24/25, 1 failed)" + + +def test_locomo_progress_tracker_records_failed_jobs(): + progress, task_id = make_three_state_progress(description="Test", transient=True) + tracker = ThreadSafeProgressTracker(progress, task_id, total=2) + + tracker.job_started() + tracker.job_finished() + tracker.job_started() + tracker.job_finished(failed=True) + + task = progress.tasks[0] + assert tracker.done == 2 + assert tracker.failed == 1 + assert task.completed == 2 + assert task.fields["succeeded"] == 1 + assert task.fields["failed"] == 1 + + +def test_async_progress_tracker_records_failed_jobs(): + progress, task_id = make_three_state_progress(description="Test", transient=True) + tracker = AsyncProgressTracker(progress, task_id, total=2) + + tracker.job_started() + tracker.job_finished(failed=True) + + task = progress.tasks[0] + assert tracker.done == 1 + assert tracker.failed == 1 + assert task.completed == 1 + assert task.fields["succeeded"] == 0 + assert task.fields["failed"] == 1 diff --git a/tests/unit/test_message.py b/tests/unit/test_message.py index 987bd4f2a6..e1727d9fb5 100644 --- a/tests/unit/test_message.py +++ b/tests/unit/test_message.py @@ -324,6 +324,20 @@ def test_unknown_type_defaults_to_text(self): # The entire dict is converted to string assert "unknown" in part.text + def test_unknown_type_with_text_preserves_text(self): + """Unknown part types with text degrade to TextPart text.""" + data = { + "type": "control", + "control_type": "batch_training_case_spec", + "payload": {"protocol": "v1"}, + "text": "control text", + } + + part = part_from_dict(data) + + assert isinstance(part, TextPart) + assert part.text == "control text" + def test_missing_type_defaults_to_text(self): """Test missing type defaults to TextPart.""" data = {"text": "Hello"} @@ -588,6 +602,43 @@ def test_from_dict_basic(self): assert isinstance(msg.parts[0], TextPart) assert msg.parts[0].text == "Hello" + def test_from_dict_unknown_type_with_text_preserves_text(self): + """Unknown serialized part types with text degrade to TextPart.""" + d = { + "id": "msg-control", + "role": "system", + "parts": [ + { + "type": "control", + "control_type": "batch_training_case_spec", + "payload": {"protocol": "v1"}, + "text": "# OpenViking Batch Training CaseSpec v1", + } + ], + "created_at": "2026-03-26T10:30:00Z", + } + + msg = Message.from_dict(d) + + assert len(msg.parts) == 1 + assert isinstance(msg.parts[0], TextPart) + assert msg.parts[0].text == "# OpenViking Batch Training CaseSpec v1" + + def test_from_dict_missing_part_type_defaults_to_text(self): + """Serialized dict parts without type should default to TextPart.""" + d = { + "id": "msg-missing-type", + "role": "user", + "parts": [{"text": "hello"}], + "created_at": "2026-03-26T10:30:00Z", + } + + msg = Message.from_dict(d) + + assert len(msg.parts) == 1 + assert isinstance(msg.parts[0], TextPart) + assert msg.parts[0].text == "hello" + def test_from_dict_with_context_part(self): """Test from_dict with ContextPart.""" d = { diff --git a/tests/unit/test_remote_rollout_executor.py b/tests/unit/test_remote_rollout_executor.py new file mode 100644 index 0000000000..e987e736d1 --- /dev/null +++ b/tests/unit/test_remote_rollout_executor.py @@ -0,0 +1,130 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 + +from __future__ import annotations + +import asyncio + +import httpx + +from openviking.session.train.components.remote import RemoteRolloutExecutor +from openviking.session.train.context import ExecutionContext +from openviking.session.train.domain import ( + Case, + Experience, + ExperienceSet, + Rubric, + RubricCriterion, +) + + +def _case() -> Case: + return Case( + name="case-1", + task_signature="booking_duplicate", + input={"user_request": "cancel duplicate booking"}, + rubric=Rubric( + name="booking_rubric", + description="Cancel only the verified duplicate booking.", + criteria=[ + RubricCriterion( + name="verify_duplicate", + description="Verify duplicate status first.", + required=True, + weight=1.0, + ) + ], + ), + ) + + +def _policy_set() -> ExperienceSet: + return ExperienceSet( + root_uri="viking://user/u/memories/experiences", + policies=[ + Experience( + name="booking_policy", + uri="viking://user/u/memories/experiences/booking_policy.md", + version=2, + status="production", + content="Always verify duplicates before cancellation.", + ) + ], + ) + + +def test_remote_rollout_executor_retries_transient_missing_execution(monkeypatch): + calls: list[str] = [] + execution_id = "rollout_exec_delayed" + + def handler(request: httpx.Request) -> httpx.Response: + calls.append(f"{request.method} {request.url.path}") + if request.method == "POST" and request.url.path == "/v1/rollouts/execute": + return httpx.Response(200, json={"execution_id": execution_id, "status": "running"}) + if request.method == "GET" and request.url.path.endswith(execution_id): + get_count = sum(1 for call in calls if call.startswith("GET ")) + if get_count == 1: + return httpx.Response(404, json={"detail": "not found yet"}) + return httpx.Response( + 200, + json={ + "execution_id": execution_id, + "status": "completed", + "rollout": { + "case": { + "name": "case-1", + "task_signature": "booking_duplicate", + "input": {"user_request": "cancel duplicate booking"}, + "rubric": { + "name": "booking_rubric", + "description": "Cancel only the verified duplicate booking.", + "criteria": [ + { + "name": "verify_duplicate", + "description": "Verify duplicate status first.", + "required": True, + "weight": 1.0, + "metadata": {}, + } + ], + "metadata": {}, + }, + "metadata": {}, + }, + "messages": [], + "policy_snapshot_id": "snapshot-1", + "evaluation": None, + "metadata": {}, + }, + }, + ) + return httpx.Response(500, json={"error": "unexpected request"}) + + original_async_client = httpx.AsyncClient + monkeypatch.setattr( + httpx, + "AsyncClient", + lambda *args, **kwargs: original_async_client( + transport=httpx.MockTransport(handler), + base_url=kwargs.get("base_url"), + timeout=kwargs.get("timeout"), + ), + ) + + executor = RemoteRolloutExecutor( + service_url="http://rollout-service", + poll_interval_seconds=0.01, + missing_execution_grace_seconds=1.0, + ) + + rollouts = asyncio.run( + executor.execute( + [_case()], + _policy_set(), + ExecutionContext(policy_snapshot_id="snapshot-1"), + ) + ) + + assert len(rollouts) == 1 + assert rollouts[0].case.name == "case-1" + assert calls.count(f"GET /v1/rollouts/executions/{execution_id}") == 2 diff --git a/tests/unit/test_vikingbot_vlm_adapter_retry.py b/tests/unit/test_vikingbot_vlm_adapter_retry.py new file mode 100644 index 0000000000..6350d84b2c --- /dev/null +++ b/tests/unit/test_vikingbot_vlm_adapter_retry.py @@ -0,0 +1,193 @@ +from types import SimpleNamespace + +import httpx +import pytest +import vikingbot.providers.vlm_adapter as vlm_adapter +from vikingbot.providers.vlm_adapter import VLMProviderAdapter +from volcenginesdkarkruntime._exceptions import ArkRateLimitError + +from openviking.utils.model_retry import ( + is_retryable_rate_limit_error, +) + + +class _DisabledLangfuse: + enabled = False + _client = None + + +class _FakeVLM: + def __init__(self, failures: list[Exception], result: str = "ok"): + self.failures = list(failures) + self.result = result + self.calls = 0 + + async def get_completion_async(self, **_kwargs): + self.calls += 1 + if self.failures: + raise self.failures.pop(0) + return self.result + + +class _AsyncChunks: + def __init__(self, chunks): + self._chunks = chunks + + def __aiter__(self): + return self._iter() + + async def _iter(self): + for chunk in self._chunks: + yield chunk + + +class _FakeStreamingCompletions: + def __init__(self, failures: list[Exception], chunks): + self.failures = list(failures) + self.chunks = chunks + self.calls = 0 + + async def create(self, **_kwargs): + self.calls += 1 + if self.failures: + raise self.failures.pop(0) + return _AsyncChunks(self.chunks) + + +class _FakeStreamingVLM: + provider = "volcengine" + model = "test-model" + thinking = False + + def __init__(self, completions: _FakeStreamingCompletions): + self._client = SimpleNamespace( + chat=SimpleNamespace(completions=completions), + ) + + def get_async_client(self): + return self._client + + +@pytest.mark.asyncio +async def test_chat_retries_rate_limit_until_success(monkeypatch): + sleep_delays: list[float] = [] + + async def _sleep(delay: float): + sleep_delays.append(delay) + + monkeypatch.setattr(vlm_adapter, "rate_limit_retry_delay", lambda attempt: attempt) + monkeypatch.setattr(vlm_adapter.asyncio, "sleep", _sleep) + + fake_vlm = _FakeVLM( + [ + RuntimeError("Error code: 429 - ModelAccountTpmRateLimitExceeded"), + RuntimeError("TooManyRequests: rate limit"), + ], + result="done", + ) + adapter = VLMProviderAdapter(fake_vlm, "test-model", langfuse_client=_DisabledLangfuse()) + + response = await adapter.chat(messages=[{"role": "user", "content": "hello"}]) + + assert response.content == "done" + assert response.finish_reason == "stop" + assert fake_vlm.calls == 3 + assert sleep_delays == [1, 2] + + +@pytest.mark.asyncio +async def test_chat_does_not_retry_errors_without_rate_limit_markers(monkeypatch): + async def _sleep(_delay: float): + raise AssertionError("non-retryable errors must not sleep/retry") + + monkeypatch.setattr(vlm_adapter.asyncio, "sleep", _sleep) + + fake_vlm = _FakeVLM([RuntimeError("AuthenticationError Unauthorized")]) + adapter = VLMProviderAdapter(fake_vlm, "test-model", langfuse_client=_DisabledLangfuse()) + + response = await adapter.chat(messages=[{"role": "user", "content": "hello"}]) + + assert response.finish_reason == "error" + assert "AuthenticationError" in response.content + assert fake_vlm.calls == 1 + + +@pytest.mark.asyncio +async def test_chat_stream_retries_rate_limit_until_success(monkeypatch): + sleep_delays: list[float] = [] + + async def _sleep(delay: float): + sleep_delays.append(delay) + + monkeypatch.setattr(vlm_adapter, "rate_limit_retry_delay", lambda attempt: attempt) + monkeypatch.setattr(vlm_adapter.asyncio, "sleep", _sleep) + + chunk = SimpleNamespace( + usage=None, + choices=[ + SimpleNamespace( + finish_reason=None, + delta=SimpleNamespace(content="streamed", reasoning_content=None), + ) + ], + ) + completions = _FakeStreamingCompletions( + [RuntimeError("Error code: 429 - ModelAccountTpmRateLimitExceeded")], + [chunk], + ) + adapter = VLMProviderAdapter( + _FakeStreamingVLM(completions), + "test-model", + langfuse_client=_DisabledLangfuse(), + ) + + events = [ + event + async for event in adapter.chat_stream( + messages=[{"role": "user", "content": "hello"}], + ) + ] + + assert completions.calls == 2 + assert sleep_delays == [1] + assert [event.type for event in events] == ["content_delta", "response"] + assert events[0].content == "streamed" + assert events[1].response.content == "streamed" + assert events[1].response.finish_reason == "stop" + + +def test_rate_limit_classifier_handles_target_error(): + assert is_retryable_rate_limit_error( + RuntimeError("Error code: 429 - ModelAccountTpmRateLimitExceeded") + ) + assert is_retryable_rate_limit_error( + RuntimeError( + "Error code: 429 - {'error': {'code': 'ModelAccountTpmRateLimitExceeded', " + "'message': 'TPM (Tokens Per Minute) limit of the model doubao-seed-2-0-pro " + "is exceeded. Please try again later Request id: " + "0217817720969006061aa40146dbf4d117b0497e84060d7ac9102', " + "'param': '', 'type': 'TooManyRequests'}}, request_id: " + "202606181641366ORRzhOSo5se81lzpolL" + ) + ) + assert is_retryable_rate_limit_error(RuntimeError("Error code: 429 - busy")) + assert not is_retryable_rate_limit_error(RuntimeError("Error code: 401 Unauthorized")) + assert not is_retryable_rate_limit_error(RuntimeError("trace_id=abc429def unrelated")) + assert not is_retryable_rate_limit_error(RuntimeError("request_id=abc429def unrelated")) + + +def test_rate_limit_classifier_handles_structured_sdk_errors(): + request = httpx.Request("POST", "https://example.test") + response = httpx.Response(429, request=request) + exc = ArkRateLimitError( + "Error code: 429 - rate limited", + response=response, + body={ + "code": "ModelAccountTpmRateLimitExceeded", + "type": "TooManyRequests", + "message": "TPM limit exceeded", + }, + request_id="0217817720969006061aa40146dbf4d117b0497e84060d7ac9102", + ) + + assert is_retryable_rate_limit_error(exc) diff --git a/uv.lock b/uv.lock index 2c5d46aa0d..ecf327a3a4 100644 --- a/uv.lock +++ b/uv.lock @@ -1038,6 +1038,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4e/8c/f3147f5c4b73e7550fe5f9352eaa956ae838d5c51eb58e7a25b9f3e2643b/decorator-5.2.1-py3-none-any.whl", hash = "sha256:d316bb415a2d9e2d2b3abcc4084c6502fc09240e292cd76a76afc106a1c8e04a", size = 9190, upload-time = "2025-02-24T04:41:32.565Z" }, ] +[[package]] +name = "defusedxml" +version = "0.7.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0f/d5/c66da9b79e5bdb124974bfe172b4daf3c984ebd9c2a06e2b8a4dc7331c72/defusedxml-0.7.1.tar.gz", hash = "sha256:1bb3032db185915b62d7c6209c5a8792be6a32ab2fedacc84e01b52c51aa3e69", size = 75520, upload-time = "2021-03-08T10:59:26.269Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/07/6c/aa3f2f849e01cb6a001cd8554a88d4c77c5c1a31c95bdf1cf9301e6d9ef4/defusedxml-0.7.1-py2.py3-none-any.whl", hash = "sha256:a352e7e428770286cc899e2542b6cdaedb2b4953ff269a210103ec58f6198a61", size = 25604, upload-time = "2021-03-08T10:59:24.45Z" }, +] + [[package]] name = "diff-match-patch" version = "20241021" @@ -1267,6 +1276,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1f/f6/da4db31001e854025ffd26bc9ba0740a9cbba2c3259695f7c5834908b336/fastuuid-0.14.0-cp314-cp314-win_amd64.whl", hash = "sha256:df61342889d0f5e7a32f7284e55ef95103f2110fee433c2ae7c2c0956d76ac8a", size = 156457, upload-time = "2025-10-19T22:33:44.579Z" }, ] +[[package]] +name = "feedparser" +version = "6.0.12" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "sgmllib3k" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/dc/79/db7edb5e77d6dfbc54d7d9df72828be4318275b2e580549ff45a962f6461/feedparser-6.0.12.tar.gz", hash = "sha256:64f76ce90ae3e8ef5d1ede0f8d3b50ce26bcce71dd8ae5e82b1cd2d4a5f94228", size = 286579, upload-time = "2025-09-10T13:33:59.486Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4e/eb/c96d64137e29ae17d83ad2552470bafe3a7a915e85434d9942077d7fd011/feedparser-6.0.12-py3-none-any.whl", hash = "sha256:6bbff10f5a52662c00a2e3f86a38928c37c48f77b3c511aedcd51de933549324", size = 81480, upload-time = "2025-09-10T13:33:58.022Z" }, +] + [[package]] name = "ffmpy" version = "1.0.0" @@ -3502,8 +3523,10 @@ dependencies = [ { name = "argon2-cffi" }, { name = "charset-normalizer" }, { name = "cryptography" }, + { name = "defusedxml" }, { name = "ebooklib" }, { name = "fastapi" }, + { name = "feedparser" }, { name = "httpx" }, { name = "jinja2" }, { name = "json-repair" }, @@ -3734,10 +3757,12 @@ requires-dist = [ { name = "datasets", marker = "extra == 'eval'", specifier = ">=2.0.0" }, { name = "datasets", marker = "extra == 'test'", specifier = ">=2.0.0" }, { name = "ddgs", marker = "extra == 'bot'", specifier = ">=9.0.0" }, + { name = "defusedxml", specifier = ">=0.7.1" }, { name = "diff-match-patch", marker = "extra == 'test'", specifier = ">=20200713" }, { name = "dingtalk-stream", marker = "extra == 'bot-dingtalk'", specifier = ">=0.4.0" }, { name = "ebooklib", specifier = ">=0.18.0" }, { name = "fastapi", specifier = ">=0.128.0" }, + { name = "feedparser", specifier = ">=6.0.0" }, { name = "fusepy", marker = "extra == 'bot-fuse'", specifier = ">=3.0.1" }, { name = "google-genai", marker = "extra == 'gemini'", specifier = ">=1.0.0" }, { name = "google-genai", marker = "extra == 'gemini-async'", specifier = ">=1.0.0" }, @@ -3758,7 +3783,7 @@ requires-dist = [ { name = "langgraph", marker = "extra == 'langgraph'", specifier = ">=1.0.0,<2.0.0" }, { name = "lark-oapi", specifier = ">=1.5.3" }, { name = "lark-oapi", marker = "extra == 'bot-feishu'", specifier = ">=1.0.0" }, - { name = "litellm", specifier = ">=1.83.7,<1.86.3" }, + { name = "litellm", specifier = ">=1.83.7,<1.89.3" }, { name = "llama-cpp-python", marker = "extra == 'local-embed'", specifier = ">=0.3.0" }, { name = "loguru", specifier = ">=0.7.3" }, { name = "markdownify", specifier = ">=0.11.0" }, @@ -5662,6 +5687,12 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5c/e1/342c4434df56aa537f6ce7647eefee521d96fbb828b08acd709865767652/setuptools_scm-10.0.5-py3-none-any.whl", hash = "sha256:f611037d8aae618221503b8fa89319f073438252ae3420e01c9ceec249131a0a", size = 21695, upload-time = "2026-03-27T15:57:03.969Z" }, ] +[[package]] +name = "sgmllib3k" +version = "1.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9e/bd/3704a8c3e0942d711c1299ebf7b9091930adae6675d7c8f476a7ce48653c/sgmllib3k-1.0.0.tar.gz", hash = "sha256:7868fb1c8bfa764c1ac563d3cf369c381d1325d36124933a726f29fcdaa812e9", size = 5750, upload-time = "2010-08-24T14:33:52.445Z" } + [[package]] name = "shellingham" version = "1.5.4"