From fea3bdef94528821606e6b0b03f66e55c74d393f Mon Sep 17 00:00:00 2001 From: chenzhwei Date: Thu, 13 Aug 2026 18:53:21 +0800 Subject: [PATCH] feat(memory): integrate multimodal video memory --- bootstrap/cli/client.py | 3 + bootstrap/core/handler.py | 101 +- bootstrap/core/server.py | 18 + bootstrap/http_server/__main__.py | 1 + examples/config_multimodal.yml | 66 + .../normalizer/normalizer_impl/__init__.py | 2 + .../normalizer_impl/routing_normalizer.py | 87 + .../normalizer/normalizer_impl/video_asr.py | 667 ++++ .../normalizer_impl/video_models.py | 175 ++ .../normalizer_impl/video_normalizer.py | 271 ++ .../normalizer_impl/video_pipeline.py | 2766 +++++++++++++++++ .../normalizer_impl/video_prompts.py | 806 +++++ .../construction/extractor_impl/__init__.py | 1 + .../extractor_impl/video_memory_extractor.py | 200 ++ .../control/engine_impl/cloud_engine.py | 11 +- .../control/engine_impl/in_memory_engine.py | 7 +- jiuwen_memory/control/ingest_job.py | 287 ++ .../retrieval/retriever_impl/__init__.py | 1 + .../retriever_impl/multimodal_retriever.py | 203 ++ pyproject.toml | 5 + .../test_video_memory_extractor.py | 57 + tests/unit/control/test_ingest_job.py | 90 + .../multimodal/test_multimodal_adapter.py | 393 +++ uv.lock | 32 +- 24 files changed, 6241 insertions(+), 9 deletions(-) create mode 100644 examples/config_multimodal.yml create mode 100644 jiuwen_memory/common/normalizer/normalizer_impl/routing_normalizer.py create mode 100644 jiuwen_memory/common/normalizer/normalizer_impl/video_asr.py create mode 100644 jiuwen_memory/common/normalizer/normalizer_impl/video_models.py create mode 100644 jiuwen_memory/common/normalizer/normalizer_impl/video_normalizer.py create mode 100644 jiuwen_memory/common/normalizer/normalizer_impl/video_pipeline.py create mode 100644 jiuwen_memory/common/normalizer/normalizer_impl/video_prompts.py create mode 100644 jiuwen_memory/construction/extractor_impl/video_memory_extractor.py create mode 100644 jiuwen_memory/control/ingest_job.py create mode 100644 jiuwen_memory/retrieval/retriever_impl/multimodal_retriever.py create mode 100644 tests/unit/construction/test_video_memory_extractor.py create mode 100644 tests/unit/control/test_ingest_job.py create mode 100644 tests/unit/multimodal/test_multimodal_adapter.py diff --git a/bootstrap/cli/client.py b/bootstrap/cli/client.py index 192e5469..3b80baed 100644 --- a/bootstrap/cli/client.py +++ b/bootstrap/cli/client.py @@ -81,6 +81,9 @@ def call(self, verb: str, payload: dict[str, Any]) -> tuple[int, dict[str, Any]] def healthz(self) -> tuple[int, dict[str, Any]]: return 200, {"status": "ok", "profile": self._srv.config.profile} + def close(self) -> None: + self._srv.close(wait=True) + class HttpClient: """Drive a running ``bootstrap`` server over HTTP (``POST /v1/``).""" diff --git a/bootstrap/core/handler.py b/bootstrap/core/handler.py index de03e320..ae4b6dc4 100644 --- a/bootstrap/core/handler.py +++ b/bootstrap/core/handler.py @@ -17,6 +17,7 @@ import os import sys +import uuid from datetime import datetime from importlib import import_module from typing import Any, Callable @@ -45,6 +46,7 @@ Modality = _type_def_module.Modality Scope = _type_def_module.Scope EvolveMode = import_module("jiuwen_memory.construction").EvolveMode +INGEST_JOB_PREFIX = import_module("jiuwen_memory.control.ingest_job").INGEST_JOB_PREFIX _control_types_module = import_module("jiuwen_memory.control.types") Action = _control_types_module.Action @@ -371,12 +373,103 @@ def _usage_view(usage) -> Body: } +def _video_unit_view(unit: MemoryUnit) -> Body: + """Serialize fields needed by video job responses.""" + return { + **_unit_view(unit), + "source": unit.source.value, + "source_ref": unit.source_ref, + "provenance": list(unit.provenance), + "metadata": dict(unit.metadata), + } + + +def _submit_video(srv, payload: Body, *, scope: Scope, identity: Scope) -> Body: + """Submit a video write through the Control-managed ingest queue.""" + uri = str(_require(payload, "uri")).strip() + if not uri: + raise ValidationError("missing required field: 'uri'") + payload_id = str(payload.get("payload_id") or uuid.uuid4()) + raw_meta = payload.get("metadata") + if not isinstance(raw_meta, dict): + raw_meta = {} + metadata = dict(raw_meta) + metadata.update( + {"infer": "true", "pipeline": "video", "payload_id": payload_id} + ) + requested_assets = payload.get("assets") + extras = requested_assets if isinstance(requested_assets, list) else [] + assets = [uri, *(str(item) for item in extras if str(item) != uri)] + + submission = srv.ingest_jobs.submit( + payload_id=payload_id, + source_ref=uri, + scope=scope, + task=lambda: srv.api.add( + uri, + scope, + Modality.VIDEO, + identity=identity, + assets=assets, + tags=payload.get("tags"), + metadata=metadata, + ), + ) + job = submission.job + return { + "ok": True, + "op": "add", + "accepted": True, + "job_id": job.id, + "video_id": payload_id, + "status": job.status, + "reused": submission.reused, + "feedback_message": ( + "已返回该视频现有的处理任务。" + if submission.reused + else "视频处理任务已提交。" + ), + } + + +def _ingest_job_status(srv, job, *, identity: Scope) -> Body: + """Adapt a Control ingest job to the shared job response shape.""" + scope = job.scope + units = [] + for unit_id in job.unit_ids: + try: + units.append(srv.api.get(unit_id, scope, identity=identity)) + except NotFoundError: + continue + items = [_video_unit_view(unit) for unit in units] + body: Body = { + "ok": True, + "op": "job", + "job_id": job.id, + "video_id": job.payload_id, + "status": job.status, + "count": len(items), + "item_ids": [item["item_id"] for item in items], + "items": items, + } + if job.status == "succeeded": + body["feedback_message"] = f"视频处理完成，共生成 {len(items)} 条多模态记忆。" + elif job.status == "failed": + body["error"] = job.error + body["feedback_message"] = "视频处理失败。" + else: + body["feedback_message"] = "视频正在处理中。" + return body + + # --- per-verb handlers ----------------------------------------------------- # def _add(srv, payload: Body) -> Body: scope, actor = _target_scope(payload), _actor_scope(payload) modality = Modality(payload.get("modality", "text")) + if modality == Modality.VIDEO: + return _submit_video(srv, payload, scope=scope, identity=actor) # metadata 透传：infer 等调用级开关经 metadata 下推到引擎（engine.write 从 # metadata["infer"]=="true" 判定是否同步走 evolve(EXTRACT) 抽取派生记忆）。 # JSON 标量原样透传（不 str 化）：数值/布尔要保持原生类型才能在索引里建 @@ -613,9 +706,13 @@ def _evolve(srv, payload: Body) -> Body: def _job(srv, payload: Body) -> Body: - """查询演进任务状态（Scheduler）。""" + """查询视频 Ingest 任务或原生 Scheduler 任务状态。""" + job_id = str(_require(payload, "job_id")) actor = _actor_scope(payload) - info = srv.api.job_status(_require(payload, "job_id"), identity=actor) + if job_id.startswith(INGEST_JOB_PREFIX): + info = srv.ingest_jobs.status(job_id, scope=_target_scope(payload)) + return _ingest_job_status(srv, info, identity=actor) + info = srv.api.job_status(job_id, identity=actor) return { "ok": True, "op": "job", diff --git a/bootstrap/core/server.py b/bootstrap/core/server.py index 43832f4d..c50fd4fb 100644 --- a/bootstrap/core/server.py +++ b/bootstrap/core/server.py @@ -35,6 +35,9 @@ Kernel = _api_module.Kernel build_kernel = _api_module.build_kernel KernelConfig = import_module("jiuwen_memory.config").Config +IngestJobController = import_module( + "jiuwen_memory.control.ingest_job" +).IngestJobController class Server: @@ -43,6 +46,17 @@ class Server: def __init__(self, config: Config, kernel: Kernel) -> None: self.config = config self.kernel = kernel + memory_config = config.settings.get("memory_api", {}) + globals_config = ( + memory_config.get("globals", {}) + if isinstance(memory_config, dict) + else {} + ) + self.ingest_jobs = IngestJobController( + max_workers=int(globals_config.get("ingest_max_workers", 1)), + max_pending_jobs=int(globals_config.get("ingest_max_pending_jobs", 2)), + kv=kernel.kv, + ) @property def api(self): @@ -75,6 +89,10 @@ def dispatch(self, verb: str, payload: Dict[str, Any]) -> Tuple[int, Dict[str, A return _dispatch(self, verb, payload) + def close(self, *, wait: bool = True) -> None: + """Release the Control-owned ingest worker pool.""" + self.ingest_jobs.close(wait=wait) + def default_spaces() -> Dict[str, Any]: """Default scope/namespace registry (none needed for the in-memory build).""" diff --git a/bootstrap/http_server/__main__.py b/bootstrap/http_server/__main__.py index 75de8f7c..52d8bc0f 100644 --- a/bootstrap/http_server/__main__.py +++ b/bootstrap/http_server/__main__.py @@ -90,6 +90,7 @@ def serve(self, host: str, port: int) -> None: sys.stderr.write("\nagent-memory server stopped\n") finally: httpd.server_close() + self.close(wait=True) def main(argv: list[str] | None = None) -> int: diff --git a/examples/config_multimodal.yml b/examples/config_multimodal.yml new file mode 100644 index 00000000..64bfb26d --- /dev/null +++ b/examples/config_multimodal.yml @@ -0,0 +1,66 @@ +memory_api: + globals: + ingest_max_workers: 1 + ingest_max_pending_jobs: 2 + + normalizer: + default: + target: routing + params: + routes: + video: + target: video + params: + whisper_model_dir: /path/to/whisper-model + whisper_batch_size: 1 + vllm_base_url: http://127.0.0.1:8000/v1 + vllm_api_key: dummy + llm_model: qwen-vl + temp_root: /tmp/agent-memory-video + + extractor: + video: + target: video_memory + + retriever: + multimodal: + target: multimodal + params: + base_retriever: default + kv_store: default + clip_top_k: 10 + event_top_k: 10 + rrf_k: 60 + + evolver: + video: + target: orchestrating + params: + extractor: video + abstractor: default + associator: default + index_builder: default + kv_store: default + graph_store: default + dedup: default + llm: default + + pipeline: + default: + target: metadata + params: + route_key: pipeline + fallback: default + routes: + video: video + profiles: + default: + index_builder: default + retriever: multimodal + evolver: default + classifier: default + video: + index_builder: default + retriever: multimodal + evolver: video + classifier: default diff --git a/jiuwen_memory/common/normalizer/normalizer_impl/__init__.py b/jiuwen_memory/common/normalizer/normalizer_impl/__init__.py index 95ada386..0d194e2c 100644 --- a/jiuwen_memory/common/normalizer/normalizer_impl/__init__.py +++ b/jiuwen_memory/common/normalizer/normalizer_impl/__init__.py @@ -8,5 +8,7 @@ from jiuwen_memory.common.normalizer.base import NormalizerProducer import_module(".passthrough_normalizer", __name__) +import_module(".routing_normalizer", __name__) +import_module(".video_normalizer", __name__) __all__ = ["NormalizerProducer"] diff --git a/jiuwen_memory/common/normalizer/normalizer_impl/routing_normalizer.py b/jiuwen_memory/common/normalizer/normalizer_impl/routing_normalizer.py new file mode 100644 index 00000000..d39072e9 --- /dev/null +++ b/jiuwen_memory/common/normalizer/normalizer_impl/routing_normalizer.py @@ -0,0 +1,87 @@ +"""Route raw payloads to a modality-specific normalizer.""" + +from __future__ import annotations + +from collections.abc import Mapping + +from jiuwen_memory.common.base import PluginType +from jiuwen_memory.common.errors import ValidationError +from jiuwen_memory.common.normalizer.base import Normalizer, NormalizerProducer +from jiuwen_memory.common.type_def import Modality, RawPayload + + +class RoutingNormalizer(Normalizer): + """Delegate normalization while keeping one Ingestor entry point.""" + + def __init__( + self, + fallback: Normalizer, + routes: Mapping[Modality, Normalizer], + ) -> None: + self._fallback = fallback + self._routes = dict(routes) + + def modalities(self) -> list[Modality]: + supported = set(self._fallback.modalities()) + supported.update(self._routes) + return sorted(supported, key=lambda item: item.value) + + def plugin_type(self) -> PluginType: + return PluginType.NORMALIZER + + def health(self) -> None: + seen: set[int] = set() + for normalizer in [self._fallback, *self._routes.values()]: + if id(normalizer) in seen: + continue + seen.add(id(normalizer)) + normalizer.health() + + def normalize(self, payload: RawPayload) -> str: + normalizer = self._routes.get(payload.modality, self._fallback) + if payload.modality not in normalizer.modalities(): + raise ValidationError( + f"no normalizer configured for modality {payload.modality.value!r}" + ) + return normalizer.normalize(payload) + + +def _build_normalizer(config, value: object, *, field: str) -> Normalizer: + if isinstance(value, str): + return NormalizerProducer.build_named(value, config.ctx) + if isinstance(value, Mapping): + target = str(value.get("target", "")).strip() + if not target: + raise ValidationError(f"routing normalizer {field!r} is missing target") + return NormalizerProducer.build( + target, + value.get("params", {}), + config.ctx, + name=str(value.get("name", "")), + ) + raise ValidationError( + f"routing normalizer {field!r} must be a named reference or component mapping" + ) + + +@NormalizerProducer.register("routing") +def _build(config): + fallback_raw = config.get("fallback", {"target": "passthrough"}) + fallback = _build_normalizer(config, fallback_raw, field="fallback") + routes_raw = config.get("routes", {}) + if not isinstance(routes_raw, Mapping): + raise ValidationError("routing normalizer routes must be a mapping") + routes: dict[Modality, Normalizer] = {} + for modality_name, raw in routes_raw.items(): + try: + modality = Modality(str(modality_name)) + except ValueError as exc: + raise ValidationError( + f"unknown routing normalizer modality {modality_name!r}" + ) from exc + routes[modality] = _build_normalizer( + config, + raw, + field=f"routes.{modality.value}", + ) + return RoutingNormalizer(fallback, routes) diff --git a/jiuwen_memory/common/normalizer/normalizer_impl/video_asr.py b/jiuwen_memory/common/normalizer/normalizer_impl/video_asr.py new file mode 100644 index 00000000..7584c6b9 --- /dev/null +++ b/jiuwen_memory/common/normalizer/normalizer_impl/video_asr.py @@ -0,0 +1,667 @@ +# ruff: noqa: E501 + +import gc +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch + +from jiuwen_memory.common.log import get_logger + +logger = get_logger(__name__) + + +@dataclass(frozen=True) +class VideoAsrConfig: + """Settings and optional output paths for one video ASR run.""" + + model_dir: Path | str | None + device: str | int | None = "auto" + language: str | None = None + batch_size: int = 1 + chunk_seconds: int = 600 + max_pause_s: float = 0.35 + max_segment_s: float = 6.0 + break_on_punct: bool = True + min_gap_s: float = 0.0 + vad_filter: bool = True + vad_min_silence_ms: int = 300 + output_json: str | Path | None = None + raw_txt_path: str | Path | None = None + cleaned_txt_path: str | Path | None = None + temp_work_dir: str | Path | None = None + cleanup: bool = True + + +def _valid_word_timestamp(timestamp: object) -> bool: + if not timestamp or not isinstance(timestamp, (list, tuple)): + return False + if len(timestamp) != 2: + return False + return timestamp[0] is not None and timestamp[1] is not None + + +def _silero_vad_timestamps( + audio_path: Path, + min_silence_duration_ms: int = 300, + threshold: float = 0.5, +) -> Optional[List[Tuple[float, float]]]: + """ + Return speech regions as [(start_s, end_s), ...] using Silero VAD. + Returns None if silero-vad is not installed (caller treats as "skip VAD"). + """ + try: + from silero_vad import get_speech_timestamps, load_silero_vad + except ImportError: + logger.info("VideoASR: silero-vad unavailable; VAD filtering disabled") + return None + + model = load_silero_vad() + # Use soundfile to avoid torchaudio 2.9+ torchcodec dependency + # audio_path is always 16kHz mono WAV (created by ensure_audio) + import soundfile as sf + import torch as _torch + + samples, sr = sf.read(str(audio_path), dtype="float32", always_2d=False) + wav = _torch.from_numpy(samples) + if sr != 16000: + # fallback resample if somehow sr differs + import torchaudio + + wav = torchaudio.functional.resample(wav.unsqueeze(0), sr, 16000).squeeze(0) + + timestamps = get_speech_timestamps( + wav, + model, + sampling_rate=16000, + min_silence_duration_ms=min_silence_duration_ms, + threshold=threshold, + return_seconds=True, + ) + return [(float(t["start"]), float(t["end"])) for t in timestamps] + + +def _apply_vad_mask_to_chunk( + chunk_path: Path, + chunk_start_s: float, + speech_regions: List[Tuple[float, float]], +) -> None: + """ + Zero out non-speech samples in a chunk WAV in-place. + speech_regions are absolute times (relative to full audio start). + Preserves audio duration so timestamps stay valid. + """ + import numpy as np + import soundfile as sf + + samples, sr = sf.read(str(chunk_path), dtype="float32", always_2d=False) + mask = np.zeros(len(samples), dtype=np.float32) + + for seg_start, seg_end in speech_regions: + i0 = max(0, int((seg_start - chunk_start_s) * sr)) + i1 = min(len(samples), int((seg_end - chunk_start_s) * sr)) + if i1 > i0: + mask[i0:i1] = 1.0 + + samples *= mask + sf.write(str(chunk_path), samples, sr) + + +def ensure_audio(video_path: Path, output_dir: Path) -> Path: + """提取视频中的音频并转为16kHz单声道WAV""" + import subprocess + + audio_path = output_dir / f"{video_path.stem}.wav" + cmd = [ + "ffmpeg", + "-y", + "-i", + str(video_path), + "-ac", + "1", + "-ar", + "16000", + "-vn", + str(audio_path), + ] + subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True) + return audio_path + + +DEFAULT_MODEL_DIR: Path | None = None + + +def _is_npu_available() -> bool: + """Return True when torch_npu backend is available and has at least one device.""" + try: + import torch_npu # noqa: F401 # injects torch.npu into torch namespace + except ImportError: + return False + try: + return torch.npu.is_available() and int(torch.npu.device_count()) > 0 + except (AttributeError, RuntimeError): + return False + + +def _auto_detect_asr_device() -> str: + """Prefer NPU, then CUDA, finally CPU.""" + if _is_npu_available(): + return "npu:0" + if torch.cuda.is_available(): + return "cuda:0" + return "cpu" + + +def _parse_device_index(device_str: str, backend: str) -> int: + if ":" not in device_str: + return 0 + raw = device_str.split(":", 1)[1].strip() + if raw == "": + return 0 + try: + return int(raw) + except (TypeError, ValueError) as e: + raise ValueError(f"Invalid {backend} device index in '{device_str}'.") from e + + +def _resolve_asr_device( + device: Optional[Union[str, int]], +) -> Tuple[Union[int, str], str]: + """ + Resolve device for transformers.pipeline. + + Returns: + pipeline_device: value passed into pipeline(device=...) + normalized_device: normalized label for logging + """ + if isinstance(device, int): + if device < 0: + return "cpu", "cpu" + if torch.cuda.is_available(): + return int(device), f"cuda:{device}" + if _is_npu_available(): + return torch.device(f"npu:{device}"), f"npu:{device}" + return "cpu", "cpu" + + if device is None: + d = "auto" + else: + d = str(device).strip().lower() + if not d: + d = "auto" + + if d == "auto": + d = _auto_detect_asr_device() + + if d.startswith("npu"): + if not _is_npu_available(): + raise RuntimeError( + "NPU is not available; ensure torch_npu is installed and NPU is visible." + ) + idx = _parse_device_index(d, "npu") + return torch.device(f"npu:{idx}"), f"npu:{idx}" + + if d.startswith("cuda"): + if not torch.cuda.is_available(): + raise RuntimeError( + "CUDA is not available; ensure the environment exposes a GPU." + ) + idx = _parse_device_index(d, "cuda") + return int(idx), f"cuda:{idx}" + + if d == "cpu": + return "cpu", "cpu" + + raise ValueError(f"Unknown device string: {device}") + + +def _get_clip_duration(clip_path: Path) -> float: + import shutil + import subprocess + + if shutil.which("ffprobe") is None: + raise FileNotFoundError("ffprobe not found in PATH. Please install ffmpeg.") + cmd = [ + "ffprobe", + "-v", + "error", + "-show_entries", + "format=duration", + "-of", + "default=noprint_wrappers=1:nokey=1", + str(clip_path), + ] + res = subprocess.run( + cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True + ) + out = res.stdout.decode().strip() + try: + return float(out) + except (TypeError, ValueError) as e: + raise RuntimeError(f"Failed to parse duration for {clip_path}: {out}") from e + + +def _seconds_to_hhmmss_int(sec: float) -> str: + """输出到整数秒 HH:MM:SS（按四舍五入）""" + sec = max(0.0, float(sec)) + t = int(round(sec)) + h = t // 3600 + m = (t % 3600) // 60 + s = t % 60 + return f"{h:02d}:{m:02d}:{s:02d}" + + +def _split_audio_to_chunks( + audio_path: Path, + out_dir: Path, + chunk_seconds: int = 600, +) -> List[Tuple[Path, float, float]]: + """ + 更稳的切片： + - 用 -ss 在输入前 + -t 指定时长（避免 -to 坑） + - asetpts=PTS-STARTPTS 强制分片从0开始 + """ + import subprocess + + out_dir.mkdir(parents=True, exist_ok=True) + duration = _get_clip_duration(audio_path) + if duration <= 0: + return [(audio_path, 0.0, duration)] + + chunks: List[Tuple[Path, float, float]] = [] + idx = 0 + start = 0.0 + + while start < duration: + end = min(duration, start + float(chunk_seconds)) + seg_dur = max(0.0, end - start) + + out_path = out_dir / f"audio_chunk_{idx:04d}.wav" + cmd = [ + "ffmpeg", + "-y", + "-ss", + f"{start:.3f}", + "-i", + str(audio_path), + "-t", + f"{seg_dur:.3f}", + "-ac", + "1", + "-ar", + "16000", + "-af", + "asetpts=PTS-STARTPTS", + str(out_path), + ] + subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True) + + chunk_duration = _get_clip_duration(out_path) + actual_end = start + chunk_duration + chunks.append((out_path, start, actual_end)) + + idx += 1 + start = end + + return chunks + + +def _group_words_to_segments( + word_chunks: List[Dict[str, Any]], + max_pause_s: float = 0.35, + max_segment_s: float = 6.0, + break_on_punct: bool = True, +) -> List[Dict[str, Any]]: + """ + word-level -> 段落： + - gap > max_pause_s 断句 + - 段长 >= max_segment_s 断句 + - 可选：句末标点断句 + """ + items: List[Dict[str, Any]] = [] + cur_text: List[str] = [] + cur_start: Optional[float] = None + last_end: Optional[float] = None + + end_punct = (".", "?", "!", "。", "？", "！") + + def flush() -> None: + nonlocal cur_text, cur_start, last_end + if cur_text and cur_start is not None and last_end is not None: + text = "".join(cur_text).strip() + if text: + items.append({"start": cur_start, "end": last_end, "text": text}) + cur_text = [] + cur_start = None + last_end = None + + for w in word_chunks: + text = w.get("text") or "" + ts = w.get("timestamp") + if not _valid_word_timestamp(ts): + continue + + st = float(ts[0]) + ed = float(ts[1]) + if ed < st: + continue + + if cur_start is None: + cur_start = st + last_end = ed + cur_text.append(text) + else: + gap = max(0.0, st - (last_end if last_end is not None else st)) + seg_len = (last_end if last_end is not None else ed) - ( + cur_start if cur_start is not None else st + ) + + if gap > max_pause_s or seg_len >= max_segment_s: + flush() + cur_start = st + last_end = ed + cur_text.append(text) + else: + cur_text.append(text) + last_end = max(last_end or 0.0, ed) + + if break_on_punct and text.strip().endswith(end_punct): + flush() + + flush() + return items + + +def _fix_monotonic_timestamps( + segments: List[Dict[str, Any]], + min_gap_s: float = 0.0, +) -> List[Dict[str, Any]]: + """ + 规则： + 1) 单句 end < start => end = start + 2) 后一句 start < 前一句 end + min_gap => start 推到该位置 + 然后若 end < start => end = start + """ + fixed: List[Dict[str, Any]] = [] + prev_end: Optional[float] = None + + for seg in segments: + st = float(seg.get("start", 0.0)) + ed = float(seg.get("end", st)) + text = str(seg.get("text", "")).strip() + if not text: + continue + + # 同一句结束小于开始 + if ed < st: + ed = st + + # 与前一句重叠 + if prev_end is not None: + need_start = prev_end + float(min_gap_s) + if st < need_start: + shift = need_start - st + st = need_start + ed = max( + ed + shift, st + ) # 把整句一起向后挪，保持相对时长；同时保证 ed>=st + + fixed.append({"start": st, "end": ed, "text": text}) + prev_end = ed + + return fixed + + +def _run_full_asr( + audio_path: Path, + config: VideoAsrConfig, + return_raw: bool = False, +) -> Union[List[Dict[str, str]], Tuple[List[Dict[str, str]], List[Dict[str, Any]]]]: + """Run Whisper and output sentence-like segments with corrected timestamps.""" + from transformers import pipeline + + model_dir = Path(config.model_dir) if config.model_dir is not None else None + if model_dir is None: + raise ValueError("ASR model_dir is required") + pipeline_device, normalized_device = _resolve_asr_device(config.device) + logger.info("VideoASR: using device=%s", normalized_device) + + if isinstance(pipeline_device, int) and torch.cuda.is_available(): + try: + torch.cuda.set_device(pipeline_device) + except (RuntimeError, ValueError) as exc: + logger.warning("VideoASR: failed to select CUDA device: %s", exc) + elif ( + isinstance(pipeline_device, torch.device) + and pipeline_device.type == "npu" + and _is_npu_available() + ): + try: + torch.npu.set_device(pipeline_device) + except (RuntimeError, ValueError) as exc: + logger.warning("VideoASR: failed to select NPU device: %s", exc) + + asr_pipe = pipeline( + task="automatic-speech-recognition", + model=str(model_dir), + dtype="auto", + device=pipeline_device, + return_timestamps="word", + generate_kwargs={ + "task": "transcribe", + **({"language": config.language} if config.language else {}), + }, + ) + _first_param = next(asr_pipe.model.parameters()) + logger.info("VideoASR: model loaded on device=%s", _first_param.device) + + segments_raw: List[Dict[str, Any]] = [] + raw_outputs: List[Dict[str, Any]] = [] + + chunk_jobs = _split_audio_to_chunks( + audio_path, + audio_path.parent / "asr_chunks", + chunk_seconds=config.chunk_seconds, + ) + + # VAD: zero out non-speech samples in each chunk to prevent hallucinations. + # Audio duration is unchanged so timestamps stay valid. + if config.vad_filter: + speech_regions = _silero_vad_timestamps( + audio_path, min_silence_duration_ms=config.vad_min_silence_ms + ) + if speech_regions is not None: + total_speech = sum(e - s for s, e in speech_regions) + logger.info( + "VideoASR: VAD found %d speech segments (%.1fs/%.1fs)", + len(speech_regions), + total_speech, + sum(c[2] - c[1] for c in chunk_jobs), + ) + for chunk_path, chunk_start, _chunk_end in chunk_jobs: + _apply_vad_mask_to_chunk(chunk_path, chunk_start, speech_regions) + + asr_bs = max(1, int(config.batch_size)) + + for i in range(0, len(chunk_jobs), asr_bs): + job_batch = chunk_jobs[i:i + asr_bs] + batch_paths = [str(j[0]) for j in job_batch] + + try: + results = asr_pipe(batch_paths, batch_size=max(1, len(job_batch))) + if isinstance(results, dict): + results = [results] + except RuntimeError as e: + err_msg = str(e) + size_mismatch = ( + "expanded size of the tensor" in err_msg.lower() + or "must match the existing size" in err_msg.lower() + ) + if not size_mismatch: + raise + + # Some NPU stacks are unstable on variable-length mixed batches. + # Fallback to per-chunk decode to keep the job progressing. + results = [] + for chunk_path in batch_paths: + single = asr_pipe(chunk_path, batch_size=1) + if isinstance(single, list): + results.extend(single) + else: + results.append(single) + + for (chunk_path, st_offset, actual_end), result in zip(job_batch, results): + if return_raw: + raw_outputs.append( + { + "chunk_path": str(chunk_path), + "offset_s": float(st_offset), + "actual_end_s": float(actual_end), + "result": result, + } + ) + + # return_timestamps="word" -> result["chunks"] 为词级 + word_chunks = result.get("chunks") or [] + if not word_chunks: + continue + + wc: List[Dict[str, Any]] = [] + for w in word_chunks: + ts = w.get("timestamp") + if not _valid_word_timestamp(ts): + continue + if float(ts[1]) < float(ts[0]): + continue + wc.append( + { + "timestamp": ( + float(ts[0]) + st_offset, + float(ts[1]) + st_offset, + ), + "text": w.get("text") or "", + } + ) + + wc.sort(key=lambda x: (x["timestamp"][0], x["timestamp"][1])) + segs = _group_words_to_segments( + wc, + max_pause_s=config.max_pause_s, + max_segment_s=config.max_segment_s, + break_on_punct=config.break_on_punct, + ) + segments_raw.extend(segs) + + # 释放显存 + del asr_pipe + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if _is_npu_available(): + try: + torch.npu.empty_cache() + except RuntimeError as exc: + logger.warning("VideoASR: failed to clear NPU cache: %s", exc) + + # 先排序（按 start） + segments_raw_sorted = sorted(segments_raw, key=lambda s: float(s.get("start", 0.0))) + + # 修正时间戳单调性（你的规则） + fixed_float = _fix_monotonic_timestamps( + segments_raw_sorted, + min_gap_s=config.min_gap_s, + ) + + # 输出到整数秒字符串 + cleaned: List[Dict[str, str]] = [] + for seg in fixed_float: + cleaned.append( + { + "start": _seconds_to_hhmmss_int(seg["start"]), + "end": _seconds_to_hhmmss_int(seg["end"]), + "text": seg["text"], + } + ) + + return (cleaned, raw_outputs) if return_raw else cleaned + + +def _write_raw_outputs(raw_outputs: List[Dict[str, Any]], txt_path: Path) -> None: + txt_path.parent.mkdir(parents=True, exist_ok=True) + with open(txt_path, "w", encoding="utf-8") as f: + for item in raw_outputs: + f.write(json.dumps(item, ensure_ascii=False)) + f.write("\n") + + +def _write_cleaned_segments(cleaned: List[Dict[str, str]], txt_path: Path) -> None: + txt_path.parent.mkdir(parents=True, exist_ok=True) + with open(txt_path, "w", encoding="utf-8") as f: + for seg in cleaned: + f.write( + f"[{seg.get('start', '')}-{seg.get('end', '')}]: {seg.get('text', '')}\n" + ) + + +def run_video_asr( + video_path: str | Path, + config: VideoAsrConfig, +) -> List[Dict[str, str]]: + video_path = Path(video_path) + if config.model_dir is None: + raise ValueError( + "whisper_model_dir is required; configure it in the video normalizer " + "configuration" + ) + output_json_path = Path(config.output_json) if config.output_json else None + raw_txt = Path(config.raw_txt_path) if config.raw_txt_path else None + cleaned_txt = Path(config.cleaned_txt_path) if config.cleaned_txt_path else None + + # Use per-task temp dir to avoid cross-process file clobbering when batch-running videos. + if config.temp_work_dir is not None: + asr_tmp_dir = Path(config.temp_work_dir) + elif output_json_path is not None: + asr_tmp_dir = output_json_path.parent / "_asr_tmp" + else: + asr_tmp_dir = video_path.parent / f".asr_tmp_{video_path.stem}" + asr_tmp_dir.mkdir(parents=True, exist_ok=True) + + audio_path = ensure_audio(video_path, asr_tmp_dir) + + if raw_txt or cleaned_txt: + asr_segments, raw_outputs = _run_full_asr(audio_path, config, return_raw=True) + if raw_txt: + _write_raw_outputs(raw_outputs, raw_txt) + if cleaned_txt: + _write_cleaned_segments(asr_segments, cleaned_txt) + else: + asr_segments = _run_full_asr(audio_path, config, return_raw=False) + + if output_json_path is not None: + output_json_path.parent.mkdir(parents=True, exist_ok=True) + with open(output_json_path, "w", encoding="utf-8") as f: + json.dump(asr_segments, f, ensure_ascii=False, indent=2) + + if config.cleanup: + import shutil + + try: + asr_chunks_dir = audio_path.parent / "asr_chunks" + if asr_chunks_dir.exists(): + shutil.rmtree(asr_chunks_dir, ignore_errors=True) + except OSError as exc: + logger.warning("VideoASR: failed to remove ASR chunks: %s", exc) + try: + if audio_path.exists() and audio_path.is_file(): + audio_path.unlink(missing_ok=True) + except OSError as exc: + logger.warning("VideoASR: failed to remove temporary audio: %s", exc) + try: + if asr_tmp_dir.exists(): + shutil.rmtree(asr_tmp_dir, ignore_errors=True) + except OSError as exc: + logger.warning("VideoASR: failed to remove ASR temporary directory: %s", exc) + + return asr_segments diff --git a/jiuwen_memory/common/normalizer/normalizer_impl/video_models.py b/jiuwen_memory/common/normalizer/normalizer_impl/video_models.py new file mode 100644 index 00000000..f5a1448f --- /dev/null +++ b/jiuwen_memory/common/normalizer/normalizer_impl/video_models.py @@ -0,0 +1,175 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Dict, List, Tuple +from uuid import UUID + + +def _ensure_float_list(values: List[float]) -> List[float]: + if not isinstance(values, list): + raise TypeError("embedding must be a list of floats") + out: List[float] = [] + for v in values: + if isinstance(v, (int, float)): + out.append(float(v)) + else: + raise TypeError("embedding must contain only numeric values") + return out + + +def _ensure_time_tuple(tp: Tuple[float, float]) -> Tuple[float, float]: + if not (isinstance(tp, (list, tuple)) and len(tp) == 2): + raise TypeError("time range/span must be a tuple/list of two numbers") + start, end = tp + if not isinstance(start, (int, float)) or not isinstance(end, (int, float)): + raise TypeError("time range/span values must be numeric") + start_f, end_f = float(start), float(end) + if end_f < start_f: + raise ValueError("end time must be >= start time") + return (start_f, end_f) + + +@dataclass +class ShortTermMemory: + """ + Short-Term Memory / Atomic Clip + + Minimal atomic storage unit for a single video clip (e.g., 30s). + Captures detailed perception and a concise visual summary for similarity. + """ + + id: UUID + video_source_path: str + time_range: Tuple[float, float] + visual_summary: str + detailed_caption: str + embedding: List[float] + asr: str + environment: str + inferred_intent: str = "" + + def __post_init__(self) -> None: + if not isinstance(self.id, UUID): + raise TypeError("id must be a UUID") + if not isinstance(self.video_source_path, str) or not self.video_source_path: + raise TypeError("video_source_path must be a non-empty string") + self.time_range = _ensure_time_tuple(self.time_range) + if not isinstance(self.visual_summary, str): + raise TypeError("visual_summary must be a string") + if not isinstance(self.inferred_intent, str): + raise TypeError("inferred_intent must be a string") + if not isinstance(self.detailed_caption, str): + raise TypeError("detailed_caption must be a string") + self.embedding = _ensure_float_list(self.embedding) + if not isinstance(self.asr, str): + raise TypeError("asr must be a string") + if not isinstance(self.environment, str): + raise TypeError("environment must be a string") + + def duration(self) -> float: + start, end = self.time_range + return float(end - start) + + def to_dict(self) -> Dict[str, Any]: + return { + "id": str(self.id), + "video_source_path": self.video_source_path, + "time_range": [self.time_range[0], self.time_range[1]], + "visual_summary": self.visual_summary, + "detailed_caption": self.detailed_caption, + "embedding": list(self.embedding), + "ASR": self.asr, + "environment": self.environment, + } + + @staticmethod + def from_dict(d: Dict[str, Any]) -> ShortTermMemory: + return ShortTermMemory( + id=UUID(d["id"]) if not isinstance(d.get("id"), UUID) else d["id"], + video_source_path=d["video_source_path"], + time_range=_ensure_time_tuple(tuple(d["time_range"])), + visual_summary=d["visual_summary"], + detailed_caption=d["detailed_caption"], + embedding=_ensure_float_list(list(d["embedding"])), + asr=d.get("ASR", ""), + environment=d.get("environment", ""), + ) + + +@dataclass +class MediumTermMemory: + """ + Medium-Term Memory / Task Session + + Represents a complete event or task, aggregating consecutive clips + into a coherent narrative with evidence via child clip references. + """ + + task_id: UUID + topic: str + time_span: Tuple[float, float] + narrative_summary: str + child_clip_ids: List[UUID] + embedding: List[float] + semantic_inference: str = "" + + def __post_init__(self) -> None: + if not isinstance(self.task_id, UUID): + raise TypeError("task_id must be a UUID") + if not isinstance(self.topic, str) or not self.topic: + raise TypeError("topic must be a non-empty string") + self.time_span = _ensure_time_tuple(self.time_span) + if not isinstance(self.narrative_summary, str): + raise TypeError("narrative_summary must be a string") + if not isinstance(self.semantic_inference, str): + raise TypeError("semantic_inference must be a string") + if not isinstance(self.child_clip_ids, list): + raise TypeError("child_clip_ids must be a list of UUIDs") + self.child_clip_ids = [ + cid if isinstance(cid, UUID) else UUID(str(cid)) + for cid in self.child_clip_ids + ] + self.embedding = _ensure_float_list(self.embedding) + + def duration(self) -> float: + start, end = self.time_span + return float(end - start) + + def add_child_clip(self, clip_id: UUID) -> None: + if not isinstance(clip_id, UUID): + raise TypeError("clip_id must be a UUID") + self.child_clip_ids.append(clip_id) + + def to_dict(self) -> Dict[str, Any]: + return { + "task_id": str(self.task_id), + "topic": self.topic, + "time_span": [self.time_span[0], self.time_span[1]], + "narrative_summary": self.narrative_summary, + "semantic_inference": self.semantic_inference, + "child_clip_ids": [str(cid) for cid in self.child_clip_ids], + "embedding": list(self.embedding), + } + + @staticmethod + def from_dict(d: Dict[str, Any]) -> MediumTermMemory: + return MediumTermMemory( + task_id=UUID(d["task_id"]) + if not isinstance(d.get("task_id"), UUID) + else d["task_id"], + topic=d["topic"], + time_span=_ensure_time_tuple(tuple(d["time_span"])), + narrative_summary=d["narrative_summary"], + semantic_inference=str(d.get("semantic_inference", "")), + child_clip_ids=[ + UUID(cid) if not isinstance(cid, UUID) else cid + for cid in d["child_clip_ids"] + ], + embedding=_ensure_float_list(list(d["embedding"])), + ) + + +__all__ = [ + "ShortTermMemory", + "MediumTermMemory", +] diff --git a/jiuwen_memory/common/normalizer/normalizer_impl/video_normalizer.py b/jiuwen_memory/common/normalizer/normalizer_impl/video_normalizer.py new file mode 100644 index 00000000..dcbc62dd --- /dev/null +++ b/jiuwen_memory/common/normalizer/normalizer_impl/video_normalizer.py @@ -0,0 +1,271 @@ +"""Video normalizer: raw video reference -> structured video-memory data.""" + +from __future__ import annotations + +import importlib.util +import json +import os +import shutil +import tempfile +from collections.abc import Callable, Mapping +from pathlib import Path +from typing import Any +from urllib.parse import unquote, urlparse + +from jiuwen_memory.common.base import PluginType +from jiuwen_memory.common.errors import BackendError, HealthCheckError, ValidationError +from jiuwen_memory.common.log import get_logger +from jiuwen_memory.common.normalizer.base import Normalizer, NormalizerProducer +from jiuwen_memory.common.type_def import Modality, RawPayload + +VideoMemoryOutput = tuple[list[dict[str, Any]], list[dict[str, Any]]] +VideoMemoryBackend = Callable[[RawPayload], VideoMemoryOutput] + +logger = get_logger(__name__) + + +class VideoNormalizer(Normalizer): + """Run video perception and return structured video-memory data.""" + + def __init__( + self, + *, + chunk_seconds: int = 30, + whisper_model_dir: str = "", + whisper_device: str = "", + whisper_language: str = "", + whisper_batch_size: int = 1, + vllm_base_url: str = "", + vllm_api_key: str = "", + llm_model: str = "", + temp_root: str = "", + backend: VideoMemoryBackend | None = None, + ) -> None: + if chunk_seconds <= 0: + raise ValidationError("chunk_seconds must be greater than zero") + if whisper_batch_size <= 0: + raise ValidationError("whisper_batch_size must be greater than zero") + self._chunk_seconds = chunk_seconds + self._whisper_model_dir = whisper_model_dir + self._whisper_device = whisper_device + self._whisper_language = whisper_language + self._whisper_batch_size = whisper_batch_size + self._vllm_base_url = vllm_base_url + self._vllm_api_key = vllm_api_key + self._llm_model = llm_model + self._temp_root = temp_root + self._backend = backend + + @classmethod + def from_config( + cls, + config: Mapping[str, Any] | None, + *, + backend: VideoMemoryBackend | None = None, + ) -> VideoNormalizer: + params = dict(config or {}) + return cls( + chunk_seconds=int(params.get("chunk_seconds", 30)), + whisper_model_dir=str(params.get("whisper_model_dir", "")), + whisper_device=str(params.get("whisper_device", "")), + whisper_language=str(params.get("whisper_language", "")), + whisper_batch_size=int(params.get("whisper_batch_size", 1)), + vllm_base_url=str(params.get("vllm_base_url", "")), + vllm_api_key=str(params.get("vllm_api_key", "")), + llm_model=str(params.get("llm_model", "")), + temp_root=str(params.get("temp_root", "")), + backend=backend, + ) + + def modalities(self) -> list[Modality]: + return [Modality.VIDEO] + + def plugin_type(self) -> PluginType: + return PluginType.NORMALIZER + + def health(self) -> None: + if self._backend is not None: + return + missing_modules = [ + name + for name in ("torch", "transformers", "soundfile") + if importlib.util.find_spec(name) is None + ] + if missing_modules: + raise HealthCheckError( + "multimodal dependencies are missing: " + ", ".join(missing_modules) + ) + missing_binaries = [ + name for name in ("ffmpeg", "ffprobe") if shutil.which(name) is None + ] + if missing_binaries: + raise HealthCheckError( + "multimodal system dependencies are missing: " + + ", ".join(missing_binaries) + ) + + def normalize(self, payload: RawPayload) -> str: + if payload.modality != Modality.VIDEO: + raise ValidationError( + f"video normalizer does not support {payload.modality.value!r}" + ) + clips, events = self._extract_video_memory(payload) + video_memory = { + "payload_id": payload.id, + "asset_uri": payload.uri, + "clips": [_normalize_clip(item) for item in clips], + "events": [_normalize_event(item) for item in events], + } + return json.dumps(video_memory, ensure_ascii=False, separators=(",", ":")) + + def _extract_video_memory(self, payload: RawPayload) -> VideoMemoryOutput: + if self._backend is not None: + return self._backend(payload) + if not payload.uri: + raise ValidationError("video normalization requires RawPayload.uri") + + video_path = _file_uri_to_path(payload.uri) + if not video_path.is_file(): + raise ValidationError(f"video file not found: {video_path}") + + temp_root = Path(self._temp_root).expanduser() if self._temp_root else None + if temp_root is not None: + try: + temp_root.mkdir(parents=True, exist_ok=True) + except OSError as exc: + raise BackendError( + f"cannot create video temp_root {temp_root}: {exc}" + ) from exc + if not temp_root.is_dir(): + raise BackendError(f"video temp_root is not a directory: {temp_root}") + try: + temporary = tempfile.TemporaryDirectory( + prefix="agent-memory-video-", + dir=temp_root, + ) + except OSError as exc: + raise BackendError( + f"cannot create temporary video directory under {temp_root}: {exc}" + ) from exc + with temporary as temp_dir: + return self._run_pipeline(video_path, Path(temp_dir)) + + def _run_pipeline(self, video_path: Path, run_root: Path) -> VideoMemoryOutput: + try: + from .video_pipeline import VideoPipelineConfig, run_video_memory_pipeline_off + + outputs = run_video_memory_pipeline_off( + video_path, + run_root, + VideoPipelineConfig( + chunk_seconds=self._chunk_seconds, + whisper_device=self._whisper_device or None, + whisper_language=self._whisper_language or None, + whisper_batch_size=self._whisper_batch_size, + whisper_model_dir=self._whisper_model_dir or None, + require_precomputed_asr=False, + vllm_base_url=self._vllm_base_url or None, + vllm_api_key=self._vllm_api_key or None, + llm_model=self._llm_model or None, + resume_from_stream=False, + cleanup=True, + ), + ) + return ( + _object_list(outputs.get("short_term"), "short_term"), + _object_list(outputs.get("medium_term"), "medium_term"), + ) + except Exception as exc: + logger.exception( + "VideoNormalizer: pipeline failed for video=%s", video_path + ) + raise BackendError(f"embedded video normalization failed: {exc}") from exc + + +def _normalize_clip(item: dict[str, Any]) -> dict[str, Any]: + source_id = _required_string(item, "id", "clip") + start, end = _time_range(item, "time_range", "clip") + return { + "id": source_id, + "start_seconds": start, + "end_seconds": end, + "visual_summary": str(item.get("visual_summary", "")).strip(), + "detailed_caption": str(item.get("detailed_caption", "")).strip(), + "asr": str(item.get("ASR", "")).strip(), + "environment": str(item.get("environment", "")).strip(), + } + + +def _normalize_event(item: dict[str, Any]) -> dict[str, Any]: + source_id = _required_string(item, "task_id", "event") + start, end = _time_range(item, "time_span", "event") + children = item.get("child_clip_ids", []) + if not isinstance(children, list): + raise BackendError("video event child_clip_ids must be a list") + return { + "id": source_id, + "start_seconds": start, + "end_seconds": end, + "topic": str(item.get("topic", "")).strip(), + "summary": str(item.get("narrative_summary", "")).strip(), + "semantic_inference": str(item.get("semantic_inference", "")).strip(), + "clip_ids": [str(child) for child in children], + } + + +def _file_uri_to_path(uri: str) -> Path: + parsed = urlparse(uri) + if parsed.scheme not in ("", "file"): + raise ValidationError( + f"video normalizer does not support URI scheme {parsed.scheme!r}" + ) + if parsed.scheme == "": + return Path(uri).expanduser() + path = unquote(parsed.path) + if parsed.netloc and parsed.netloc not in ("", "localhost"): + path = f"//{parsed.netloc}{path}" + is_windows_drive = len(path) >= 3 and path[0] == "/" and path[2] == ":" + if os.name == "nt" and is_windows_drive: + path = path[1:] + return Path(path) + + +def _object_list(value: Any, label: str) -> list[dict[str, Any]]: + if not isinstance(value, list) or not all(isinstance(item, dict) for item in value): + raise BackendError(f"video memory {label} must be a list of objects") + return value + + +def _required_string(item: dict[str, Any], key: str, label: str) -> str: + value = str(item.get(key, "")).strip() + if not value: + raise BackendError(f"video {label} is missing {key}") + return value + + +def _time_range(item: dict[str, Any], key: str, label: str) -> tuple[float, float]: + value = item.get(key) + if not isinstance(value, (list, tuple)) or len(value) != 2: + raise BackendError(f"video {label} {key} must contain start and end") + try: + start, end = float(value[0]), float(value[1]) + except (TypeError, ValueError) as exc: + raise BackendError(f"video {label} {key} must be numeric") from exc + if end < start: + raise BackendError(f"video {label} {key} end must be >= start") + return start, end + + +@NormalizerProducer.register("video") +def _build(config): + return VideoNormalizer( + chunk_seconds=int(config.get("chunk_seconds", 30)), + whisper_model_dir=str(config.get("whisper_model_dir", "")), + whisper_device=str(config.get("whisper_device", "")), + whisper_language=str(config.get("whisper_language", "")), + whisper_batch_size=int(config.get("whisper_batch_size", 1)), + vllm_base_url=str(config.get("vllm_base_url", "")), + vllm_api_key=str(config.get("vllm_api_key", "")), + llm_model=str(config.get("llm_model", "")), + temp_root=str(config.get("temp_root", "")), + ) diff --git a/jiuwen_memory/common/normalizer/normalizer_impl/video_pipeline.py b/jiuwen_memory/common/normalizer/normalizer_impl/video_pipeline.py new file mode 100644 index 00000000..e2f48122 --- /dev/null +++ b/jiuwen_memory/common/normalizer/normalizer_impl/video_pipeline.py @@ -0,0 +1,2766 @@ +# ruff: noqa: E501 + +import atexit +import gc +import json +import os +import re +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple +from uuid import UUID, uuid4 + +import torch +import torch.distributed as dist +from openai import OpenAI + +from jiuwen_memory.common.log import get_logger + +from .video_asr import VideoAsrConfig, run_video_asr +from .video_models import MediumTermMemory, ShortTermMemory +from .video_prompts import ( + CAPTION_PROMPT, + EVENT_LINK_WITH_ET_PROMPT_AIO, + SESSION_SUMMARY_PROMPT_TEMPLATE, + UPDATE_EVENT_TABLE_PROMPT, +) +from .video_prompts import ( + CHAPTER_CHUNK as CHAPTER_SEGMENT_PROMPT_JSON, +) + +# Reduce CUDA fragmentation; must be set before CUDA init +os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") +# New unified allocator env var (PyTorch >= 2.4) +os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True") + +JSON_API_MAX_ATTEMPTS = 3 + +logger = get_logger(__name__) + + +@dataclass(frozen=True) +class ModelServiceConfig: + base_url: str + api_key: str + model: str + + +@dataclass(frozen=True) +class EventLinkRequest: + pre_event: Optional[Dict[str, Any]] + anchor: ShortTermMemory + pending: ShortTermMemory + segmentation_confidence: str + chapters: List[Dict[str, Any]] + candidate_source: str + + +@dataclass(frozen=True) +class SegmentPolicy: + target_len_s: float = 30.0 + min_len_s: float = 20.0 + max_len_s: float = 40.0 + snap_window_s: float = 10.0 + + +@dataclass(frozen=True) +class VideoPipelineConfig: + chunk_seconds: int = 30 + whisper_model_dir: Path | str | None = None + whisper_device: str | None = None + whisper_language: str | None = None + whisper_batch_size: int = 1 + require_precomputed_asr: bool = False + vllm_base_url: str | None = None + vllm_api_key: str | None = None + llm_model: str | None = None + resume_from_stream: bool = True + cleanup: bool = True + +UPDATE_EVENT_TABLE_PROMPT_FORCE_CONTINUE = ( + UPDATE_EVENT_TABLE_PROMPT + + """ + # FORCE CONTINUE MODE (HARD) + - You MUST output delta in CONTINUE format. + - You MUST NOT output SHIFT under any circumstance. + - You MAY slightly adjust event_title (event_identity) to appropriately encompass the current state of the ongoing event, while keeping it consistent with the core semantic of pre_ET.event_title. + - The adjusted event_title must remain local, stable, and not expand to a global/session-level scope; it should only refine the original title to fit cumulative information of the same unit. + """ +) + + +def _cleanup_dist() -> None: + """Best-effort destroy any torch distributed/NCCL process group at exit.""" + try: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + except RuntimeError as exc: + logger.warning("VideoPipeline: failed to destroy process group: %s", exc) + + +atexit.register(_cleanup_dist) + + +def _build_vllm_client( + base_url: str, + api_key: str, +) -> OpenAI: + """OpenAI-compatible client pointing to the local vLLM server.""" + return OpenAI(api_key=api_key, base_url=base_url) + + +def _parse_json_safe(text: str) -> Dict[str, Any]: + text = (text or "").strip() + if text.startswith("```"): + first_newline = text.find("\n") + if first_newline != -1: + text = text[first_newline + 1:] + if text.endswith("```"): + text = text[:-3].strip() + try: + return json.loads(text) + except (TypeError, json.JSONDecodeError) as exc: + start = text.find("{") + end = text.rfind("}") + if start != -1 and end != -1 and end > start: + candidate = text[start:end + 1] + try: + return json.loads(candidate) + except (TypeError, json.JSONDecodeError): + pass + raise ValueError("Model did not return valid JSON.") from exc + + +def _raw_snippet(text: Any, limit: int = 200) -> str: + snippet = str(text or "").replace("\n", " ").strip() + if len(snippet) > limit: + return snippet[:limit] + "..." + return snippet + + +def _call_openai_json_with_retry( + *, + model: str, + prompt: str, + max_tokens: int, + vllm_base_url: str, + vllm_api_key: str, + temperature: Optional[float] = None, + validator: Optional[Any] = None, + call_name: str = "json_call", +) -> Tuple[Optional[Dict[str, Any]], str, str]: + client = _build_vllm_client(vllm_base_url, vllm_api_key) + last_error = "" + last_raw = "" + + for attempt in range(1, JSON_API_MAX_ATTEMPTS + 1): + retry_prompt = prompt + if attempt > 1: + retry_prompt = ( + f"{prompt}\n\n" + "Your previous response could not be consumed by the pipeline. " + "Return valid JSON only, matching the requested schema exactly. " + "Do not include explanations, markdown fences, comments, or extra keys." + ) + + kwargs: Dict[str, Any] = { + "model": model, + "messages": [{"role": "user", "content": retry_prompt}], + "max_tokens": max_tokens, + } + if temperature is not None: + kwargs["temperature"] = temperature + + try: + resp = client.chat.completions.create(**kwargs) + raw = (resp.choices[0].message.content or "").strip() + last_raw = raw + data = _parse_json_safe(raw) + if validator is not None: + validator(data) + return data, raw, "" + except Exception as e: + last_error = str(e) + logger.warning( + "VideoPipeline: %s failed (attempt %d/%d): %s; raw=%s", + call_name, + attempt, + JSON_API_MAX_ATTEMPTS, + last_error, + _raw_snippet(last_raw or e), + ) + + return None, last_raw, last_error + + +def _validate_summary_payload(data: Dict[str, Any]) -> None: + if not isinstance(data, dict): + raise ValueError("summary payload is not a dict") + required = ("topic_label", "full_narrative", "semantic_inference") + missing = [key for key in required if key not in data] + if missing: + raise ValueError(f"summary payload missing keys: {missing}") + + +def _validate_event_link_payload(data: Dict[str, Any]) -> None: + if not isinstance(data, dict): + raise ValueError("event link payload is not a dict") + if "is_same_event" not in data: + raise ValueError("event link payload missing is_same_event") + if not isinstance(data.get("is_same_event"), bool): + raise ValueError("event link payload is_same_event must be bool") + split_raw = data.get("split_point") + if split_raw is None: + split_raw = data.get("split_points", []) + if split_raw is not None and not isinstance(split_raw, list): + raise ValueError("event link payload split_point(s) must be list") + + +def _validate_chapter_payload(data: Dict[str, Any]) -> None: + if not isinstance(data, dict): + raise ValueError("chapter payload is not a dict") + chapters = data.get("chapters") + if not isinstance(chapters, list) or not chapters: + raise ValueError("chapter payload missing non-empty chapters list") + seg_conf = str(data.get("segmentation_confidence", "")).strip().lower() + if seg_conf not in {"high", "medium", "low"}: + raise ValueError(f"invalid segmentation_confidence: {seg_conf}") + for idx, ch in enumerate(chapters, start=1): + if not isinstance(ch, dict): + raise ValueError(f"chapter {idx} is not a dict") + missing = [k for k in ("start_time", "title", "summary") if k not in ch] + if missing: + raise ValueError(f"chapter {idx} missing keys: {missing}") + + +def _normalize_chapters_start_only( + chapters: List[Dict[str, Any]], + video_duration_s: float, +) -> List[Dict[str, Any]]: + """Normalize and keep chapter fields in start-time-only format.""" + normalized: List[Dict[str, Any]] = [] + for idx, ch in enumerate(chapters, start=1): + if not isinstance(ch, dict): + continue + st_raw = str(ch.get("start_time", "00:00:00")) + st_s = _hhmmss_to_seconds(st_raw) + st_s = max(0.0, min(st_s, max(0.0, float(video_duration_s)))) + title = str(ch.get("title", "")).strip() or f"Chapter {idx}" + summary = str(ch.get("summary", "")).strip() + normalized.append( + { + "chapter_id": int(ch.get("chapter_id", idx) or idx), + "start_time": _seconds_to_hhmmss(st_s), + "title": title, + "summary": summary, + } + ) + + if not normalized: + return [ + { + "chapter_id": 1, + "start_time": _seconds_to_hhmmss(0.0), + "title": "Full Video", + "summary": "Auto-generated chapter due to missing ASR/segmentation", + } + ] + + normalized.sort( + key=lambda c: _hhmmss_to_seconds(str(c.get("start_time", "00:00:00"))) + ) + + deduped: List[Dict[str, Any]] = [] + for ch in normalized: + st = _hhmmss_to_seconds(str(ch.get("start_time", "00:00:00"))) + if deduped: + prev_st = _hhmmss_to_seconds(str(deduped[-1].get("start_time", "00:00:00"))) + if abs(st - prev_st) <= 1e-6: + # Keep the richer chapter content if starts are duplicated. + prev_summary = str(deduped[-1].get("summary", "")).strip() + curr_summary = str(ch.get("summary", "")).strip() + if len(curr_summary) > len(prev_summary): + deduped[-1] = ch + continue + deduped.append(ch) + + if _hhmmss_to_seconds(str(deduped[0].get("start_time", "00:00:00"))) > 0.0: + deduped[0]["start_time"] = _seconds_to_hhmmss(0.0) + + for i, ch in enumerate(deduped, start=1): + ch["chapter_id"] = i + + return deduped + + +def _validate_event_table_payload(data: Dict[str, Any]) -> None: + if not isinstance(data, dict): + raise ValueError("event table payload is not a dict") + required = { + "event_identity", + "event_summary", + "entities", + "open_questions", + "delta", + } + if not required.issubset(set(data.keys())): + missing = sorted(required - set(data.keys())) + raise ValueError(f"event table payload missing keys: {missing}") + identity = str(data.get("event_identity") or data.get("event_intent") or "").strip() + if not identity: + raise ValueError("event table payload missing event_identity") + if not isinstance(data.get("entities"), list): + raise ValueError("event table payload entities must be list") + if not isinstance(data.get("open_questions"), list): + raise ValueError("event table payload open_questions must be list") + + +def _fill_prompt_template(template: str, **kwargs: Any) -> str: + """Fill {name} placeholders without interpreting unrelated JSON braces.""" + out = str(template) + for key, value in kwargs.items(): + out = out.replace("{" + str(key) + "}", str(value)) + return out + + +def _resolve_video_path(video_path: Path) -> Path: + """Resolve migrated dataset paths by searching common roots and name variants.""" + if video_path.exists() and video_path.is_file(): + return video_path + + suffix = video_path.suffix or ".mp4" + stem = video_path.stem if video_path.suffix else video_path.name + base_name = video_path.name if video_path.suffix else f"{stem}{suffix}" + + stem_variants = [stem] + if stem.startswith("_"): + stem_variants.append(stem.lstrip("_")) + else: + stem_variants.append(f"_{stem}") + + name_variants: List[str] = [base_name] + for s in stem_variants: + n = f"{s}{suffix}" + if n not in name_variants: + name_variants.append(n) + + parent = video_path.parent + if parent.exists() and parent.is_dir(): + for name in name_variants: + cand = parent / name + if cand.exists() and cand.is_file(): + logger.info( + "VideoPipeline: resolved video path %s -> %s", video_path, cand + ) + return cand + + raise FileNotFoundError( + f"Input video not found: {video_path}. Checked common dataset roots and filename variants: {name_variants}" + ) + + +def _resolve_whisper_model_dir(whisper_model_dir: str | Path | None) -> str: + """Resolve the configured Whisper model directory.""" + raw = Path(str(whisper_model_dir)) if whisper_model_dir else None + if raw and raw.exists(): + return str(raw) + checked = [str(raw)] if raw else [] + raise FileNotFoundError( + f"Whisper model dir not found. Tried: {checked}. " + "Configure whisper_model_dir in the video normalizer." + ) + + +def _append_jsonl(path: Path, item: Dict[str, Any]) -> None: + def _drop_embeddings(obj: Any) -> Any: + if isinstance(obj, dict): + out: Dict[str, Any] = {} + for k, v in obj.items(): + if k == "embedding": + continue + out[k] = _drop_embeddings(v) + return out + if isinstance(obj, list): + return [_drop_embeddings(x) for x in obj] + return obj + + path.parent.mkdir(parents=True, exist_ok=True) + payload = _drop_embeddings(item) if path.name.endswith(".stream.jsonl") else item + with open(path, "a", encoding="utf-8") as f: + f.write(json.dumps(payload, ensure_ascii=True) + "\n") + + +# -------------------- Video helpers -------------------- + + +def _get_clip_duration(clip_path: Path) -> float: + import shutil + import subprocess + + if shutil.which("ffprobe") is None: + raise FileNotFoundError("ffprobe not found in PATH. Please install ffmpeg.") + cmd = [ + "ffprobe", + "-v", + "error", + "-show_entries", + "format=duration", + "-of", + "default=noprint_wrappers=1:nokey=1", + str(clip_path), + ] + res = subprocess.run( + cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True + ) + out = res.stdout.decode().strip() + try: + return float(out) + except (TypeError, ValueError) as e: + raise RuntimeError(f"Failed to parse duration for {clip_path}: {out}") from e + + +def _mmss_to_seconds(ts: str) -> float: + try: + parts = ts.split(":") + if len(parts) != 2: + return 0.0 + m, s = parts + return float(m) * 60 + float(s) + except (AttributeError, TypeError, ValueError): + return 0.0 + + +def _hhmmss_to_seconds(ts: str) -> float: + try: + parts = ts.split(":") + if len(parts) == 2: + return _mmss_to_seconds(ts) + if len(parts) != 3: + return 0.0 + h, m, s = parts + return float(h) * 3600 + float(m) * 60 + float(s) + except (AttributeError, TypeError, ValueError): + return 0.0 + + +def _seconds_to_hhmmss(sec: float) -> str: + sec = max(0.0, float(sec)) + h = int(sec // 3600) + m = int((sec % 3600) // 60) + s = int(sec % 60) + return f"{h:02d}:{m:02d}:{s:02d}" + + +def _normalize_to_hhmmss(ts: str) -> str: + return _seconds_to_hhmmss(_hhmmss_to_seconds(ts)) + + +# -------------------- Session / memory helpers -------------------- + + +def _summarize_session( + details: List[str], + *, + vllm_base_url: str, + vllm_api_key: str, + llm_model: str, +) -> Dict[str, str]: + detail_list = "\n".join(f"- {d}" for d in details) + prompt = _fill_prompt_template( + SESSION_SUMMARY_PROMPT_TEMPLATE, detail_list=detail_list + ) + data, raw, _ = _call_openai_json_with_retry( + model=llm_model, + prompt=prompt, + max_tokens=2048, + validator=_validate_summary_payload, + call_name="summarize_session", + vllm_base_url=vllm_base_url, + vllm_api_key=vllm_api_key, + ) + if data is None: + return { + "topic_label": "", + "full_narrative": raw, + "semantic_inference": "", + } + return { + "topic_label": str(data.get("topic_label", "")).strip(), + "full_narrative": str(data.get("full_narrative", "")).strip(), + "semantic_inference": str(data.get("semantic_inference", "")).strip(), + } + + +def _build_event_link_prompt(request: EventLinkRequest) -> str: + pre_et_payload = _event_link_pre_et_payload(request.pre_event) + pre_et_text = ( + "null" + if not pre_et_payload + else json.dumps(pre_et_payload, ensure_ascii=False, indent=2) + ) + seg_conf = str(request.segmentation_confidence or "").strip().lower() or "unknown" + + chapter_context_lines: List[str] = [f"segmentation_confidence: {seg_conf}"] + if request.chapters: + chapter_context_lines.append( + "chapter_list:\n" + + json.dumps(request.chapters, ensure_ascii=False, indent=2) + ) + else: + chapter_context_lines.append("chapter_list: []") + + chapter_context_text = "\n\n".join(chapter_context_lines) + + base = EVENT_LINK_WITH_ET_PROMPT_AIO + base = base.replace( + "{candidate_source}", + str(request.candidate_source or "").strip(), + ) + base = base.replace("{asr_confidence}", seg_conf) + base = base.replace("{chapter_context}", chapter_context_text) + base = base.replace("{pre_et}", pre_et_text) + base = base.replace("{anchor_summary}", request.anchor.visual_summary) + base = base.replace("{anchor_caption}", request.anchor.detailed_caption) + base = base.replace("{anchor_ASR}", request.anchor.asr) + base = base.replace("{pending_summary}", request.pending.visual_summary) + base = base.replace("{pending_caption}", request.pending.detailed_caption) + base = base.replace("{pending_ASR}", request.pending.asr) + return base + + +def _event_link_pre_et_payload( + pre_et: Optional[Dict[str, Any]], +) -> Optional[Dict[str, Any]]: + if not isinstance(pre_et, dict): + return None + return { + "event_identity": pre_et.get("event_identity") + or pre_et.get("event_intent", ""), + "event_summary": pre_et.get("event_summary", ""), + "entities": pre_et.get("entities") or [], + } + + +def _event_link_current_segments_payload( + anchor_stm: ShortTermMemory, + pending_stm: ShortTermMemory, +) -> Dict[str, Any]: + return { + "anchor_segment": { + "clip_id": str(anchor_stm.id), + "time_range": [anchor_stm.time_range[0], anchor_stm.time_range[1]], + "summary": anchor_stm.visual_summary, + "caption": anchor_stm.detailed_caption, + "asr": anchor_stm.asr, + }, + "pending_segment": { + "clip_id": str(pending_stm.id), + "time_range": [pending_stm.time_range[0], pending_stm.time_range[1]], + "summary": pending_stm.visual_summary, + "caption": pending_stm.detailed_caption, + "asr": pending_stm.asr, + }, + } + + +def _parse_split_points(data: Dict[str, Any]) -> List[Dict[str, str]]: + points: List[Dict[str, str]] = [] + if not isinstance(data, dict): + return points + raw = data.get("split_point") + if raw is None: + raw = data.get("split_points", []) + if not isinstance(raw, list): + return points + for item in raw: + if not isinstance(item, dict): + continue + t = str(item.get("t", "")).strip() + reason = str(item.get("reason", "")).strip() + if t: + points.append({"t": t, "reason": reason}) + return points + + +def _format_et_shift_candidate(delta: str, evidence: str) -> str: + raw = str(delta or "").strip() + if not raw: + return f"ET shift:SHIFT: unknown -> unknown | evidence={evidence}." + + upper = raw.upper() + if upper.startswith("SHIFT:"): + body = raw + else: + body = f"SHIFT: {raw}" + + marker = "| evidence=" + idx = body.lower().find(marker) + if idx != -1: + body = body[:idx].rstrip() + + if body.endswith("."): + body = body[:-1].rstrip() + + return f"ET shift:{body} | evidence={evidence}." + + +def _is_initialization_shift(delta: str) -> bool: + """Detect SHIFT deltas that are initialization-like (null/none/empty -> event).""" + raw = str(delta or "").strip() + if not raw: + return False + upper = raw.upper() + if not upper.startswith("SHIFT:"): + return False + body = raw.split(":", 1)[1].strip() if ":" in raw else "" + if "->" not in body: + return False + lhs = body.split("->", 1)[0].strip().strip(":").strip() + if not lhs: + return True + lhs_token = re.sub(r"[^A-Z0-9]+", "", lhs.upper()) + init_tokens = { + "NONE", + "NULL", + "NIL", + "NA", + "NAN", + "UNKNOWN", + "UNSET", + "EMPTY", + "INITIAL", + "INITIALSTATE", + "INIT", + "NOEVENT", + } + return lhs_token in init_tokens + + +def _judge_event_with_et( + request: EventLinkRequest, + *, + model_service: ModelServiceConfig, + log_path: Optional[Path] = None, + meta: Optional[Dict[str, Any]] = None, +) -> Tuple[bool, List[Dict[str, str]]]: + pre_et_payload = _event_link_pre_et_payload(request.pre_event) + current_segments_payload = _event_link_current_segments_payload( + request.anchor, + request.pending, + ) + prompt = _build_event_link_prompt(request) + data, content, err = _call_openai_json_with_retry( + model=model_service.model, + prompt=prompt, + max_tokens=256, + validator=_validate_event_link_payload, + call_name="judge_event_with_et", + vllm_base_url=model_service.base_url, + vllm_api_key=model_service.api_key, + ) + if data is None: + if log_path: + _append_jsonl( + log_path, + { + "ts": datetime.utcnow().isoformat() + "Z", + "meta": meta or {}, + "candidate_source": request.candidate_source, + "pre_ET": pre_et_payload, + "current_segments": current_segments_payload, + "raw_response": content, + "error": err or "judge_event_with_et failed after retries", + }, + ) + logger.warning( + "VideoPipeline: event-link inference exhausted retries; raw=%s", + _raw_snippet(content), + ) + return False, [] + try: + if log_path: + _append_jsonl( + log_path, + { + "ts": datetime.utcnow().isoformat() + "Z", + "meta": meta or {}, + "candidate_source": request.candidate_source, + "pre_ET": pre_et_payload, + "current_segments": current_segments_payload, + "raw_response": content, + "parsed": data, + }, + ) + return bool(data.get("is_same_event")), _parse_split_points(data) + except (AttributeError, TypeError, ValueError) as e: + if log_path: + _append_jsonl( + log_path, + { + "ts": datetime.utcnow().isoformat() + "Z", + "meta": meta or {}, + "candidate_source": request.candidate_source, + "pre_ET": pre_et_payload, + "current_segments": current_segments_payload, + "raw_response": content, + "error": str(e), + }, + ) + logger.warning( + "VideoPipeline: event-link response parsing failed: %s; raw=%s", + e, + _raw_snippet(content), + ) + return False, [] + + +def _stm_to_event_payload(stm: ShortTermMemory) -> Dict[str, Any]: + return { + "clip_id": str(stm.id), + "time_range": [stm.time_range[0], stm.time_range[1]], + "visual_summary": stm.visual_summary, + "detailed_caption": stm.detailed_caption, + "ASR": stm.asr, + "environment": stm.environment, + } + + +def _fallback_event_table( + pre_et: Optional[Dict[str, Any]], stm: ShortTermMemory +) -> Dict[str, Any]: + identity = (stm.visual_summary or "ongoing activity").strip() + summary = stm.detailed_caption.strip() or stm.visual_summary.strip() or "" + if pre_et and isinstance(pre_et, dict): + return { + "event_identity": pre_et.get("event_identity") + or pre_et.get("event_intent") + or identity, + "event_summary": pre_et.get("event_summary") or summary, + "entities": pre_et.get("entities") or [], + "open_questions": pre_et.get("open_questions") or [], + "delta": "No update due to parse failure.", + } + return { + "event_identity": identity, + "event_summary": summary, + "entities": [], + "open_questions": [], + "delta": "Initialize event from current clip.", + } + + +def _coerce_event_table( + data: Any, pre_et: Optional[Dict[str, Any]], stm: ShortTermMemory +) -> Dict[str, Any]: + if not isinstance(data, dict): + return _fallback_event_table(pre_et, stm) + required = {"event_summary", "entities", "open_questions", "delta"} + if not required.issubset(set(data.keys())): + return _fallback_event_table(pre_et, stm) + identity = str(data.get("event_identity") or data.get("event_intent") or "").strip() + if not identity: + return _fallback_event_table(pre_et, stm) + if not isinstance(data.get("entities"), list) or not isinstance( + data.get("open_questions"), list + ): + return _fallback_event_table(pre_et, stm) + return { + "event_identity": identity, + "event_summary": str(data.get("event_summary", "")).strip(), + "entities": data.get("entities") or [], + "open_questions": data.get("open_questions") or [], + "delta": str(data.get("delta", "")).strip(), + } + + +def _update_event_table( + pre_et: Optional[Dict[str, Any]], + stm: ShortTermMemory, + force_continue: bool = False, + *, + vllm_base_url: str, + vllm_api_key: str, + llm_model: str, +) -> Dict[str, Any]: + pre_et_text = ( + "null" if not pre_et else json.dumps(pre_et, ensure_ascii=False, indent=2) + ) + stm_text = json.dumps(_stm_to_event_payload(stm), ensure_ascii=False, indent=2) + base_prompt = ( + UPDATE_EVENT_TABLE_PROMPT_FORCE_CONTINUE + if force_continue + else UPDATE_EVENT_TABLE_PROMPT + ) + prompt = f"{base_prompt}\n\npre_ET:\n{pre_et_text}\n\nSTM_k:\n{stm_text}" + data, content, err = _call_openai_json_with_retry( + model=llm_model, + prompt=prompt, + max_tokens=1024, + validator=_validate_event_table_payload, + call_name="update_event_table", + vllm_base_url=vllm_base_url, + vllm_api_key=vllm_api_key, + ) + if data is None: + logger.warning( + "VideoPipeline: event-table inference exhausted retries: %s; raw=%s", + err, + _raw_snippet(content), + ) + return _fallback_event_table(pre_et, stm) + return _coerce_event_table(data, pre_et, stm) + + +# -------------------- Chapter segmentation -------------------- + + +def _segment_chapters( + cleaned_asr_segments: List[Dict[str, str]], + video_duration_s: float, + *, + vllm_base_url: str, + vllm_api_key: str, + llm_model: str, +) -> Dict[str, Any]: + chapters_only_schema = ( + "You must return JSON only in the following schema:\n" + "{\n" + ' "segmentation_confidence": string, // one of: high, medium, low\n' + ' "chapters": [\n' + " {\n" + ' "chapter_id": number,\n' + ' "start_time": "HH:MM:SS",\n' + ' "title": string,\n' + ' "summary": string\n' + " }\n" + " ]\n" + "}\n" + ) + transcript_lines = [] + for seg in cleaned_asr_segments: + st = _normalize_to_hhmmss(seg.get("start", "00:00")) + text = str(seg.get("text", "")).strip() + transcript_lines.append(f"[{st}] {text}") + + if not transcript_lines: + raise RuntimeError( + "ASR chaptering aborted: no ASR transcript lines available for chapter segmentation." + ) + + prompt = f"{CHAPTER_SEGMENT_PROMPT_JSON}\n\n{chapters_only_schema}\n\n" + "\n".join( + transcript_lines + ) + data, _, err = _call_openai_json_with_retry( + model=llm_model, + prompt=prompt, + max_tokens=8000, + temperature=0, + validator=_validate_chapter_payload, + call_name="segment_chapters", + vllm_base_url=vllm_base_url, + vllm_api_key=vllm_api_key, + ) + if data is not None: + seg_conf = ( + str(data.get("segmentation_confidence", "")).strip().lower() or "unknown" + ) + chapters = _normalize_chapters_start_only( + data.get("chapters") or [], video_duration_s + ) + return {"chapters": chapters, "segmentation_confidence": seg_conf} + raise RuntimeError( + "ASR chaptering failed after " + f"{JSON_API_MAX_ATTEMPTS} attempts; aborting this task. " + f"error={err}; asr_segments={len(cleaned_asr_segments)}" + ) + + +def _should_fallback_to_legacy( + asr_segments: List[Dict[str, str]], + cleaned_asr_segments: List[Dict[str, str]], + chapters: List[Dict[str, Any]], + video_duration_s: float, + segmentation_confidence: Optional[str] = None, +) -> Tuple[bool, str]: + try: + if segmentation_confidence and segmentation_confidence.lower() in ("low"): + return ( + True, + f"chapter segmentation_confidence {segmentation_confidence.lower()}", + ) + + text_len = sum(len(str(seg.get("text", ""))) for seg in cleaned_asr_segments) + if not asr_segments or not cleaned_asr_segments or text_len < 200: + return True, "ASR empty or too short" + + if not chapters: + return True, "chaptering returned empty" + + durations = [] + titles_norm: List[str] = [] + for idx, ch in enumerate(chapters): + st = _hhmmss_to_seconds(str(ch.get("start_time", "00:00:00"))) + if idx + 1 < len(chapters): + ed = _hhmmss_to_seconds( + str(chapters[idx + 1].get("start_time", "00:00:00")) + ) + else: + ed = video_duration_s + st = max(0.0, min(st, video_duration_s)) + ed = max(st, min(ed, video_duration_s)) + durations.append(ed - st) + title = str(ch.get("title", "")).strip().lower() + titles_norm.append(" ".join(title.split())) + + if len(chapters) == 1 and video_duration_s >= 600: + return True, "single chapter for long video" + + # 注释掉章节过长的回退判断 + # if durations and max(durations) > 900: + # return True, "chapter too long" + + if titles_norm: + from collections import Counter + + ctr = Counter(titles_norm) + most_common = ctr.most_common(1)[0][1] + if most_common / max(len(titles_norm), 1) >= 0.70: + return True, "chapter titles repetitive" + + except (AttributeError, IndexError, TypeError, ValueError) as e: + return True, f"chapter validation error: {e}" + + return False, "ok" + + +def _collect_asr_boundaries(cleaned_asr_segments: List[Dict[str, str]]) -> List[float]: + boundaries: List[float] = [] + for seg in cleaned_asr_segments: + st = _hhmmss_to_seconds(seg.get("start", "00:00:00")) + ed = _hhmmss_to_seconds(seg.get("end", "00:00:00")) + if st >= 0: + boundaries.append(st) + if ed >= 0: + boundaries.append(ed) + return sorted(set(boundaries)) + + +def _collect_chapter_boundaries(chapters: List[Dict[str, Any]]) -> List[float]: + boundaries: List[float] = [] + for idx, ch in enumerate(chapters): + if idx == 0: + continue + st = _hhmmss_to_seconds(str(ch.get("start_time", "00:00:00"))) + if st >= 0: + boundaries.append(st) + return sorted(set(boundaries)) + + +def _chapter_time_span( + chapters: List[Dict[str, Any]], + index: int, + video_duration_s: float, +) -> Tuple[float, float]: + """Derive chapter [start, end) from start-only chapter list.""" + start_s = _hhmmss_to_seconds(str(chapters[index].get("start_time", "00:00:00"))) + if index + 1 < len(chapters): + end_s = _hhmmss_to_seconds( + str(chapters[index + 1].get("start_time", "00:00:00")) + ) + else: + end_s = max(0.0, float(video_duration_s)) + start_s = max(0.0, min(start_s, max(0.0, float(video_duration_s)))) + end_s = max(start_s, min(end_s, max(0.0, float(video_duration_s)))) + return start_s, end_s + + +def make_asr_aligned_segments( + start_s: float, + end_s: float, + boundaries: List[float], + policy: SegmentPolicy = SegmentPolicy(), +) -> List[Dict[str, float]]: + segments: List[Dict[str, float]] = [] + cur = start_s + while cur < end_s - 1e-6: + remaining = end_s - cur + if remaining <= 5.0: + if segments: + segments[-1]["end_s"] = end_s + else: + segments.append({"segment_id": 1, "start_s": start_s, "end_s": end_s}) + break + + proposed = min(end_s, cur + policy.target_len_s) + candidates: List[Tuple[float, float]] = [] # (distance, boundary) + for b in boundaries: + if b <= cur or b >= end_s: + continue + seg_len = b - cur + if seg_len < policy.min_len_s or seg_len > policy.max_len_s: + continue + if abs(b - proposed) <= policy.snap_window_s: + candidates.append((abs(b - proposed), b)) + + chosen = proposed + if candidates: + candidates.sort(key=lambda x: x[0]) + chosen = candidates[0][1] + + seg_end = min(end_s, chosen) + if seg_end - cur < policy.min_len_s and remaining > policy.min_len_s: + seg_end = min(end_s, cur + policy.min_len_s) + if seg_end - cur > policy.max_len_s: + seg_end = min(end_s, cur + policy.max_len_s) + + if seg_end - cur < 5.0 and segments: + segments[-1]["end_s"] = end_s + break + + segments.append( + {"segment_id": len(segments) + 1, "start_s": cur, "end_s": seg_end} + ) + cur = seg_end + + return segments + + +def _get_video_duration(path: Path) -> float: + return _get_clip_duration(path) + + +def _extract_video_subclip( + video_path: Path, + out_path: Path, + start_s: float, + end_s: float, + timeout_s: int = 180, +) -> Path: + import shutil + import subprocess + + out_path.parent.mkdir(parents=True, exist_ok=True) + if shutil.which("ffmpeg") is None: + raise FileNotFoundError("ffmpeg not found in PATH. Please install ffmpeg.") + start_s = max(0.0, float(start_s)) + end_s = max(start_s, float(end_s)) + cmd = [ + "ffmpeg", + "-y", + "-nostdin", + "-hide_banner", + "-loglevel", + "error", + "-ss", + f"{start_s:.3f}", + "-to", + f"{end_s:.3f}", + "-i", + str(video_path), + "-c:v", + "libx264", + "-preset", + "fast", + "-crf", + "23", + "-c:a", + "aac", + "-b:a", + "128k", + str(out_path), + ] + try: + subprocess.run( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=True, + timeout=int(timeout_s), + ) + except subprocess.TimeoutExpired as e: + raise RuntimeError( + f"ffmpeg subclip timeout after {timeout_s}s: [{start_s:.3f}, {end_s:.3f}] {video_path}" + ) from e + return out_path + + +def _asr_for_range( + asr_segments: List[Dict[str, str]], start_s: float, end_s: float +) -> List[str]: + lines: List[str] = [] + for seg in asr_segments: + seg_start = _hhmmss_to_seconds(seg.get("start", "00:00:00")) + seg_end = _hhmmss_to_seconds(seg.get("end", "00:00:00")) + # Use half-open interval [start_s, end_s) to avoid boundary duplication + if seg_end <= start_s or seg_start >= end_s: + continue + text = str(seg.get("text", "")).strip() + if not text: + continue + st_norm = _normalize_to_hhmmss(seg.get("start", "00:00:00")) + ed_norm = _normalize_to_hhmmss(seg.get("end", "00:00:00")) + lines.append(f"[{st_norm}-{ed_norm}]:{text}") + return lines + + +def _merge_short_segments( + segments: List[Dict[str, Any]], + min_len_s: float, + chapter_start_s: float, + chapter_end_s: float, +) -> List[Dict[str, Any]]: + merged: List[Dict[str, Any]] = [] + for seg in segments: + st_s = max(chapter_start_s, float(seg.get("start_s", 0.0))) + ed_s = min(chapter_end_s, float(seg.get("end_s", st_s))) + if ed_s < st_s: + st_s, ed_s = ed_s, ed_s + dur = ed_s - st_s + if dur < min_len_s and merged: + merged[-1]["end_s"] = max(merged[-1]["end_s"], ed_s) + else: + merged.append( + {"segment_id": seg.get("segment_id"), "start_s": st_s, "end_s": ed_s} + ) + return merged + + +# -------------------- Offline Qwen-VL inference -------------------- + + +def _offline_vu( + clip_path: Path, + prompt: str, + vllm_client: OpenAI, + llm_model: str, +) -> Dict[str, Any]: + """Call vLLM server for vision-language inference on a video clip.""" + max_retries = 4 + generated_text = "" + for attempt in range(max_retries + 1): + try: + retry_prompt = ( + prompt + if attempt == 0 + else ( + f"{prompt}\n\nPlease strictly output valid JSON only. Do not include any extra text." + ) + ) + resp = vllm_client.chat.completions.create( + model=llm_model, + messages=[ + { + "role": "user", + "content": [ + { + "type": "video_url", + "video_url": {"url": f"file://{clip_path.resolve()}"}, + }, + {"type": "text", "text": retry_prompt}, + ], + } + ], + max_tokens=512, + temperature=0.4, + ) + generated_text = (resp.choices[0].message.content or "").strip() + return _parse_json_safe(generated_text) + except Exception as e: + snippet = (generated_text or str(e)).replace("\n", " ") + logger.warning( + "VideoPipeline: VLM inference failed (attempt %d/%d): %s", + attempt + 1, + max_retries + 1, + _raw_snippet(snippet), + ) + if attempt < max_retries: + continue + raise RuntimeError("VLM inference failed after retries") from e + + +# -------------------- Main pipeline -------------------- + + +def run_video_memory_pipeline_off( + video_path: str | Path, + work_root: str | Path, + config: VideoPipelineConfig, +) -> Dict[str, Any]: + video_path = _resolve_video_path(Path(video_path)) + chunk_seconds = config.chunk_seconds + whisper_model_dir = config.whisper_model_dir + whisper_device = config.whisper_device + whisper_language = config.whisper_language + whisper_batch_size = config.whisper_batch_size + require_precomputed_asr = config.require_precomputed_asr + vllm_base_url = config.vllm_base_url + vllm_api_key = config.vllm_api_key + llm_model = config.llm_model + resume_from_stream = config.resume_from_stream + cleanup = config.cleanup + if whisper_device is None: + whisper_device = "auto" + whisper_model_dir = _resolve_whisper_model_dir(whisper_model_dir) + if not str(vllm_base_url or "").strip(): + raise ValueError("configure vllm_base_url in the video normalizer") + if not str(vllm_api_key or "").strip(): + raise ValueError("configure vllm_api_key in the video normalizer") + if not str(llm_model or "").strip(): + raise ValueError("configure llm_model in the video normalizer") + logger.info( + "VideoPipeline: starting video=%s work_root=%s chunk_seconds=%d", + video_path, + work_root, + chunk_seconds, + ) + work_root = Path(work_root) + segments_dir = work_root / "segments" + outputs_dir = work_root / "outputs" + segments_dir.mkdir(parents=True, exist_ok=True) + outputs_dir.mkdir(parents=True, exist_ok=True) + event_link_log_path = outputs_dir / "event_link_logs.jsonl" + st_stream_path = outputs_dir / "short_term.stream.jsonl" + mt_stream_path = outputs_dir / "medium_term.stream.jsonl" + et_stream_path = outputs_dir / "event_table.stream.jsonl" + resume_active = bool(resume_from_stream) and any( + p.exists() and p.is_file() and p.stat().st_size > 0 + for p in (st_stream_path, mt_stream_path, et_stream_path) + ) + + if resume_active: + logger.info("VideoPipeline: resuming from existing stream state") + else: + with open(event_link_log_path, "w", encoding="utf-8") as f: + f.write("") + for p in (st_stream_path, mt_stream_path, et_stream_path): + with open(p, "w", encoding="utf-8") as f: + f.write("") + + # Ensure files exist for append mode even when partially resumed. + for p in (st_stream_path, mt_stream_path, et_stream_path, event_link_log_path): + if not p.exists(): + with open(p, "w", encoding="utf-8") as f: + f.write("") + + # ASR once via UASR, write only asr_segments.json + asr_save_path = outputs_dir / "asr_segments.json" + asr_tmp_work_dir = work_root / "asr_tmp" + asr_segments: List[Dict[str, str]] = [] + reused_asr = False + if asr_save_path.exists() and asr_save_path.is_file(): + try: + with open(asr_save_path, "r", encoding="utf-8") as f: + cached = json.load(f) + if isinstance(cached, list) and cached: + asr_segments = cached + reused_asr = True + logger.info( + "VideoPipeline: reusing ASR transcript path=%s segments=%d", + asr_save_path, + len(asr_segments), + ) + elif isinstance(cached, list) and not cached: + logger.warning( + "VideoPipeline: cached ASR transcript is empty; rerunning path=%s", + asr_save_path, + ) + else: + logger.warning( + "VideoPipeline: cached ASR transcript is invalid; rerunning path=%s", + asr_save_path, + ) + except (OSError, json.JSONDecodeError) as e: + logger.warning( + "VideoPipeline: failed to load cached ASR transcript; rerunning: %s", + e, + ) + + if require_precomputed_asr and not reused_asr: + raise RuntimeError( + f"ASR prerequisite missing or empty: {asr_save_path}. Run batch_ASR first and retry." + ) + + if not reused_asr: + asr_segments = run_video_asr( + video_path, + VideoAsrConfig( + model_dir=whisper_model_dir, + device=whisper_device, + language=whisper_language, + batch_size=max(1, int(whisper_batch_size)), + output_json=asr_save_path, + temp_work_dir=asr_tmp_work_dir, + cleanup=True, + ), + ) + logger.info( + "VideoPipeline: ASR completed path=%s segments=%d", + asr_save_path, + len(asr_segments), + ) + + video_duration_s = _get_video_duration(video_path) + + chapter_out_path = outputs_dir / "chapter_segmentation.json" + chapter_json: Dict[str, Any] = {} + reused_chapters = False + if chapter_out_path.exists() and chapter_out_path.is_file(): + try: + with open(chapter_out_path, "r", encoding="utf-8") as f: + cached_chapters = json.load(f) + if isinstance(cached_chapters, dict) and isinstance( + cached_chapters.get("chapters"), list + ): + chapter_json = cached_chapters + reused_chapters = True + logger.info( + "VideoPipeline: reusing chapter segmentation path=%s chapters=%d", + chapter_out_path, + len(cached_chapters.get("chapters", [])), + ) + else: + logger.warning( + "VideoPipeline: cached chapter segmentation is invalid; " + "rerunning path=%s", + chapter_out_path, + ) + except (OSError, json.JSONDecodeError) as e: + logger.warning( + "VideoPipeline: failed to load cached chapter segmentation; " + "rerunning: %s", + e, + ) + + if not reused_chapters: + chapter_json = _segment_chapters( + asr_segments, + video_duration_s, + vllm_base_url=vllm_base_url, + vllm_api_key=vllm_api_key, + llm_model=llm_model, + ) + + seg_conf = "unknown" + if isinstance(chapter_json, dict): + seg_conf = ( + str(chapter_json.get("segmentation_confidence", "")).strip().lower() + or "unknown" + ) + chapters = ( + chapter_json.get("chapters", []) if isinstance(chapter_json, dict) else [] + ) + chapters = _normalize_chapters_start_only(chapters, video_duration_s) + if chapters: + _, last_chapter_end_before = _chapter_time_span( + chapters, len(chapters) - 1, video_duration_s + ) + _, last_chapter_end_after = _chapter_time_span( + chapters, len(chapters) - 1, video_duration_s + ) + else: + last_chapter_end_before = 0.0 + last_chapter_end_after = 0.0 + + try: + with open(chapter_out_path, "w", encoding="utf-8") as f: + json.dump( + { + "segmentation_confidence": seg_conf, + "chapters": chapters, + }, + f, + ensure_ascii=False, + indent=2, + ) + logger.debug( + "VideoPipeline: saved chapter segmentation path=%s", chapter_out_path + ) + except OSError as save_err: + logger.warning( + "VideoPipeline: failed to save temporary chapter segmentation: %s", + save_err, + ) + + boundaries = _collect_asr_boundaries(asr_segments) + chapter_boundaries = _collect_chapter_boundaries(chapters) + last_asr_ts = max(boundaries) if boundaries else 0.0 + max_ch_dur = 0.0 + if chapters: + for i, _ in enumerate(chapters): + st, ed = _chapter_time_span(chapters, i, video_duration_s) + max_ch_dur = max(max_ch_dur, ed - st) + logger.debug( + "VideoPipeline: duration=%.3fs last_asr=%.3fs " + "last_chapter_before=%.3fs last_chapter_after=%.3fs max_chapter=%.3fs", + video_duration_s, + last_asr_ts, + last_chapter_end_before, + last_chapter_end_after, + max_ch_dur, + ) + logger.info( + "VideoPipeline: chapters=%d confidence=%s asr_text_length=%d", + len(chapters), + seg_conf, + sum(len(str(seg.get("text", ""))) for seg in asr_segments), + ) + fallback_needed, fallback_reason = _should_fallback_to_legacy( + asr_segments, asr_segments, chapters, video_duration_s, seg_conf + ) + logger.debug( + "VideoPipeline: fixed-segment fallback=%s reason=%s", + fallback_needed, + fallback_reason, + ) + if fallback_needed: + logger.info( + "VideoPipeline: using unified ASR-aligned fallback: %s", + fallback_reason, + ) + + vllm_client = _build_vllm_client(vllm_base_url, vllm_api_key) + logger.info("VideoPipeline: using VLM endpoint=%s", vllm_base_url) + + def _update_event_table_for_run( + pre_et: Optional[Dict[str, Any]], + stm: ShortTermMemory, + force_continue: bool = False, + ) -> Dict[str, Any]: + return _update_event_table( + pre_et, + stm, + force_continue=force_continue, + vllm_base_url=vllm_base_url, + vllm_api_key=vllm_api_key, + llm_model=llm_model, + ) + + short_terms: List[ShortTermMemory] = [] + stms_by_id: Dict[UUID, ShortTermMemory] = {} + medium_terms: List[MediumTermMemory] = [] + event_tables: List[Dict[str, Any]] = [] + + session_clip_ids: List[UUID] = [] + session_details: List[str] = [] + prev_stm: Optional[ShortTermMemory] = None + current_event_table: Optional[Dict[str, Any]] = None + current_event_id: Optional[UUID] = None + last_accepted_split_time_s = -1.0 + min_split_gap_s = float(chunk_seconds) + pending_active = False + pending_anchor_info: Optional[Dict[str, Any]] = None + pending_stms: List[ShortTermMemory] = [] + pending_event_table: Optional[Dict[str, Any]] = None + + def _iter_jsonl(path: Path) -> List[Dict[str, Any]]: + rows: List[Dict[str, Any]] = [] + if not path.exists() or not path.is_file(): + return rows + with open(path, "r", encoding="utf-8") as f: + for raw in f: + line = raw.strip() + if not line: + continue + try: + obj = json.loads(line) + except (TypeError, json.JSONDecodeError): + continue + if isinstance(obj, dict): + rows.append(obj) + return rows + + def _safe_uuid(raw: Any) -> Optional[UUID]: + try: + return raw if isinstance(raw, UUID) else UUID(str(raw)) + except (AttributeError, TypeError, ValueError): + return None + + def _stm_from_stream_row(d: Dict[str, Any]) -> Optional[ShortTermMemory]: + try: + sid = _safe_uuid(d.get("id")) + tr = d.get("time_range", [0.0, 0.0]) + if sid is None or not isinstance(tr, (list, tuple)) or len(tr) != 2: + return None + st = float(tr[0]) + ed = float(tr[1]) + if ed < st: + return None + visual_summary = str(d.get("visual_summary", "")) + detailed_caption = str(d.get("detailed_caption", "")) + asr_text = str(d.get("ASR", "")) + emb_raw = d.get("embedding", []) + emb = emb_raw if isinstance(emb_raw, list) else [] + emb_clean = [float(x) for x in emb if isinstance(x, (int, float))] + return ShortTermMemory( + id=sid, + video_source_path=str(d.get("video_source_path", video_path)), + time_range=(st, ed), + visual_summary=visual_summary, + detailed_caption=detailed_caption, + embedding=emb_clean, + asr=asr_text, + environment=str(d.get("environment", "")), + inferred_intent=str(d.get("inferred_intent", "")), + ) + except (AttributeError, TypeError, ValueError): + return None + + def _mtm_from_stream_row(d: Dict[str, Any]) -> Optional[MediumTermMemory]: + try: + tid = _safe_uuid(d.get("task_id")) + sp = d.get("time_span", [0.0, 0.0]) + if tid is None or not isinstance(sp, (list, tuple)) or len(sp) != 2: + return None + st = float(sp[0]) + ed = float(sp[1]) + if ed < st: + return None + child_raw = d.get("child_clip_ids", []) + child_ids: List[UUID] = [] + if isinstance(child_raw, list): + for x in child_raw: + uid = _safe_uuid(x) + if uid is not None: + child_ids.append(uid) + topic = str(d.get("topic", "") or "event") + narrative_summary = str(d.get("narrative_summary", "")) + semantic_inference = str(d.get("semantic_inference", "")) + emb_raw = d.get("embedding", []) + emb = emb_raw if isinstance(emb_raw, list) else [] + emb_clean = [float(x) for x in emb if isinstance(x, (int, float))] + return MediumTermMemory( + task_id=tid, + topic=topic, + time_span=(st, ed), + narrative_summary=narrative_summary, + semantic_inference=semantic_inference, + child_clip_ids=child_ids, + embedding=emb_clean, + ) + except (AttributeError, TypeError, ValueError): + return None + + def _record_stm(stm: ShortTermMemory) -> None: + short_terms.append(stm) + stms_by_id[stm.id] = stm + _append_jsonl(st_stream_path, stm.to_dict()) + + def _mtm_detail_from_stm(stm: ShortTermMemory) -> str: + visual = str(stm.detailed_caption or "").strip() + asr_text = str(stm.asr or "").strip() + environment = str(stm.environment or "").strip() + env_line = f"\n[ENV] {environment}" if environment else "" + if asr_text and asr_text != "无可用音频或转录失败": + return f"[VISUAL] {visual}{env_line}\n[ASR] {asr_text}".strip() + return f"[VISUAL] {visual}{env_line}".strip() + + def _finalize_session() -> None: + if not session_clip_ids: + return + first_stm = stms_by_id[session_clip_ids[0]] + last_stm = stms_by_id[session_clip_ids[-1]] + topic = first_stm.visual_summary or "事件" + time_span = (first_stm.time_range[0], last_stm.time_range[1]) + summary = _summarize_session( + session_details, + vllm_base_url=vllm_base_url, + vllm_api_key=vllm_api_key, + llm_model=llm_model, + ) + topic = summary.get("topic_label") or topic + narrative = summary.get("full_narrative") or "\n".join(session_details) + semantic_inference = summary.get("semantic_inference", "") + mtm = MediumTermMemory( + task_id=uuid4(), + topic=topic, + time_span=time_span, + narrative_summary=narrative, + semantic_inference=semantic_inference, + child_clip_ids=session_clip_ids, + embedding=[], + ) + medium_terms.append(mtm) + _append_jsonl(mt_stream_path, mtm.to_dict()) + + def _chapters_for_segment(start_s: float, end_s: float) -> List[Dict[str, Any]]: + if not chapters: + return [] + overlapping: List[Dict[str, Any]] = [] + for i, ch in enumerate(chapters): + st, ed = _chapter_time_span(chapters, i, video_duration_s) + if ed >= start_s and st <= end_s: + enriched = dict(ch) + enriched["start_time"] = _seconds_to_hhmmss(st) + enriched["end_time"] = _seconds_to_hhmmss(ed) + enriched["time_range"] = [ + _seconds_to_hhmmss(st), + _seconds_to_hhmmss(ed), + ] + overlapping.append(enriched) + return overlapping + + def _chapter_shift_text(boundary_t: float) -> str: + if not chapters: + return "ASR shift:unknown -> unknown | evidence=asr." + eps = 1e-3 + next_idx = -1 + for i in range(1, len(chapters)): + st = _hhmmss_to_seconds(str(chapters[i].get("start_time", "00:00:00"))) + if abs(st - boundary_t) <= 1.0 + eps: + next_idx = i + break + if next_idx <= 0: + for i in range(1, len(chapters)): + st = _hhmmss_to_seconds(str(chapters[i].get("start_time", "00:00:00"))) + if st >= boundary_t - eps: + next_idx = i + break + if next_idx <= 0: + return "ASR shift:unknown -> unknown | evidence=asr." + prev_ch = chapters[next_idx - 1] + next_ch = chapters[next_idx] + prev_summary = ( + str(prev_ch.get("summary") or prev_ch.get("title") or "unknown").strip() + or "unknown" + ) + next_summary = ( + str(next_ch.get("summary") or next_ch.get("title") or "unknown").strip() + or "unknown" + ) + return f"ASR shift:{prev_summary}->{next_summary} | evidence=asr." + + def _fixed_length_segments() -> List[Dict[str, float]]: + step = max(1.0, float(chunk_seconds)) + segments: List[Dict[str, float]] = [] + start_t = 0.0 + while start_t < video_duration_s - 1e-6: + end_t = min(video_duration_s, start_t + step) + segments.append({"start_s": start_t, "end_s": end_t}) + start_t = end_t + if not segments and video_duration_s > 0.0: + segments.append({"start_s": 0.0, "end_s": video_duration_s}) + for idx, seg in enumerate(segments, start=1): + seg["segment_id"] = idx + return segments + + def _prepare_segments() -> List[Dict[str, float]]: + asr_text_len = sum( + len(str(seg.get("text", "")).strip()) for seg in asr_segments + ) + if asr_text_len <= 0: + logger.info( + "VideoPipeline: no usable ASR text; using fixed-length segments" + ) + return _fixed_length_segments() + + segs = make_asr_aligned_segments( + 0.0, + video_duration_s, + boundaries, + SegmentPolicy( + target_len_s=float(chunk_seconds), + min_len_s=max(1.0, float(chunk_seconds) - 5.0), + max_len_s=float(chunk_seconds) + 5.0, + snap_window_s=10.0, + ), + ) + segs = _merge_short_segments( + segs, min_len_s=5.0, chapter_start_s=0.0, chapter_end_s=video_duration_s + ) + merged: List[Dict[str, float]] = [] + for seg in segs: + st = float(seg.get("start_s", 0.0)) + ed = float(seg.get("end_s", st)) + asr_lines = _asr_for_range(asr_segments, st, ed) + if not asr_lines and merged: + merged[-1]["end_s"] = max(merged[-1]["end_s"], ed) + else: + merged.append({"start_s": st, "end_s": ed}) + for idx, seg in enumerate(merged, start=1): + seg["segment_id"] = idx + return merged + + def _is_segment_already_covered(start_t: float, end_t: float) -> bool: + """Check whether existing STMs already cover this base segment interval.""" + if not short_terms: + return False + eps = 1e-3 + intervals: List[Tuple[float, float]] = [] + for stm in short_terms: + st = float(stm.time_range[0]) + ed = float(stm.time_range[1]) + if ed <= start_t + eps or st >= end_t - eps: + continue + cs = max(start_t, st) + ce = min(end_t, ed) + if ce - cs > eps: + intervals.append((cs, ce)) + + if not intervals: + return False + + intervals.sort(key=lambda x: x[0]) + merged: List[Tuple[float, float]] = [] + for st, ed in intervals: + if not merged or st > merged[-1][1] + 0.05: + merged.append((st, ed)) + else: + merged[-1] = (merged[-1][0], max(merged[-1][1], ed)) + + cursor = start_t + for st, ed in merged: + if st > cursor + 0.25: + return False + cursor = max(cursor, ed) + if cursor >= end_t - 0.25: + return True + return cursor >= end_t - 0.25 + + if resume_active: + restored_st = 0 + restored_mt = 0 + restored_et = 0 + + st_rows = _iter_jsonl(st_stream_path) + parsed_st: List[ShortTermMemory] = [] + seen_st_ids: set[UUID] = set() + for item in st_rows: + stm = _stm_from_stream_row(item) + if stm is None or stm.id in seen_st_ids: + continue + parsed_st.append(stm) + seen_st_ids.add(stm.id) + + # If a crashed/buggy rerun appended fresh clips from t~0 at tail, + # keep only the monotonic historical prefix and drop the rollback suffix. + rollback_idx = -1 + prev_end = -1.0 + for idx, stm in enumerate(parsed_st): + st_v = float(stm.time_range[0]) + ed_v = float(stm.time_range[1]) + if idx > 0 and st_v + 1e-3 < prev_end - 1.0: + rollback_idx = idx + break + prev_end = max(prev_end, ed_v) + + rollback_trimmed = rollback_idx >= 0 + if rollback_trimmed: + logger.warning( + "VideoPipeline: discarded rollback-appended clip tail from row=%d", + rollback_idx + 1, + ) + parsed_st = parsed_st[:rollback_idx] + + for stm in parsed_st: + short_terms.append(stm) + stms_by_id[stm.id] = stm + restored_st += 1 + + short_terms.sort( + key=lambda s: (float(s.time_range[0]), float(s.time_range[1]), str(s.id)) + ) + + st_resume_end_s = max( + (float(s.time_range[1]) for s in short_terms), default=-1.0 + ) + valid_st_ids = set(stms_by_id.keys()) + + seen_mtm_ids: set[UUID] = set() + for item in _iter_jsonl(mt_stream_path): + mtm = _mtm_from_stream_row(item) + if mtm is None or mtm.task_id in seen_mtm_ids: + continue + if st_resume_end_s >= 0 and float(mtm.time_span[1]) > st_resume_end_s + 1.0: + continue + if mtm.child_clip_ids and any( + cid not in valid_st_ids for cid in mtm.child_clip_ids + ): + continue + medium_terms.append(mtm) + seen_mtm_ids.add(mtm.task_id) + restored_mt += 1 + + for item in _iter_jsonl(et_stream_path): + cid = _safe_uuid(item.get("clip_id")) + if cid is None or cid not in valid_st_ids: + continue + tr = item.get("time_range", []) + if isinstance(tr, (list, tuple)) and len(tr) == 2 and st_resume_end_s >= 0: + try: + if float(tr[1]) > st_resume_end_s + 1.0: + continue + except (TypeError, ValueError): + continue + event_tables.append(item) + restored_et += 1 + + if rollback_trimmed: + # Rewrite sanitized streams so future resumes are deterministic. + with open(st_stream_path, "w", encoding="utf-8") as f: + for s in short_terms: + f.write(json.dumps(s.to_dict(), ensure_ascii=True) + "\n") + with open(mt_stream_path, "w", encoding="utf-8") as f: + for m in medium_terms: + f.write(json.dumps(m.to_dict(), ensure_ascii=True) + "\n") + with open(et_stream_path, "w", encoding="utf-8") as f: + for rec in event_tables: + f.write(json.dumps(rec, ensure_ascii=True) + "\n") + + # Rebuild active session state from last committed ET event chain. + if event_tables: + valid_chain: List[Tuple[Dict[str, Any], UUID]] = [] + for rec in event_tables: + cid = _safe_uuid(rec.get("clip_id")) + if cid is None or cid not in stms_by_id: + continue + valid_chain.append((rec, cid)) + + if valid_chain: + last_rec, _ = valid_chain[-1] + evt_raw = str(last_rec.get("event_id", "")).strip() + current_event_id = _safe_uuid(evt_raw) + committed = last_rec.get("event_table_committed") + if not isinstance(committed, dict): + committed = last_rec.get("event_table") + current_event_table = committed if isinstance(committed, dict) else None + + if current_event_id is not None: + for rec, cid in valid_chain: + if str(rec.get("event_id", "")).strip() == str( + current_event_id + ): + session_clip_ids.append(cid) + + session_details = [ + _mtm_detail_from_stm(stms_by_id[cid]) + for cid in session_clip_ids + if cid in stms_by_id + ] + prev_stm = ( + stms_by_id[session_clip_ids[-1]] if session_clip_ids else None + ) + + for rec in event_tables: + decision = rec.get("decision") + if not isinstance(decision, dict): + continue + split_t = decision.get("split_t") + if split_t is None: + continue + try: + split_v = float(split_t) + except (TypeError, ValueError): + continue + if split_v > last_accepted_split_time_s: + last_accepted_split_time_s = split_v + + logger.info( + "VideoPipeline: restored clips=%d events=%d event_tables=%d " + "active_session_clips=%d", + restored_st, + restored_mt, + restored_et, + len(session_clip_ids), + ) + + def _append_event_table_record( + *, + event_id: str, + clip_id: str, + time_range: List[float], + committed_et: Dict[str, Any], + proposed_et: Optional[Dict[str, Any]] = None, + decision: Optional[Dict[str, Any]] = None, + ) -> None: + """Persist both pre-judge proposal and committed ET for traceability.""" + rec = { + "event_id": event_id, + "clip_id": clip_id, + "time_range": time_range, + # Backward compatible field consumed by existing scripts. + "event_table": committed_et, + "event_table_proposed": proposed_et + if proposed_et is not None + else committed_et, + "event_table_committed": committed_et, + } + if decision is not None: + rec["decision"] = decision + event_tables.append(rec) + _append_jsonl(et_stream_path, rec) + + def _is_et_only_cooldown_active( + start_t: float, candidate_sources: List[str] + ) -> bool: + """Suppress ET-only re-splitting immediately after a confirmed cut.""" + if candidate_sources != ["ET"]: + return False + if last_accepted_split_time_s < 0: + return False + gap = start_t - last_accepted_split_time_s + return 0.0 <= gap < min_split_gap_s + + def _apply_split(split_t: float, start_t: float, end_t: float, seg_id: int) -> None: + nonlocal session_clip_ids + nonlocal session_details + nonlocal prev_stm + nonlocal current_event_table + nonlocal current_event_id + nonlocal last_accepted_split_time_s + + pre_stm, _, _ = _run_stm_for_segment(start_t, split_t, f"{seg_id}_a") + if pre_stm: + _record_stm(pre_stm) + session_clip_ids.append(pre_stm.id) + session_details.append(_mtm_detail_from_stm(pre_stm)) + prev_stm = pre_stm + pre_proposed = _update_event_table_for_run(current_event_table, pre_stm) + current_event_table = pre_proposed + _append_event_table_record( + event_id=str(current_event_id) if current_event_id else "", + clip_id=str(pre_stm.id), + time_range=[start_t, split_t], + proposed_et=pre_proposed, + committed_et=current_event_table, + decision={"mode": "split_pre", "split_t": split_t}, + ) + + if session_clip_ids: + _finalize_session() + session_clip_ids = [] + session_details = [] + prev_stm = None + current_event_table = None + current_event_id = None + last_accepted_split_time_s = split_t + + post_stm, _, _ = _run_stm_for_segment(split_t, end_t, f"{seg_id}_b") + if post_stm: + _record_stm(post_stm) + session_clip_ids = [post_stm.id] + session_details = [_mtm_detail_from_stm(post_stm)] + prev_stm = post_stm + current_event_id = uuid4() + post_proposed = _update_event_table_for_run(None, post_stm) + current_event_table = post_proposed + _append_event_table_record( + event_id=str(current_event_id), + clip_id=str(post_stm.id), + time_range=[split_t, end_t], + proposed_et=post_proposed, + committed_et=current_event_table, + decision={"mode": "split_post", "split_t": split_t}, + ) + + def _parse_split_seconds(raw_t: str) -> Optional[float]: + t = str(raw_t or "").strip() + if not t: + return None + try: + return float(t) + except (TypeError, ValueError): + pass + try: + return _hhmmss_to_seconds(t) + except (AttributeError, TypeError, ValueError): + return None + + def _pick_split_t( + split_points: List[Dict[str, str]], + start_t: float, + end_t: float, + fallback_t: Optional[float] = None, + ) -> Optional[float]: + eps = 1e-3 + for sp in split_points: + t_s = _parse_split_seconds(sp.get("t", "")) + if t_s is None: + continue + # Allow split at the left edge (new clip starts a new event), + # but keep right edge exclusive to avoid empty post-split segment. + if start_t - eps <= t_s < end_t - eps: + return t_s + if fallback_t is not None and start_t - eps <= fallback_t < end_t - eps: + return float(fallback_t) + return None + + def _run_stm_for_segment( + start_t: float, end_t: float, seg_tag: str + ) -> Tuple[Optional[ShortTermMemory], List[float], str]: + out_clip = segments_dir / f"seg_{seg_tag}.mp4" + logger.debug( + "VideoPipeline: extracting segment=%s start=%.3f end=%.3f", + seg_tag, + start_t, + end_t, + ) + try: + _extract_video_subclip(video_path, out_clip, start_t, end_t) + except Exception as e: + logger.exception( + "VideoPipeline: failed to extract segment=%s start=%.3f end=%.3f", + seg_tag, + start_t, + end_t, + ) + raise RuntimeError(f"failed to extract video segment {seg_tag}") from e + logger.debug("VideoPipeline: segment=%s ready path=%s", seg_tag, out_clip) + + asr_lines = _asr_for_range(asr_segments, start_t, end_t) + asr_text = "\n".join(asr_lines) if asr_lines else "无可用音频或转录失败" + pre_et_text = ( + "null" + if not current_event_table + else json.dumps(current_event_table, ensure_ascii=False, indent=2) + ) + prompt = CAPTION_PROMPT.replace("{event_table}", pre_et_text) + logger.debug("VideoPipeline: running VLM inference segment=%s", seg_tag) + cap = _offline_vu( + out_clip, + prompt, + vllm_client=vllm_client, + llm_model=llm_model, + ) + logger.debug("VideoPipeline: VLM inference completed segment=%s", seg_tag) + + visual_summary = str(cap.get("visual_summary", "")).strip() + detailed_caption = str(cap.get("detailed_caption", "")).strip() + environment = str(cap.get("environment", "")).strip() + asr_corrected = asr_text + + stm = ShortTermMemory( + id=uuid4(), + video_source_path=str(video_path), + time_range=(start_t, end_t), + visual_summary=visual_summary, + detailed_caption=detailed_caption, + embedding=[], + asr=asr_corrected, + environment=environment, + ) + return stm, [], asr_text + + def _rebuild_event_table_from_old( + old_et: Optional[Dict[str, Any]], + pending_items: List[ShortTermMemory], + ) -> Dict[str, Any]: + rebuilt = old_et + for item in pending_items: + rebuilt = _update_event_table_for_run( + rebuilt, + item, + force_continue=True, + ) + if isinstance(rebuilt, dict): + return rebuilt + if pending_items: + return _update_event_table_for_run( + old_et, + pending_items[-1], + force_continue=True, + ) + return old_et or {} + + try: + segments = _prepare_segments() + logger.info("VideoPipeline: processing segments=%d", len(segments)) + for seg in segments: + start_t = round(float(seg["start_s"]), 3) + end_t = round(float(seg["end_s"]), 3) + if resume_active and _is_segment_already_covered(start_t, end_t): + logger.debug( + "VideoPipeline: skipping restored segment=%d start=%.3f end=%.3f", + int(seg.get("segment_id", 0) or 0), + start_t, + end_t, + ) + continue + segment_chapter_boundaries = [ + b for b in chapter_boundaries if start_t <= b < end_t + ] + + stm_full, _, _ = _run_stm_for_segment( + start_t, end_t, f"{seg['segment_id']}" + ) + if stm_full is None: + continue + + if current_event_table is None: + if segment_chapter_boundaries: + boundary_time = min(segment_chapter_boundaries) + _append_jsonl( + event_link_log_path, + { + "ts": datetime.utcnow().isoformat() + "Z", + "meta": { + "segment_id": int(seg.get("segment_id", 0) or 0), + "time_range": [start_t, end_t], + "boundary_source": "chapter_segment", + "boundary_time": boundary_time, + "boundary_time_hhmmss": _seconds_to_hhmmss( + boundary_time + ), + "lookahead_clips": 1, + "init_pre_et": True, + "init_action": "direct_split", + }, + "candidate_source": _chapter_shift_text(boundary_time), + "pre_ET": None, + "current_segments": _event_link_current_segments_payload( + stm_full, stm_full + ), + "note": "Chapter boundary observed during initialization; split applied immediately.", + }, + ) + _apply_split( + boundary_time, + start_t, + end_t, + int(seg.get("segment_id", 0) or 0), + ) + continue + + _record_stm(stm_full) + session_clip_ids = [stm_full.id] + session_details = [_mtm_detail_from_stm(stm_full)] + prev_stm = stm_full + current_event_id = uuid4() + init_proposed = _update_event_table_for_run(None, stm_full) + current_event_table = init_proposed + _append_event_table_record( + event_id=str(current_event_id), + clip_id=str(stm_full.id), + time_range=[start_t, end_t], + proposed_et=init_proposed, + committed_et=current_event_table, + decision={"mode": "init_from_null_pre_et"}, + ) + continue + + if pending_active and pending_anchor_info is not None: + pending_stms.append(stm_full) + pending_event_table = _update_event_table_for_run( + pending_event_table, + stm_full, + ) + + left_event_table = pending_anchor_info.get("left_event_table") + left_event_id = pending_anchor_info.get("left_event_id") + left_session_ids = list( + pending_anchor_info.get("left_session_clip_ids") or [] + ) + left_session_detail_list = list( + pending_anchor_info.get("left_session_details") or [] + ) + + anchor_stm = pending_stms[0] + pending_confirm_stm = pending_stms[-1] + pending_start_t = float(pending_stms[0].time_range[0]) + pending_end_t = float(pending_stms[-1].time_range[1]) + pending_chapters = _chapters_for_segment(pending_start_t, pending_end_t) + + effective_candidate_sources = list( + pending_anchor_info.get("candidate_sources") or [] + ) + anchor_end_t = float( + pending_anchor_info.get("anchor_end_t", pending_start_t) + ) + pending_new_chapter_boundaries = [ + b for b in chapter_boundaries if anchor_end_t <= b < pending_end_t + ] + if ( + pending_new_chapter_boundaries + and "chapter_segment" not in effective_candidate_sources + ): + effective_candidate_sources.append("chapter_segment") + + has_asr_source = "chapter_segment" in effective_candidate_sources + has_et_source = "ET" in effective_candidate_sources + if has_asr_source and pending_new_chapter_boundaries: + pending_boundary_time = min(pending_new_chapter_boundaries) + else: + pending_boundary_time = float( + pending_anchor_info.get("boundary_time", pending_start_t) + ) + pending_source_tag = ( + "ET+chapter_segment" + if len(effective_candidate_sources) > 1 + else ( + effective_candidate_sources[0] + if effective_candidate_sources + else "unknown" + ) + ) + parts: List[str] = [] + if has_asr_source: + parts.append(_chapter_shift_text(pending_boundary_time)) + if has_et_source: + et_evidence = "both" if has_asr_source else "et" + parts.append( + _format_et_shift_candidate( + str(pending_anchor_info.get("trigger_delta", "")), + et_evidence, + ) + ) + pending_candidate_source_text = ( + " | ".join(parts) + if parts + else str( + pending_anchor_info.get("candidate_source_text", "unknown") + ) + ) + pending_anchor_info["candidate_sources"] = effective_candidate_sources + pending_anchor_info["candidate_source_text"] = ( + pending_candidate_source_text + ) + pending_anchor_info["boundary_source"] = pending_source_tag + pending_anchor_info["boundary_time"] = pending_boundary_time + + same_event, split_points = _judge_event_with_et( + EventLinkRequest( + pre_event=left_event_table, + anchor=anchor_stm, + pending=pending_confirm_stm, + segmentation_confidence=seg_conf, + chapters=pending_chapters, + candidate_source=pending_candidate_source_text, + ), + model_service=ModelServiceConfig( + base_url=vllm_base_url, + api_key=vllm_api_key, + model=llm_model, + ), + log_path=event_link_log_path, + meta={ + "segment_id": int(seg.get("segment_id", 0) or 0), + "time_range": [pending_start_t, pending_end_t], + "boundary_source": "pending_delayed_gate", + "pending_candidate_source": pending_source_tag, + "boundary_time": pending_boundary_time, + "boundary_time_hhmmss": _seconds_to_hhmmss( + pending_boundary_time + ), + "anchor_segment_id": pending_anchor_info.get( + "anchor_segment_id" + ), + "anchor_time_range": [ + pending_anchor_info.get("anchor_start_t", pending_start_t), + pending_anchor_info.get("anchor_end_t", pending_end_t), + ], + "trigger_delta": pending_anchor_info.get("trigger_delta", ""), + }, + ) + logger.debug( + "VideoPipeline: boundary confirmation anchor=%s segment=%s " + "same_event=%s", + pending_anchor_info.get("anchor_segment_id"), + seg.get("segment_id"), + same_event, + ) + + if not same_event: + split_t = _pick_split_t( + split_points, + pending_start_t, + pending_end_t, + fallback_t=float( + pending_anchor_info.get("boundary_time", pending_start_t) + ), + ) + split_idx = -1 + split_boundary_idx = -1 + split_seg_start = 0.0 + split_seg_end = 0.0 + if split_t is not None: + eps = 1e-3 + for i, pstm in enumerate(pending_stms): + st_i = float(pstm.time_range[0]) + ed_i = float(pstm.time_range[1]) + if st_i + eps < split_t < ed_i - eps: + split_idx = i + split_seg_start = st_i + split_seg_end = ed_i + break + # split point lands exactly on clip boundary: split by clip index + if abs(ed_i - split_t) <= eps: + split_boundary_idx = i + 1 + break + if abs(st_i - split_t) <= eps: + split_boundary_idx = i + break + + session_clip_ids = left_session_ids + session_details = left_session_detail_list + prev_stm = ( + stms_by_id[session_clip_ids[-1]] + if session_clip_ids and session_clip_ids[-1] in stms_by_id + else None + ) + current_event_table = ( + left_event_table + if isinstance(left_event_table, dict) + else current_event_table + ) + current_event_id = ( + left_event_id + if isinstance(left_event_id, UUID) + else current_event_id + ) + + if split_idx >= 0 and split_t is not None: + for pstm in pending_stms[:split_idx]: + _record_stm(pstm) + session_clip_ids.append(pstm.id) + session_details.append(_mtm_detail_from_stm(pstm)) + prev_stm = pstm + current_event_table = _update_event_table_for_run( + current_event_table, pstm, force_continue=True + ) + _append_event_table_record( + event_id=str(current_event_id) + if current_event_id + else "", + clip_id=str(pstm.id), + time_range=[pstm.time_range[0], pstm.time_range[1]], + proposed_et=pending_event_table, + committed_et=current_event_table, + decision={ + "mode": "pending_accept_pre_split", + "split_t": split_t, + "pending_index": split_idx, + "candidate_sources": pending_anchor_info.get( + "candidate_sources" + ) + or [], + }, + ) + + _apply_split( + split_t, + split_seg_start, + split_seg_end, + int(seg.get("segment_id", 0) or 0), + ) + + for pstm in pending_stms[split_idx + 1:]: + _record_stm(pstm) + session_clip_ids.append(pstm.id) + session_details.append(_mtm_detail_from_stm(pstm)) + prev_stm = pstm + proposed_after = _update_event_table_for_run( + current_event_table, pstm, force_continue=True + ) + current_event_table = proposed_after + _append_event_table_record( + event_id=str(current_event_id) + if current_event_id + else "", + clip_id=str(pstm.id), + time_range=[pstm.time_range[0], pstm.time_range[1]], + proposed_et=proposed_after, + committed_et=current_event_table, + decision={ + "mode": "split_post_attach", + "split_t": split_t, + "candidate_sources": pending_anchor_info.get( + "candidate_sources" + ) + or [], + }, + ) + elif split_boundary_idx >= 0 and split_t is not None: + # Boundary split: no need to re-run clip inference; split by existing clip boundaries. + for pstm in pending_stms[:split_boundary_idx]: + _record_stm(pstm) + session_clip_ids.append(pstm.id) + session_details.append(_mtm_detail_from_stm(pstm)) + prev_stm = pstm + current_event_table = _update_event_table_for_run( + current_event_table, pstm, force_continue=True + ) + _append_event_table_record( + event_id=str(current_event_id) + if current_event_id + else "", + clip_id=str(pstm.id), + time_range=[pstm.time_range[0], pstm.time_range[1]], + proposed_et=pending_event_table, + committed_et=current_event_table, + decision={ + "mode": "pending_accept_pre_boundary", + "split_t": split_t, + "split_boundary_idx": split_boundary_idx, + "candidate_sources": pending_anchor_info.get( + "candidate_sources" + ) + or [], + }, + ) + + if session_clip_ids: + _finalize_session() + + current_event_id = uuid4() + current_event_table = None + session_clip_ids = [] + session_details = [] + prev_stm = None + + for pstm in pending_stms[split_boundary_idx:]: + _record_stm(pstm) + session_clip_ids.append(pstm.id) + session_details.append(_mtm_detail_from_stm(pstm)) + prev_stm = pstm + proposed_after = _update_event_table_for_run( + current_event_table, pstm + ) + current_event_table = proposed_after + _append_event_table_record( + event_id=str(current_event_id), + clip_id=str(pstm.id), + time_range=[pstm.time_range[0], pstm.time_range[1]], + proposed_et=proposed_after, + committed_et=current_event_table, + decision={ + "mode": "pending_accept_post_boundary", + "split_t": split_t, + "split_boundary_idx": split_boundary_idx, + "candidate_sources": pending_anchor_info.get( + "candidate_sources" + ) + or [], + }, + ) + else: + if session_clip_ids: + _finalize_session() + + current_event_id = uuid4() + current_event_table = ( + pending_event_table + if isinstance(pending_event_table, dict) + else _update_event_table_for_run(None, pending_stms[0]) + ) + session_clip_ids = [] + session_details = [] + for idx, pstm in enumerate(pending_stms, start=1): + _record_stm(pstm) + session_clip_ids.append(pstm.id) + session_details.append(_mtm_detail_from_stm(pstm)) + _append_event_table_record( + event_id=str(current_event_id), + clip_id=str(pstm.id), + time_range=[pstm.time_range[0], pstm.time_range[1]], + proposed_et=current_event_table, + committed_et=current_event_table, + decision={ + "mode": "pending_accept_new_event", + "pending_index": idx, + "pending_size": len(pending_stms), + "candidate_sources": pending_anchor_info.get( + "candidate_sources" + ) + or [], + }, + ) + prev_stm = pending_stms[-1] if pending_stms else None + last_accepted_split_time_s = float( + split_t + if split_t is not None + else pending_anchor_info.get("anchor_start_t", pending_start_t) + ) + else: + running_old_et = left_event_table + for idx, pstm in enumerate(pending_stms, start=1): + _record_stm(pstm) + running_old_et = _update_event_table_for_run( + running_old_et, pstm, force_continue=True + ) + _append_event_table_record( + event_id=str(left_event_id) if left_event_id else "", + clip_id=str(pstm.id), + time_range=[pstm.time_range[0], pstm.time_range[1]], + proposed_et=pending_event_table, + committed_et=running_old_et, + decision={ + "mode": "pending_reject_force_continue", + "pending_index": idx, + "pending_size": len(pending_stms), + "candidate_sources": pending_anchor_info.get( + "candidate_sources" + ) + or [], + }, + ) + current_event_table = running_old_et + current_event_id = ( + left_event_id + if isinstance(left_event_id, UUID) + else current_event_id + ) + session_clip_ids = left_session_ids + [s.id for s in pending_stms] + session_details = left_session_detail_list + [ + _mtm_detail_from_stm(s) for s in pending_stms + ] + prev_stm = pending_stms[-1] + + pending_active = False + pending_anchor_info = None + pending_stms = [] + pending_event_table = None + continue + + proposed_event_table = _update_event_table_for_run( + current_event_table, + stm_full, + ) + delta = str(proposed_event_table.get("delta", "")).strip() + delta_upper = delta.upper() + is_candidate = ( + current_event_table is not None + and delta_upper.startswith("SHIFT:") + and "NONE ->" not in delta_upper + and not _is_initialization_shift(delta) + ) + candidate_sources: List[str] = [] + if is_candidate: + candidate_sources.append("ET") + if segment_chapter_boundaries: + candidate_sources.append("chapter_segment") + + et_only_cooldown_active = _is_et_only_cooldown_active( + start_t, candidate_sources + ) + + if et_only_cooldown_active: + _record_stm(stm_full) + session_clip_ids.append(stm_full.id) + session_details.append(_mtm_detail_from_stm(stm_full)) + prev_stm = stm_full + current_event_table = _update_event_table_for_run( + current_event_table, stm_full, force_continue=True + ) + _append_event_table_record( + event_id=str(current_event_id) if current_event_id else "", + clip_id=str(stm_full.id), + time_range=[start_t, end_t], + proposed_et=proposed_event_table, + committed_et=current_event_table, + decision={ + "mode": "et_only_cooldown_suppressed", + "candidate_sources": candidate_sources, + "last_accepted_split_time_s": last_accepted_split_time_s, + "cooldown_gap_s": start_t - last_accepted_split_time_s, + "cooldown_window_s": min_split_gap_s, + }, + ) + _append_jsonl( + event_link_log_path, + { + "ts": datetime.utcnow().isoformat() + "Z", + "meta": { + "segment_id": int(seg.get("segment_id", 0) or 0), + "time_range": [start_t, end_t], + "delta": delta, + "boundary_source": "ET", + "boundary_time": start_t, + "boundary_time_hhmmss": _seconds_to_hhmmss(start_t), + "last_accepted_split_time_s": last_accepted_split_time_s, + "cooldown_gap_s": start_t - last_accepted_split_time_s, + "cooldown_window_s": min_split_gap_s, + }, + "candidate_source": _format_et_shift_candidate(delta, "et"), + "pre_ET": _event_link_pre_et_payload(current_event_table), + "current_segments": _event_link_current_segments_payload( + stm_full, stm_full + ), + "proposed_event_table": proposed_event_table, + "committed_event_table": current_event_table, + "note": "ET-only candidate suppressed by nearby cooldown after a recent confirmed cut.", + }, + ) + continue + + if candidate_sources and current_event_table is not None: + boundary_time = ( + min(segment_chapter_boundaries) + if segment_chapter_boundaries + else start_t + ) + source_tag = ( + "ET+chapter_segment" + if len(candidate_sources) > 1 + else candidate_sources[0] + ) + parts: List[str] = [] + has_asr_source = "chapter_segment" in candidate_sources + has_et_source = "ET" in candidate_sources + if has_asr_source: + parts.append(_chapter_shift_text(boundary_time)) + if has_et_source: + et_evidence = "both" if has_asr_source else "et" + parts.append(_format_et_shift_candidate(delta, et_evidence)) + prompt_candidate_source = " | ".join(parts) if parts else "unknown" + pending_active = True + pending_anchor_info = { + "left_event_table": current_event_table, + "left_event_id": current_event_id, + "left_session_clip_ids": list(session_clip_ids), + "left_session_details": list(session_details), + "candidate_sources": list(candidate_sources), + "candidate_source_text": prompt_candidate_source, + "trigger_delta": delta, + "anchor_segment_id": int(seg.get("segment_id", 0) or 0), + "anchor_start_t": start_t, + "anchor_end_t": end_t, + "boundary_source": source_tag, + "boundary_time": boundary_time, + } + pending_stms = [stm_full] + pending_event_table = _update_event_table_for_run(None, stm_full) + _append_jsonl( + event_link_log_path, + { + "ts": datetime.utcnow().isoformat() + "Z", + "meta": { + "segment_id": int(seg.get("segment_id", 0) or 0), + "time_range": [start_t, end_t], + "delta": delta, + "boundary_source": source_tag, + "boundary_time": boundary_time, + "boundary_time_hhmmss": _seconds_to_hhmmss(boundary_time), + "lookahead_clips": 1, + }, + "candidate_source": prompt_candidate_source, + "pre_ET": _event_link_pre_et_payload(current_event_table), + "current_segments": _event_link_current_segments_payload( + stm_full, stm_full + ), + "proposed_event_table": proposed_event_table, + "pending_event_table": pending_event_table, + "note": "Pending opened: delayed confirmation will judge after one right clip.", + }, + ) + continue + else: + _record_stm(stm_full) + session_clip_ids.append(stm_full.id) + session_details.append(_mtm_detail_from_stm(stm_full)) + prev_stm = stm_full + if _is_initialization_shift(delta): + current_event_table = _update_event_table_for_run( + current_event_table, stm_full, force_continue=True + ) + decision_mode = "no_judge_force_continue_init_shift" + else: + current_event_table = proposed_event_table + decision_mode = "no_judge_direct_commit" + _append_event_table_record( + event_id=str(current_event_id) if current_event_id else "", + clip_id=str(stm_full.id), + time_range=[start_t, end_t], + proposed_et=proposed_event_table, + committed_et=current_event_table, + decision={"mode": decision_mode}, + ) + + if pending_active and pending_anchor_info is not None and pending_stms: + left_event_table = pending_anchor_info.get("left_event_table") + left_event_id = pending_anchor_info.get("left_event_id") + left_session_ids = list( + pending_anchor_info.get("left_session_clip_ids") or [] + ) + left_session_detail_list = list( + pending_anchor_info.get("left_session_details") or [] + ) + current_event_table = _rebuild_event_table_from_old( + left_event_table, pending_stms + ) + current_event_id = ( + left_event_id if isinstance(left_event_id, UUID) else current_event_id + ) + for pstm in pending_stms: + _record_stm(pstm) + session_clip_ids = left_session_ids + [s.id for s in pending_stms] + session_details = left_session_detail_list + [ + _mtm_detail_from_stm(s) for s in pending_stms + ] + prev_stm = pending_stms[-1] + for idx, pstm in enumerate(pending_stms, start=1): + _append_event_table_record( + event_id=str(current_event_id) if current_event_id else "", + clip_id=str(pstm.id), + time_range=[pstm.time_range[0], pstm.time_range[1]], + proposed_et=pending_event_table, + committed_et=current_event_table, + decision={ + "mode": "pending_flush_reject_no_lookahead", + "pending_index": idx, + "pending_size": len(pending_stms), + "candidate_sources": pending_anchor_info.get( + "candidate_sources" + ) + or [], + }, + ) + pending_active = False + pending_anchor_info = None + pending_stms = [] + pending_event_table = None + + if session_clip_ids: + _finalize_session() + finally: + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + # Ensure any distributed/NCCL groups are torn down to avoid leaks + try: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + except RuntimeError as e: + logger.warning("VideoPipeline: failed to destroy process group: %s", e) + + logger.info( + "VideoPipeline: completed clips=%d events=%d", + len(short_terms), + len(medium_terms), + ) + # -------------------- Cleanup temporary artifacts -------------------- + if cleanup: + import shutil + + # Remove extracted segment mp4s + try: + if segments_dir.exists(): + shutil.rmtree(segments_dir, ignore_errors=True) + except OSError as e: + logger.warning("VideoPipeline: failed to remove segments: %s", e) + + # All pipeline files are temporary; memories have already been materialized. + try: + if outputs_dir.exists(): + shutil.rmtree(outputs_dir, ignore_errors=True) + except OSError as e: + logger.warning("VideoPipeline: failed to remove outputs: %s", e) + + try: + if asr_tmp_work_dir.exists(): + shutil.rmtree(asr_tmp_work_dir, ignore_errors=True) + except OSError as e: + logger.warning("VideoPipeline: failed to remove ASR temporary files: %s", e) + + # Legacy pipeline may create work_root/chunks — remove if present + try: + legacy_chunks = Path(work_root) / "chunks" + if legacy_chunks.exists(): + shutil.rmtree(legacy_chunks, ignore_errors=True) + except OSError as e: + logger.warning("VideoPipeline: failed to remove legacy chunks: %s", e) + + return { + "short_term": [memory.to_dict() for memory in short_terms], + "medium_term": [memory.to_dict() for memory in medium_terms], + "event_table": event_tables, + } diff --git a/jiuwen_memory/common/normalizer/normalizer_impl/video_prompts.py b/jiuwen_memory/common/normalizer/normalizer_impl/video_prompts.py new file mode 100644 index 00000000..ea0de5bb --- /dev/null +++ b/jiuwen_memory/common/normalizer/normalizer_impl/video_prompts.py @@ -0,0 +1,806 @@ +# ruff: noqa: E501 + +CHAPTER_CHUNK = """ +You are given a time-aligned transcript of a video. +Each line contains a timestamp and the corresponding spoken content. + +Your task is to segment the transcript into coherent CHAPTERS based on semantic, structural, or intent transitions. +Do NOT divide chapters into arbitrarily small segments. + +--------------------------------------- +Core Segmentation Principles +--------------------------------------- + +A chapter should represent a coherent and self-contained unit of discourse. + +Chapters may be formed based on: + +1) Structural Unit Shift (HIGHEST PRIORITY) + - Explicit section/unit markers: + numbered items, steps, parts, episodes, cases, + years/phases, labeled segments. + - Clear transition phrases: + "next", "now let's", "moving on", "part two", etc. + - When a new named unit becomes dominant, + this strongly indicates a boundary. + +2) Semantic Topic Shift + - Clear change in subject, goal, or organizing focus. + - A new discussion thread or objective becomes primary. +--------------------------------------- +Granularity Control +--------------------------------------- + +- Do NOT segment by fixed time intervals. +- Do NOT over-segment minor tone or sentence changes. +- Prefer meaningful structural or semantic transitions. +- In clearly structured formats, align chapters with major units. + +--------------------------------------- +Low-Information Handling +--------------------------------------- + +If semantic signal is weak: + +- Still attempt segmentation based on any identifiable transitions. +- Always output the inferred chapters. +- Use segmentation_confidence to reflect boundary reliability, + NOT to suppress segmentation. + +--------------------------------------- +Coverage Rules +--------------------------------------- + +- Chapters must fully cover the entire timeline. +- No gaps or overlaps. +- Timestamp gaps belong to the previous chapter. +- Chapters must remain minimally coherent. + +--------------------------------------- +Required Output Fields +--------------------------------------- + +Each chapter must include: +- chapter_id +- start_time +- title (<= 12 words) +- summary (1 concise sentence) + +--------------------------------------- +Segmentation Confidence (REDEFINED) +--------------------------------------- + +segmentation_confidence reflects the RELIABILITY of ASR-based boundary evidence, +especially whether boundaries are structurally explicit and precisely alignable. + +Set as: + +Assign ASR segmentation confidence using only two levels: HIGH or LOW. + +HIGH: +- ASR provides reliable evidence for semantic progression, topic shift, task shift, or stage transition, +- and it is meaningfully useful for chapter segmentation, +- even if the structure is not explicitly labeled with phrases like "part 1" or "next section". + +LOW: +- ASR is not a reliable basis for chapter segmentation. +- This includes cases where the audio is dominated by: + - lyrics, singing, repeated chorus-like language, + - pantomime / near-silent performance, + - music or environmental sound with little usable speech, + - extremely sparse, noisy, or badly recognized speech, + - casual fragmented dialogue that does not provide stable structural cues. +- In these cases, chapter boundaries would be mostly subjective or visually inferred, and ASR offers little reliable anchoring. + +Important: +Use LOW only when ASR is genuinely weak or uninformative for segmentation, +not merely because explicit section markers are absent. + +Important: +LOW means ASR boundary evidence is weak, +NOT that segmentation should be avoided. + +--------------------------------------- +Output Requirements +--------------------------------------- + +- Output MUST be valid JSON. +- Output ONLY the JSON object. +- No explanations or extra text. + +JSON schema: + +{ + "segmentation_confidence": "high|low", + "chapters": [ + { + "chapter_id": number, + "start_time": "HH:MM:SS", + "title": string, + "summary": string + } + ] +} +""" + +CAPTION_PROMPT = """ +# Role +You convert the CURRENT video segment(clip) into a precise visual memory unit. + +# Event Table (ET) +ET records the ongoing event context from earlier clips. +Use ET only as a weak continuity cue for the CURRENT clip. + +You may use ET to: +- keep naming of visibly recurring people / objects / places consistent +- recognize when the same visible event or setup continues +- maintain light local continuity in description + +If the clip clearly continues the same visible event, you may use light continuity phrases. + +Rules: +- ET is NOT ground truth and may be outdated. +- NEVER describe something only because ET mentions it. +- NEVER copy ET text. +- If ET conflicts with the current visuals, trust the current clip. +- If something is not visually present now, do not include it just for continuity. +- If ET is empty, ignore it. + +{event_table} + +# Task +Describe the CURRENT segment as completely and accurately as possible. + +Focus on: +- people, objects, animals +- actions and interactions +- scene layout and environment +- on-screen text or graphics +- visible temporal progression + +# Priority +1. The CURRENT clip is the source of truth. +2. This is NOT an event summary. +3. Do not omit visible details for the sake of continuity. + +# Requirements +- Capture the main visible content faithfully. +- Preserve important concrete details that may matter for later QA or retrieval. +- Prefer temporally ordered description when actions unfold over time. +- Be factual and visually grounded. +- Prefer concrete visual descriptors (clothing, color, objects, gestures). +- Do not infer unseen events. + +# Output JSON +{ + "detailed_caption": "Temporally ordered visual description grounded in the current clip.", + "visual_summary": "Short subject-verb-object phrase for indexing.", + "environment": "Scene description." +} + +# Constraints +- Output JSON only. +- Do not add new keys. +- visual_summary ≤ 15 words. +""" + + +SESSION_SUMMARY_PROMPT_TEMPLATE = """ +# Role +You are an event memory consolidation and reasoning assistant. +Your task is to merge multiple short video segment descriptions into ONE coherent EVENT-level memory. + +# Input +The following atomic segment descriptions are ordered by time: +{detail_list} + +Each segment may contain: +- visual observations of actions or scene changes +- ASR transcripts capturing spoken dialogue + +ASR transcripts may contain recognition errors, but they can provide useful clues about goals, relationships, decisions, or conflicts that are not fully visible. + +# Objective +Construct an event memory with TWO layers: + +1) Event Narrative (what happened) +Reconstruct the episode as a coherent sequence with clear temporal and causal structure by integrating both visual actions and spoken dialogue when relevant. + +2) Semantic Inference (what higher-level knowledge can be inferred from the event) +Provide a compact semantic memory that captures the event’s likely goal, relationship clues, roles, habits, or deeper task structure based on the full sequence. + +# Instructions + +## 1. Topic Label +Provide a concise label capturing the core episode or objective. +Avoid generic labels like "people talking" or "walking". + +## 2. Event Narrative +Merge all segments into a single coherent narrative. + +The narrative should clearly describe: +- Initial context or situation +- Trigger or turning point +- Key actions, interactions, or attempts +- Important spoken exchanges when they affect the meaning of the event +- Final consequence or state change + +Preserve chronological order but organize actions into meaningful phases rather than listing clips. + +Use causal or temporal connectors when appropriate. + +Preserve important factual details useful for QA, including: +entities, objects, locations, numbers, significant visual changes, and relevant spoken information. + +## 3. Semantic Inference +Based on the entire event, provide a compact high-level semantic memory. + +Focus on extracting useful higher-level knowledge beyond the surface narrative, such as: +- the apparent goal or objective of the event +- possible relationship dynamics between key entities +- roles, habits, preferences, or capabilities suggested by repeated behavior +- whether repeated actions are serving a larger task, rescue, conflict, plan, or decision process + +Use dialogue as supporting evidence when it helps reveal goals, relationships, decisions, or conflicts that are not fully visible. + +IMPORTANT: +- Do not simply restate the narrative. +- Prefer useful semantic abstraction over general commentary. +- Avoid vague statements about symbolism, creativity, atmosphere, or moral meaning. +- Separate direct evidence from interpretation. +- Use cautious wording when uncertain. +- Do NOT convert weak inference into confirmed facts. + +## 4. Evidence Discipline +- Only infer motivations or causal explanations supported by the provided descriptions. +- Treat ASR transcripts cautiously if they appear noisy, incomplete, or isolated. +- If evidence is incomplete, preserve uncertainty rather than forcing conclusions. +- Do NOT hallucinate missing events or identities. + +# Output Format +Return ONLY the following JSON: + +{ + "topic_label": "...", + "full_narrative": "...", + "semantic_inference": "..." +} + +# Constraints +1) Do not output a list of clips; produce one unified event memory. +2) Focus on causal structure and state transitions rather than exhaustive low-level detail. +3) Semantic inference must be higher-level than the narrative and should add new useful knowledge. +4) Output JSON only. No explanations or markdown. +""" + +TIME_SUMMARY_PROMPT = """ +# Role +You are a neutral video summarizer responsible for consolidating multiple atomic video segments +that fall within the SAME FIXED TIME WINDOW into a coherent temporal summary. + +# Input Context +Below is a list of atomic segment descriptions ordered by time. +All segments come from a continuous, fixed-length time window of the video +(not necessarily aligned with a single real-world event): +{detail_list} + +# Instructions +1. Window Labeling: + - Provide a short descriptive label summarizing what is MOSTLY shown in this time window. + - Do NOT assume the window corresponds to a complete task or a well-defined event. + - Prefer descriptive phrases over intentional or goal-oriented wording. + +2. Temporal Consolidation: + - Merge the segments into a fluent, time-ordered narrative. + - Remove obvious redundancy while preserving chronological progression. + - Describe what happens, changes, or appears over time within this window. + - Retain important observable details such as objects, people, actions, locations, + quantities, colors, on-screen text, and notable interactions. + +# Output Format +Return ONLY the following JSON object: + +{{ + "topic_label": "Descriptive label for this time window", + "full_narrative": "A coherent, time-ordered summary of what occurs within this fixed window." +}} +""" + + +EVENT_RELATIONSHIP = """ +You are an event relation classifier for long-form narrative videos. + +You will be given TWO events (Event A happens before Event B). +Each event is a summarized description produced by an upstream pipeline. + +Your task is to determine whether Event B is narratively connected to Event A, +and if so, how. + +The event graph should reflect STORY CONTINUITY within the video itself, +not generic similarity or real-world common sense. + +It is better to produce a sparse but meaningful graph +than a dense graph with weak or generic edges. + +-------------------------------- +RELATION DEFINITIONS +-------------------------------- +- "causes": + Event A directly leads to Event B. + This requires explicit or clearly described causal linkage. + +- "enables": + Event A establishes a concrete prerequisite for Event B + (e.g., a decision, an action, or new information that allows B to happen). + +- "elaborates": + Event B continues the SAME specific storyline or subplot as Event A. + This must be traceable as a continuation of the same situation, + not merely a similar type of activity. + +- "contrasts": + Event B represents a clear narrative shift away from Event A + (e.g., escalation vs resolution, cooperation vs conflict). + +- "unrelated": + No narratively traceable connection. + +-------------------------------- +STORY ANCHOR RULE (CRITICAL) +-------------------------------- +You may output "causes", "enables", or "elaborates" +ONLY IF there is a SHARED STORY ANCHOR. + +A valid story anchor is: +- A recurring character or group implicitly or explicitly referenced, +- OR a recurring concrete storyline element + (such as a specific situation, decision, plan, accusation, object, or outcome). + +The anchor must be specific enough that a viewer could recognize +that both events belong to the same storyline. + +-------------------------------- +INVALID CONNECTIONS (DO NOT USE) +-------------------------------- +The following are NOT valid reasons to connect events: +- Shared setting or environment. +- Shared profession or role. +- Shared emotional tone or intensity. +- Shared activity type that appears frequently in the video. +- General world knowledge or genre conventions. + +If the connection relies primarily on these, +the correct relation is "unrelated". + +-------------------------------- +SOFT CONTINUITY (RESTRICTED) +-------------------------------- +Soft continuity is allowed ONLY when: +- Both events clearly belong to the same identifiable subplot, +- AND you can explicitly explain how Event B advances or revisits + the situation introduced in Event A. + +If you cannot answer: +"How does Event B specifically follow from Event A in this video?" +then you MUST choose "unrelated". + +-------------------------------- +CONFIDENCE RULES +-------------------------------- +- Clear shared storyline element: 0.65–0.90 +- Implicit but traceable continuation: 0.45–0.65 +- Weak but defensible continuity: 0.35–0.45 +- No clear story anchor: ≤ 0.30 and MUST be "unrelated" + +-------------------------------- +EVIDENCE RULES +-------------------------------- +- Quotes must reference concrete storyline elements, + not abstract descriptions or generic activities. +- The "reason" MUST explicitly name the shared story anchor. +- Reasons based on generic similarity are INVALID. + +-------------------------------- +INPUT +-------------------------------- +Event A: +- task_id: {a_id} +- time_span: {a_start}-{a_end} +- topic: {a_topic} +- narrative_summary: {a_narr} + +Event B: +- task_id: {b_id} +- time_span: {b_start}-{b_end} +- topic: {b_topic} +- narrative_summary: {b_narr} + +Also provided: +- time_gap_seconds: {gap_s} + +-------------------------------- +OUTPUT +-------------------------------- +Output a JSON object with this exact schema: +{{ + "src_task_id": "{a_id}", + "dst_task_id": "{b_id}", + "relation": "causes|enables|elaborates|contrasts|unrelated", + "confidence": 0.0, + "evidence": {{ + "src_quote": "", + "dst_quote": "", + "reason": "" + }} +}} +""" + + +UPDATE_EVENT_TABLE_PROMPT = """ +Update the Event Table (ET) after each clip. + +ET represents ONE ongoing chapter-worthy local unit in the video. +It stores the cumulative state of that unit so far. + +ET is NOT a per-clip summary. +ET is NOT a full-video summary. +ET should preserve a stable local anchor across adjacent clips whenever the same unit continues. + +Input: +(A) pre_ET (or null) +(B) STM_k (current clip summary) + +Return STRICT JSON: + +{ + "event_identity": "...", + "event_summary": "...", + "entities": [...], + "open_questions": [...], + "delta": "..." +} + +-------------------------------------------------- +INITIALIZATION + +If pre_ET is null: +- create a new local unit from STM_k +- event_identity should be short and local +- event_summary should describe the current unit state +- delta MUST be: + +"SHIFT: NONE -> ." + +-------------------------------------------------- +CORE IDEA + +ET is the working state of the CURRENT chapter-worthy local unit. + +A good ET should help preserve a video structure that is: +- locally coherent +- viewer-friendly +- navigable +- not fragmented by tiny visual variation +- not so broad that multiple meaningful parts are collapsed together + +If the same unit continues: +- keep event_identity unchanged by default +- update event_summary by integrating prior state + new info +- do NOT rewrite ET to mirror STM_k +- do NOT broaden ET into a global or session-level summary + +ET should stay local, stable, and moderately slow-changing. + +-------------------------------------------------- +EVENT IDENTITY + +event_identity is the short title of the current local unit. + +It should name the stable semantic unit currently unfolding in the video. + +It must be: +- local +- stable +- specific enough to anchor the current unit +- not tied to one shot, object, or camera angle +- not so broad that it absorbs multiple meaningful parts of the video + +Good identities are usually: +- a current task +- a phase +- a section +- a stage +- a segment focus +- a unit being explained, built, performed, presented, or completed + +If several consecutive clips still serve the same chapter-worthy unit, keep the same identity. + +-------------------------------------------------- +WHEN TO CONTINUE + +Prefer CONTINUE when: +- the same task, section, phase, project, stage, or focal unit is still unfolding +- the clip adds examples, sub-steps, local details, reactions, or new views of the same unit +- the visual focus changes but the same unit still explains the clip well + +-------------------------------------------------- +WHEN TO SHIFT + +Use SHIFT only when the current clip is better understood as the start of a NEW local unit, such as: + +- a new section, phase, stage, or named unit begins +- a different task, function, or focal activity becomes dominant +- a different artifact, outcome, segment focus, or block becomes central +- the video clearly moves into a new part such as setup, explanation, execution, showcase, recap, closing, award, or outro +- the clip contains a clear opener or reset for a new unit rather than just more detail of the old one +- a clear visual chapter/title card appears for the first time, such as a title screen, large on-screen heading, chapter card, or OCR-visible segment title +- a strong visible marker explicitly tells the viewer that a new part of the video has begun + +A good SHIFT should make the video structure clearer and more viewer-friendly. + +Do NOT use SHIFT for weak wording drift alone. +Do NOT keep everything under one broad umbrella just because it is loosely related. + +-------------------------------------------------- +DO NOT SHIFT FOR + +- new examples of the same concept +- local sub-steps inside the same task / phase / project +- new objects used within the same ongoing unit +- camera, actor, or scene changes +- temporary inserts, montage, credits, subscribe prompts, or transition screens + unless they clearly mark a real new unit + +-------------------------------------------------- +SUMMARY RULE + +event_summary = compressed cumulative state of the same local unit. + +Update it by: +- starting from pre_ET.event_summary +- integrating new relevant information from STM_k +- compressing back into a short state description + +Keep only information still relevant to the SAME unit. +Do not let the summary become a broad session summary. + +-------------------------------------------------- +DELTA FORMAT + +CONTINUE: +"CONTINUE: ." + +SHIFT: +"SHIFT: -> ." + +If SHIFT is used: +- must match event_identity +- SHIFT means a proposed unit change, not a final segmentation decision + +-------------------------------------------------- +LIMITS + +event_identity <= 10 tokens +event_summary <= 80 tokens +entities <= 4 +open_questions <= 2 +delta <= 20 tokens + +Return JSON only. +""" + +EVENT_LINK_WITH_ET_PROMPT_AIO = """ +# Role + +You are a delayed-confirmation boundary gate for video chaptering. + +A candidate boundary was proposed earlier at candidate_t. + +Your task: +Given the previous ET state, the ancor segment that triggered the candidate, +and the later pending segment used for delayed confirmation, +decide whether the candidate boundary should be: + +- REJECTED: KEEP the same chapter-worthy local unit as pre_ET +- ACCEPTED: SPLIT into a new chapter-worthy local unit + +If splitting, output exactly ONE split point. + +Your goal is to produce boundaries that are: +- structurally clear +- viewer-friendly +- navigable +- not fragmented by tiny local variation +- not so coarse that multiple meaningful parts are collapsed together + +-------------------------------------------------- +# Inputs + +candidate_source: {candidate_source} +asr_confidence: {asr_confidence} + +asr_chapter_context: +{chapter_context} + +pre_ET: +{pre_et} + +anchor_segment: +Summary: {anchor_summary} +Caption: {anchor_caption} +ASR lines: +{anchor_ASR} + +pending_segment: +Summary: {pending_summary} +Caption: {pending_caption} +ASR lines: +{pending_ASR} + +-------------------------------------------------- +# Core Definition + +SAME UNIT means: +The anchor segment and pending segment are still best understood +as continuation of the same chapter-worthy local unit represented by pre_ET, +even if examples, objects, views, or local steps change. + +SPLIT means: +The anchor segment is better treated as the beginning of a new local unit, +and the pending segment continues that new unit rather than the old one. + +A boundary should be accepted when it makes the video structure clearer for a viewer. + +-------------------------------------------------- +# Main Decision Principle + +Use pre_ET as the old-unit anchor. +Use anchor_segment as the candidate-triggering evidence. +Use pending_segment as right-context confirmation. + +Ask: + +1) Does pre_ET still explain both segments well as one coherent local unit? +2) Or does anchor_segment already begin a different local unit, + which pending_segment then continues? +3) Would splitting here create a clearer and more navigable chapter structure? + +If the right-side evidence mainly continues the same underlying unit, +prefer KEEP. + +If the right-side evidence confirms a different section / phase / stage / task / segment focus, +prefer SPLIT. + +-------------------------------------------------- +# Evidence Use + +Use all available evidence, with these priorities: + +1) clear structural or semantic change across pre_ET -> anchor_segment -> pending_segment +2) raw ASR when it is informative +3) pre_ET.event_identity and event_summary as the old-unit anchor +4) visual summaries / captions as supporting evidence + +Important: +- ASR is strong evidence when it is informative, but not automatically decisive. +- pre_ET is a useful anchor, but not absolute truth. +- visual evidence alone is usually weak, unless it clearly marks a real new unit. + +-------------------------------------------------- +# ASR Confidence Rule + +If asr_confidence is HIGH: +- treat ASR as strong evidence for section / task / topic / stage progression +- a clear semantic shift in ASR may support SPLIT even without explicit labels + +If asr_confidence is LOW: +- treat ASR as weak evidence only +- do NOT split mainly because of lyrics, repeated sung lines, sparse speech, + noisy ASR, or fragmented casual dialogue +- in LOW-confidence cases, require stronger ET or visual support before splitting + +-------------------------------------------------- +# Strong Split Evidence + +Accept SPLIT when one or more of the following is clearly true: + +A) anchor_segment already begins a new section, phase, stage, segment focus, or local unit, + and pending_segment continues that new unit + +B) the dominant task, function, outcome, or video focus changes, + and pending_segment confirms the new direction + +C) there is a clear visible or spoken opener, reset, title, marker, or transition + for a new local unit + +D) a clear visual chapter/title card appears for the first time, such as a title screen, + chapter card, large on-screen heading, or OCR-visible segment title that explicitly marks a new part + +E) keeping both segments under pre_ET would make the chapter structure too coarse or less clear + +-------------------------------------------------- +# Weak Evidence (Not sufficient alone) + +Do NOT split based only on: +- camera or shot changes +- different people or objects +- new examples +- local sub-steps inside the same task / section +- wording drift between summaries +- short inserts, montage, credits, subscribe prompts + +-------------------------------------------------- +# Anti-Oversegmentation Rule + +Do NOT split merely because: +- the new segment could have a slightly different short title +- the clip looks different +- a new object appears +- a small sub-step appears inside the same broader local unit + +-------------------------------------------------- +# Anti-Undersplitting Rule + +Do NOT keep merely because: +- the scene is similar +- the same people remain +- the topic is broadly related +- everything belongs to the same broad activity domain + +If anchor_segment starts a meaningfully new local unit and pending_segment confirms it, +split. + +-------------------------------------------------- +# Split Point Selection + +If is_same_event = false, output exactly ONE split point. + +Choose in this order: + +1) the START timestamp of the raw ASR line in anchor_segment or pending_segment that clearly begins the new unit +2) otherwise the earliest explicit visible section onset, title card, chapter card, large heading, or OCR-visible segment title that marks the new unit +3) otherwise candidate_t + +split_point.reason must be one of: +- structural_transition +- goal_shift +- anchor_shift +- entity_shift +- scene_break + +Use: +- structural_transition for explicit section / phase / stage / opener / closing / title-card transition +- goal_shift for clear task / function reset +- anchor_shift for a new dominant local unit without explicit marker +- entity_shift only when a new main entity truly replaces the old core focus +- scene_break only for a real semantic jump in scene / time / location + +-------------------------------------------------- +# Output + +Return STRICT JSON ONLY. + +If KEEP: +{ + "rationale": "One short sentence.", + "is_same_event": true, + "split_point": [] +} + +If SPLIT: +{ + "rationale": "One short sentence.", + "is_same_event": false, + "split_point": [ + { + "t": "HH:MM:SS", + "reason": "anchor_shift|goal_shift|entity_shift|scene_break|structural_transition" + } + ] +} + +Rules: +- If is_same_event=true -> split_point must be [] +- If is_same_event=false -> exactly one split point +- Return JSON only +""" diff --git a/jiuwen_memory/construction/extractor_impl/__init__.py b/jiuwen_memory/construction/extractor_impl/__init__.py index 29927b18..ae99d97e 100644 --- a/jiuwen_memory/construction/extractor_impl/__init__.py +++ b/jiuwen_memory/construction/extractor_impl/__init__.py @@ -11,5 +11,6 @@ import_module(".keyword_extractor", __name__) import_module(".llm_extractor", __name__) import_module(".dynamic_llm_extractor", __name__) +import_module(".video_memory_extractor", __name__) __all__ = ["ExtractorProducer"] diff --git a/jiuwen_memory/construction/extractor_impl/video_memory_extractor.py b/jiuwen_memory/construction/extractor_impl/video_memory_extractor.py new file mode 100644 index 00000000..1b7f7508 --- /dev/null +++ b/jiuwen_memory/construction/extractor_impl/video_memory_extractor.py @@ -0,0 +1,200 @@ +"""Construction extractor for hierarchical video memories.""" + +from __future__ import annotations + +import json +import uuid +from copy import deepcopy +from typing import Any + +from jiuwen_memory.common.errors import BackendError +from jiuwen_memory.common.type_def import MemoryTier, MemoryUnit, Segment +from jiuwen_memory.construction.base import ExtractContext, OperatorType +from jiuwen_memory.construction.extractor import Extractor, ExtractorProducer + + +class VideoMemoryExtractor(Extractor): + """Convert normalized video data into CLM and ELM MemoryUnits.""" + + def operator_type(self) -> OperatorType: + return OperatorType.EXTRACTOR + + def health(self) -> None: + return None + + def extract( + self, + units: list[MemoryUnit], + *, + context: ExtractContext | None = None, + ) -> list[MemoryUnit]: + del context + derived: list[MemoryUnit] = [] + for source in units: + video_data = _load_video_data(source) + clips = _build_clips(source, video_data) + events = _build_events(source, video_data, clips) + derived.extend(clips.values()) + derived.extend(events) + return derived + + +def _build_clips( + source: MemoryUnit, + video_data: dict[str, Any], +) -> dict[str, MemoryUnit]: + video_id = str(video_data.get("payload_id") or source.source_ref) + clips: dict[str, MemoryUnit] = {} + for clip in video_data["clips"]: + source_id = str(clip.get("id", "")).strip() + if not source_id: + continue + if source_id in clips: + raise BackendError(f"duplicate video clip id: {source_id!r}") + metadata = _memory_metadata( + source, + level="clm", + video_id=video_id, + source_id=source_id, + start=clip.get("start_seconds"), + end=clip.get("end_seconds"), + ) + clips[source_id] = MemoryUnit( + id=str(uuid.uuid4()), + scope=source.scope, + tier=MemoryTier.EPISODIC, + segments=[ + Segment( + content=_memory_content( + ("Visual summary", clip.get("visual_summary")), + ("Detailed caption", clip.get("detailed_caption")), + ("Speech transcript", clip.get("asr")), + ("Environment", clip.get("environment")), + ), + assets=list(source.assets), + source=source.source, + ) + ], + source_ref=source.source_ref, + temporal=deepcopy(source.temporal), + provenance=[source.id], + tags=list(source.tags), + metadata=metadata, + ) + return clips + + +def _build_events( + source: MemoryUnit, + video_data: dict[str, Any], + clips: dict[str, MemoryUnit], +) -> list[MemoryUnit]: + video_id = str(video_data.get("payload_id") or source.source_ref) + events: list[MemoryUnit] = [] + for event in video_data["events"]: + source_id = str(event.get("id", "")).strip() + child_ids = event.get("clip_ids", []) + if not source_id or not isinstance(child_ids, list): + continue + normalized_child_ids = [str(child) for child in child_ids] + missing = [child for child in normalized_child_ids if child not in clips] + if missing: + raise BackendError( + f"video event {source_id!r} references missing clips: {missing}" + ) + metadata = _memory_metadata( + source, + level="elm", + video_id=video_id, + source_id=source_id, + start=event.get("start_seconds"), + end=event.get("end_seconds"), + ) + metadata["child_clm_source_ids"] = json.dumps( + normalized_child_ids, ensure_ascii=False, separators=(",", ":") + ) + events.append( + MemoryUnit( + id=str(uuid.uuid4()), + scope=source.scope, + tier=MemoryTier.EPISODIC, + segments=[ + Segment( + content=_memory_content( + ("Topic", event.get("topic")), + ("Event summary", event.get("summary")), + ("Semantic inference", event.get("semantic_inference")), + ), + assets=list(source.assets), + source=source.source, + ) + ], + source_ref=source.source_ref, + temporal=deepcopy(source.temporal), + provenance=[source.id], + tags=list(source.tags), + metadata=metadata, + ) + ) + return events + + +def _load_video_data(unit: MemoryUnit) -> dict[str, Any]: + try: + value = json.loads(unit.content) + except json.JSONDecodeError as exc: + raise BackendError("video memory content is not valid JSON") from exc + if not isinstance(value, dict): + raise BackendError("video memory content must be an object") + for field in ("clips", "events"): + items = value.get(field) + if not isinstance(items, list) or not all(isinstance(item, dict) for item in items): + raise BackendError(f"video memory {field} must be a list of objects") + return value + + +def _memory_metadata( + source: MemoryUnit, + *, + level: str, + video_id: str, + source_id: str, + start: object, + end: object, +) -> dict[str, Any]: + metadata = dict(source.metadata) + metadata.update( + { + "modal_type": "multimodal", + "memory_level": level, + "video_id": video_id, + "source_memory_id": source_id, + "start_seconds": _memory_float(start, field=f"{level}.start_seconds"), + "end_seconds": _memory_float(end, field=f"{level}.end_seconds"), + } + ) + return metadata + + +def _memory_content(*parts: tuple[str, object]) -> str: + lines = [ + f"{label}: {str(value).strip()}" + for label, value in parts + if str(value or "").strip() + ] + if not lines: + raise BackendError("video memory item has no textual content") + return "\n".join(lines) + + +def _memory_float(value: object, *, field: str) -> float: + try: + return float(value) + except (TypeError, ValueError) as exc: + raise BackendError(f"video memory {field} must be numeric") from exc + + +@ExtractorProducer.register("video_memory") +def _build(config): + del config + return VideoMemoryExtractor() diff --git a/jiuwen_memory/control/engine_impl/cloud_engine.py b/jiuwen_memory/control/engine_impl/cloud_engine.py index 9dee85c0..7db9234d 100644 --- a/jiuwen_memory/control/engine_impl/cloud_engine.py +++ b/jiuwen_memory/control/engine_impl/cloud_engine.py @@ -11,9 +11,9 @@ import asyncio import copy import uuid -from typing import Any from dataclasses import dataclass from datetime import datetime, timezone +from typing import Any from jiuwen_memory.common.errors import AgentMemoryError, NotFoundError, ValidationError from jiuwen_memory.common.log import get_logger @@ -247,15 +247,18 @@ async def write( *, assets: list[str] | None = None, tags: list[str] | None = None, - metadata: dict[str, str] | None = None, + metadata: dict[str, Any] | None = None, occurred_at: datetime | None = None, ) -> list[MemoryUnit]: meta = self._normalized_metadata(metadata) + is_video = source == Modality.VIDEO + payload_id = str(meta.get("payload_id", "")).strip() if is_video else "" payload = RawPayload( - id=str(uuid.uuid4()), + id=payload_id or str(uuid.uuid4()), scope=scope, modality=source, - data=content.encode("utf-8"), + data=b"" if is_video else content.encode("utf-8"), + uri=content if is_video else "", metadata=meta, occurred_at=occurred_at, ) diff --git a/jiuwen_memory/control/engine_impl/in_memory_engine.py b/jiuwen_memory/control/engine_impl/in_memory_engine.py index 35e0daf3..a224ac38 100644 --- a/jiuwen_memory/control/engine_impl/in_memory_engine.py +++ b/jiuwen_memory/control/engine_impl/in_memory_engine.py @@ -270,11 +270,14 @@ def _is_true(key: str) -> bool: "metadata.middle=true requires infer=true (middle 是 infer 下的二级开关)" ) + is_video = source == Modality.VIDEO + payload_id = str(meta.get("payload_id", "")).strip() if is_video else "" payload = RawPayload( - id=str(uuid.uuid4()), + id=payload_id or str(uuid.uuid4()), scope=scope, modality=source, - data=content.encode("utf-8"), + data=b"" if is_video else content.encode("utf-8"), + uri=content if is_video else "", metadata=meta, occurred_at=occurred_at, ) diff --git a/jiuwen_memory/control/ingest_job.py b/jiuwen_memory/control/ingest_job.py new file mode 100644 index 00000000..6905b9ce --- /dev/null +++ b/jiuwen_memory/control/ingest_job.py @@ -0,0 +1,287 @@ +"""In-process background jobs for long-running ingest requests.""" + +from __future__ import annotations + +import json +import uuid +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, replace +from datetime import datetime, timezone +from threading import BoundedSemaphore, RLock +from typing import Callable + +from jiuwen_memory.common.errors import ( + BackendError, + ConflictError, + NotFoundError, + ValidationError, +) +from jiuwen_memory.common.log import get_logger +from jiuwen_memory.common.type_def import MemoryUnit, Scope +from jiuwen_memory.storage.kv import KVStore + +logger = get_logger(__name__) + +INGEST_JOB_PREFIX = "ing_" +_JOB_KEY_PREFIX = "/ingest/jobs/" +_PAYLOAD_KEY_PREFIX = "/ingest/payloads/" +IngestTask = Callable[[], list[MemoryUnit]] + + +@dataclass(frozen=True) +class IngestJob: + id: str + payload_id: str + source_ref: str + scope: Scope + status: str + created_at: datetime + updated_at: datetime + unit_ids: tuple[str, ...] = () + error: str = "" + + +@dataclass(frozen=True) +class IngestSubmission: + job: IngestJob + reused: bool + + +@dataclass(frozen=True) +class _PayloadKey: + org: str + space: str + user: str + agent: str + session: str + payload_id: str + + +class IngestJobController: + """Own queueing, status persistence and payload idempotency for ingest.""" + + def __init__( + self, + *, + max_workers: int = 1, + max_pending_jobs: int = 2, + kv: KVStore | None = None, + ) -> None: + if max_workers <= 0: + raise ValidationError("max_workers must be greater than zero") + if max_pending_jobs < 0: + raise ValidationError("max_pending_jobs must be non-negative") + self._kv = kv + self._executor = ThreadPoolExecutor( + max_workers=max_workers, + thread_name_prefix="memory-ingest", + ) + self._capacity = BoundedSemaphore(max_workers + max_pending_jobs) + self._jobs: dict[str, IngestJob] = {} + self._job_id_by_payload: dict[_PayloadKey, str] = {} + self._lock = RLock() + self._closed = False + + def submit( + self, + *, + payload_id: str, + source_ref: str, + scope: Scope, + task: IngestTask, + ) -> IngestSubmission: + key = _payload_key(scope, payload_id) + now = datetime.now(timezone.utc) + job = IngestJob( + id=f"{INGEST_JOB_PREFIX}{uuid.uuid4().hex}", + payload_id=payload_id, + source_ref=source_ref, + scope=scope, + status="pending", + created_at=now, + updated_at=now, + ) + previous_job_id: str | None = None + with self._lock: + if self._closed: + raise BackendError("ingest job controller is closed") + existing = self._find_existing(scope, payload_id, key) + if existing is not None: + if existing.source_ref != source_ref: + raise ConflictError( + "ingest_payload", + payload_id, + "the same payload_id cannot point to a different source", + ) + if existing.status in {"pending", "running", "succeeded"}: + return IngestSubmission(existing, reused=True) + previous_job_id = existing.id + if not self._capacity.acquire(blocking=False): + raise BackendError("ingest job queue is full") + self._jobs[job.id] = job + self._job_id_by_payload[key] = job.id + try: + self._persist(job) + except Exception: + self._capacity.release() + self._jobs.pop(job.id, None) + if previous_job_id is None: + self._job_id_by_payload.pop(key, None) + else: + self._job_id_by_payload[key] = previous_job_id + raise + try: + future = self._executor.submit(self._run, job.id, task) + future.add_done_callback(lambda _future: self._capacity.release()) + except RuntimeError as exc: + self._capacity.release() + self._update(job.id, status="failed", error=str(exc)) + raise BackendError(f"failed to submit ingest job: {exc}") from exc + return IngestSubmission(job, reused=False) + + def status(self, job_id: str, *, scope: Scope) -> IngestJob: + with self._lock: + job = self._jobs.get(job_id) or self._load(scope, job_id) + if job is not None: + self._jobs[job.id] = job + self._job_id_by_payload[_payload_key(scope, job.payload_id)] = job.id + if job is None or job.scope != scope: + raise NotFoundError("ingest_job", job_id) + return job + + def close(self, *, wait: bool = True) -> None: + with self._lock: + if self._closed: + return + self._closed = True + self._executor.shutdown(wait=wait, cancel_futures=False) + + def _find_existing( + self, + scope: Scope, + payload_id: str, + key: _PayloadKey, + ) -> IngestJob | None: + job_id = self._job_id_by_payload.get(key) + existing = self._jobs.get(job_id or "") + if existing is None and self._kv is not None: + mapping_key = _payload_storage_key(payload_id) + if self._kv.exists(scope, mapping_key): + job_id = self._kv.get(scope, mapping_key).decode("utf-8") + existing = self._load(scope, job_id) + if existing is not None: + self._jobs[existing.id] = existing + self._job_id_by_payload[key] = existing.id + return existing + + def _run(self, job_id: str, task: IngestTask) -> None: + self._update(job_id, status="running") + try: + units = task() + except Exception as exc: + logger.exception("Ingest job failed: job_id=%s", job_id) + self._update(job_id, status="failed", error=str(exc)) + return + self._update( + job_id, + status="succeeded", + unit_ids=tuple(unit.id for unit in units), + ) + + def _update( + self, + job_id: str, + *, + status: str, + unit_ids: tuple[str, ...] = (), + error: str = "", + ) -> None: + with self._lock: + current = self._jobs[job_id] + updated = replace( + current, + status=status, + updated_at=datetime.now(timezone.utc), + unit_ids=unit_ids, + error=error, + ) + self._jobs[job_id] = updated + self._persist(updated) + + def _persist(self, job: IngestJob) -> None: + if self._kv is None: + return + value = json.dumps( + { + "id": job.id, + "payload_id": job.payload_id, + "source_ref": job.source_ref, + "status": job.status, + "created_at": job.created_at.isoformat(), + "updated_at": job.updated_at.isoformat(), + "unit_ids": list(job.unit_ids), + "error": job.error, + }, + ensure_ascii=False, + separators=(",", ":"), + ).encode("utf-8") + job_key = _job_storage_key(job.id) + if self._kv.exists(job.scope, job_key): + self._kv.update(job.scope, job_key, value) + else: + self._kv.insert(job.scope, job_key, value) + payload_key = _payload_storage_key(job.payload_id) + payload_value = job.id.encode("utf-8") + if self._kv.exists(job.scope, payload_key): + self._kv.update(job.scope, payload_key, payload_value) + else: + self._kv.insert(job.scope, payload_key, payload_value) + + def _load(self, scope: Scope, job_id: str) -> IngestJob | None: + if self._kv is None or not self._kv.exists(scope, _job_storage_key(job_id)): + return None + try: + data = json.loads( + self._kv.get(scope, _job_storage_key(job_id)).decode("utf-8") + ) + job = IngestJob( + id=str(data["id"]), + payload_id=str(data["payload_id"]), + source_ref=str(data["source_ref"]), + scope=scope, + status=str(data["status"]), + created_at=datetime.fromisoformat(str(data["created_at"])), + updated_at=datetime.fromisoformat(str(data["updated_at"])), + unit_ids=tuple(str(item) for item in data.get("unit_ids", [])), + error=str(data.get("error", "")), + ) + except (KeyError, TypeError, ValueError) as exc: + raise BackendError(f"invalid persisted ingest job {job_id!r}") from exc + if job.status in {"pending", "running"}: + job = replace( + job, + status="failed", + updated_at=datetime.now(timezone.utc), + error="ingest job was interrupted by server restart", + ) + self._persist(job) + return job + + +def _payload_key(scope: Scope, payload_id: str) -> _PayloadKey: + return _PayloadKey( + org=scope.org, + space=scope.space, + user=scope.user, + agent=scope.agent, + session=scope.session, + payload_id=payload_id, + ) + + +def _job_storage_key(job_id: str) -> str: + return f"{_JOB_KEY_PREFIX}{job_id}" + + +def _payload_storage_key(payload_id: str) -> str: + return f"{_PAYLOAD_KEY_PREFIX}{payload_id}" diff --git a/jiuwen_memory/retrieval/retriever_impl/__init__.py b/jiuwen_memory/retrieval/retriever_impl/__init__.py index b934db93..279e130a 100644 --- a/jiuwen_memory/retrieval/retriever_impl/__init__.py +++ b/jiuwen_memory/retrieval/retriever_impl/__init__.py @@ -8,5 +8,6 @@ from jiuwen_memory.retrieval.retriever import RetrieverProducer import_module(".pipeline_retriever", __name__) +import_module(".multimodal_retriever", __name__) __all__ = ["RetrieverProducer"] diff --git a/jiuwen_memory/retrieval/retriever_impl/multimodal_retriever.py b/jiuwen_memory/retrieval/retriever_impl/multimodal_retriever.py new file mode 100644 index 00000000..2cdf01d8 --- /dev/null +++ b/jiuwen_memory/retrieval/retriever_impl/multimodal_retriever.py @@ -0,0 +1,203 @@ +"""Parallel native, CLM and ELM retrieval for multimodal memory.""" + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import replace + +from jiuwen_memory.common.errors import ValidationError +from jiuwen_memory.common.log import get_logger +from jiuwen_memory.common.type_def import ( + MEMORY_KEY_PREFIX, + FilterClause, + FilterOp, + LifecycleState, + Scope, + and_merge, +) +from jiuwen_memory.common.type_def.memory_codec import loads +from jiuwen_memory.retrieval.base import RetrievalOperatorType +from jiuwen_memory.retrieval.retriever import Retriever, RetrieverProducer +from jiuwen_memory.retrieval.types import ( + RetrievalQuery, + RetrievalResult, + RetrievedItem, + TrajectoryStep, +) +from jiuwen_memory.storage.kv import KvProducer, KVStore + +logger = get_logger(__name__) + + +class MultimodalRetriever(Retriever): + """Compose the native retriever with independent CLM and ELM branches.""" + + def __init__( + self, + base_retriever: Retriever, + kv: KVStore, + *, + clip_top_k: int = 10, + event_top_k: int = 10, + rrf_k: int = 60, + ) -> None: + if clip_top_k <= 0 or event_top_k <= 0: + raise ValidationError("clip_top_k and event_top_k must be greater than zero") + if rrf_k <= 0: + raise ValidationError("rrf_k must be greater than zero") + self._base = base_retriever + self._kv = kv + self._clip_top_k = clip_top_k + self._event_top_k = event_top_k + self._rrf_k = rrf_k + + def operator_type(self) -> RetrievalOperatorType: + return RetrievalOperatorType.RETRIEVER + + def health(self) -> None: + self._base.health() + self._kv.health() + + def retrieve(self, scope: Scope, query: RetrievalQuery) -> RetrievalResult: + if not self._has_multimodal_memory(scope, query.include_archived): + return self._base.retrieve(scope, query) + + queries = { + "native": _with_filters( + query, + FilterClause("source", FilterOp.NE, "video"), + top_k=query.top_k, + ), + "multimodal_clip": _with_filters( + query, + FilterClause("modal_type", FilterOp.EQ, "multimodal"), + FilterClause("memory_level", FilterOp.EQ, "clm"), + top_k=self._clip_top_k, + ), + "multimodal_event": _with_filters( + query, + FilterClause("modal_type", FilterOp.EQ, "multimodal"), + FilterClause("memory_level", FilterOp.EQ, "elm"), + top_k=self._event_top_k, + ), + } + results: dict[str, RetrievalResult] = {} + degraded: dict[str, str] = {} + with ThreadPoolExecutor(max_workers=len(queries)) as executor: + futures = { + executor.submit(self._base.retrieve, scope, branch_query): branch + for branch, branch_query in queries.items() + } + for future in as_completed(futures): + branch = futures[future] + try: + results[branch] = future.result() + except Exception as exc: + logger.warning( + "MultimodalRetriever branch %s failed: %s", + branch, + exc, + ) + results[branch] = RetrievalResult() + degraded[branch] = type(exc).__name__ + + branch_order = ("native", "multimodal_clip", "multimodal_event") + items = _rrf_merge( + [results[branch].items for branch in branch_order], + top_k=query.top_k, + rrf_k=self._rrf_k, + ) + trajectory: list[TrajectoryStep] = [] + if query.with_trajectory: + for branch in branch_order: + trajectory.extend(_branch_trajectory(results[branch].trajectory, branch)) + if branch in degraded: + trajectory.append( + TrajectoryStep( + stage="recall", + candidate_count=0, + detail={ + "branch": branch, + "degraded": degraded[branch], + }, + ) + ) + trajectory.append( + TrajectoryStep( + stage="fuse", + candidate_count=len(items), + detail={ + "strategy": "rrf", + "branches": ",".join(branch_order), + "rrf_k": str(self._rrf_k), + }, + ) + ) + return RetrievalResult(items=items, trajectory=trajectory) + + def _has_multimodal_memory(self, scope: Scope, include_archived: bool) -> bool: + for _, raw in self._kv.scan(scope, MEMORY_KEY_PREFIX): + unit = loads(raw) + if unit is None or unit.metadata.get("modal_type") != "multimodal": + continue + if unit.metadata.get("memory_level") not in {"clm", "elm"}: + continue + if unit.lifecycle == LifecycleState.ACTIVE: + return True + if include_archived and unit.lifecycle == LifecycleState.ARCHIVED: + return True + return False + + +def _with_filters( + query: RetrievalQuery, + *filters: FilterClause, + top_k: int, +) -> RetrievalQuery: + return replace( + query, + filters=and_merge(query.filters, list(filters)), + top_k=top_k, + ) + + +def _rrf_merge( + branches: list[list[RetrievedItem]], + *, + top_k: int, + rrf_k: int, +) -> list[RetrievedItem]: + by_id: dict[str, RetrievedItem] = {} + scores: dict[str, float] = {} + for branch in branches: + for rank, item in enumerate(branch, start=1): + by_id.setdefault(item.unit_id, item) + scores[item.unit_id] = scores.get(item.unit_id, 0.0) + 1.0 / (rrf_k + rank) + ranked_ids = sorted(scores, key=scores.get, reverse=True)[:top_k] + merged: list[RetrievedItem] = [] + for unit_id in ranked_ids: + item = by_id.get(unit_id) + if item is not None: + merged.append(replace(item, score=scores.get(unit_id, 0.0))) + return merged + + +def _branch_trajectory( + steps: list[TrajectoryStep], + branch: str, +) -> list[TrajectoryStep]: + return [ + replace(step, detail={**step.detail, "branch": branch}) + for step in steps + ] + + +@RetrieverProducer.register("multimodal") +def _build(config): + return MultimodalRetriever( + RetrieverProducer.dep(config, "base_retriever", default="pipeline"), + KvProducer.dep(config, "kv_store", default="memory"), + clip_top_k=int(config.get("clip_top_k", 10)), + event_top_k=int(config.get("event_top_k", 10)), + rrf_k=int(config.get("rrf_k", 60)), + ) diff --git a/pyproject.toml b/pyproject.toml index 1df12b51..c9974f57 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,6 +30,11 @@ embed = [ # transformers 5.x 移除了 FlagReranker 仍调用的 prepare_for_model，会令 bge reranker 崩；固定在 4.x。 "transformers>=4.39,<5", ] +multimodal = [ + "torch>=2.0", + "transformers>=4.39,<5", + "soundfile>=0.12", +] # 接真后端部署所需的客户端库（Milvus / Elasticsearch / Redis / PostgreSQL）+ YAML 配置解析。 deploy = [ "pymilvus>=2.4", diff --git a/tests/unit/construction/test_video_memory_extractor.py b/tests/unit/construction/test_video_memory_extractor.py new file mode 100644 index 00000000..1af3b284 --- /dev/null +++ b/tests/unit/construction/test_video_memory_extractor.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +import json + +from jiuwen_memory.common.type_def import MemoryUnit, Modality, Scope, Segment +from jiuwen_memory.common.type_def.memory_codec import dumps, loads +from jiuwen_memory.construction.extractor_impl.video_memory_extractor import ( + VideoMemoryExtractor, +) + + +def test_video_memory_times_are_float_metadata() -> None: + source = MemoryUnit( + id="source-1", + scope=Scope(user="user-1"), + segments=[ + Segment( + content=json.dumps( + { + "payload_id": "video-1", + "clips": [ + { + "id": "clip-1", + "start_seconds": 1.9, + "end_seconds": 30.8, + "visual_summary": "A person enters the room.", + } + ], + "events": [ + { + "id": "event-1", + "start_seconds": 1.9, + "end_seconds": 30.8, + "topic": "Room entry", + "clip_ids": ["clip-1"], + } + ], + } + ), + source=Modality.VIDEO, + ) + ], + source_ref="video-1", + ) + + units = VideoMemoryExtractor().extract([source]) + + assert len(units) == 2 + assert all(unit.metadata["start_seconds"] == 1.9 for unit in units) + assert all(unit.metadata["end_seconds"] == 30.8 for unit in units) + assert all(isinstance(unit.metadata["start_seconds"], float) for unit in units) + assert all(isinstance(unit.metadata["end_seconds"], float) for unit in units) + + restored = loads(dumps(units[0])) + assert restored is not None + assert isinstance(restored.metadata["start_seconds"], float) + assert isinstance(restored.metadata["end_seconds"], float) diff --git a/tests/unit/control/test_ingest_job.py b/tests/unit/control/test_ingest_job.py new file mode 100644 index 00000000..81390f5b --- /dev/null +++ b/tests/unit/control/test_ingest_job.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +import time +from threading import Event + +import pytest + +from jiuwen_memory.common.errors import ConflictError +from jiuwen_memory.common.type_def import MemoryUnit, Modality, Scope, Segment +from jiuwen_memory.control.ingest_job import IngestJobController + + +def test_ingest_job_runs_in_background_and_reuses_payload() -> None: + scope = Scope(user="user-1") + started = Event() + release = Event() + controller = IngestJobController(max_workers=1, max_pending_jobs=1) + + def task() -> list[MemoryUnit]: + started.set() + assert release.wait(2) + return [ + MemoryUnit( + id="unit-1", + scope=scope, + segments=[Segment(content="video memory", source=Modality.VIDEO)], + ) + ] + + try: + first = controller.submit( + payload_id="video-1", + source_ref="file:///data/demo.mp4", + scope=scope, + task=task, + ) + assert first.job.id.startswith("ing_") + assert started.wait(1) + + duplicate = controller.submit( + payload_id="video-1", + source_ref="file:///data/demo.mp4", + scope=scope, + task=task, + ) + assert duplicate.reused is True + assert duplicate.job.id == first.job.id + + with pytest.raises(ConflictError): + controller.submit( + payload_id="video-1", + source_ref="file:///data/other.mp4", + scope=scope, + task=task, + ) + + release.set() + deadline = time.monotonic() + 2 + while time.monotonic() < deadline: + job = controller.status(first.job.id, scope=scope) + if job.status == "succeeded": + break + time.sleep(0.01) + + assert job.status == "succeeded" + assert job.unit_ids == ("unit-1",) + finally: + release.set() + controller.close() + + +def test_ingest_payload_idempotency_isolated_by_space() -> None: + controller = IngestJobController(max_workers=1, max_pending_jobs=1) + try: + first = controller.submit( + payload_id="video-1", + source_ref="file:///data/a.mp4", + scope=Scope(org="org-1", space="space-a", user="user-1"), + task=lambda: [], + ) + second = controller.submit( + payload_id="video-1", + source_ref="file:///data/b.mp4", + scope=Scope(org="org-1", space="space-b", user="user-1"), + task=lambda: [], + ) + assert second.reused is False + assert second.job.id != first.job.id + finally: + controller.close() diff --git a/tests/unit/multimodal/test_multimodal_adapter.py b/tests/unit/multimodal/test_multimodal_adapter.py new file mode 100644 index 00000000..6cdf07c1 --- /dev/null +++ b/tests/unit/multimodal/test_multimodal_adapter.py @@ -0,0 +1,393 @@ +from __future__ import annotations + +import importlib +import json +import sys +import time +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import yaml + +from jiuwen_memory.api import build_kernel +from jiuwen_memory.common.normalizer.normalizer_impl.passthrough_normalizer import ( + PassthroughNormalizer, +) +from jiuwen_memory.common.normalizer.normalizer_impl.routing_normalizer import ( + RoutingNormalizer, +) +from jiuwen_memory.common.normalizer.normalizer_impl.video_normalizer import VideoNormalizer +from jiuwen_memory.common.type_def import ( + Context, + FilterClause, + FilterOp, + MemoryUnit, + Modality, + RawPayload, + Scope, + Segment, + iter_clauses, + memory_key, +) +from jiuwen_memory.common.type_def.memory_codec import dumps +from jiuwen_memory.config import Config +from jiuwen_memory.control.ingest_job import IngestJobController +from jiuwen_memory.retrieval.base import RetrievalOperatorType +from jiuwen_memory.retrieval.retriever import Retriever +from jiuwen_memory.retrieval.retriever_impl.multimodal_retriever import ( + MultimodalRetriever, +) +from jiuwen_memory.retrieval.types import RetrievalQuery, RetrievalResult +from jiuwen_memory.storage.kv_impl.in_memory_kv_store import InMemoryKVStore + +_BOOTSTRAP_CORE = Path(__file__).parents[3] / "bootstrap" / "core" +if str(_BOOTSTRAP_CORE) not in sys.path: + sys.path.append(str(_BOOTSTRAP_CORE)) +handler = importlib.import_module("handler") + + +def _video_memory_output(_payload: RawPayload): + return ( + [ + { + "id": "clip-1", + "time_range": [1.25, 30.75], + "visual_summary": "A presenter opens a deployment diagram.", + "detailed_caption": "A topology slide is visible.", + "ASR": "Deploy the embedding service first.", + "environment": "Meeting room", + } + ], + [ + { + "task_id": "event-1", + "topic": "Deployment plan", + "time_span": [1.25, 30.75], + "narrative_summary": "The team validates before production.", + "semantic_inference": "Testing is a release gate.", + "child_clip_ids": ["clip-1"], + } + ], + ) + + +def test_routing_normalizer_keeps_text_and_routes_video() -> None: + scope = Scope(user="user-1") + normalizer = RoutingNormalizer( + PassthroughNormalizer(), + {Modality.VIDEO: VideoNormalizer(backend=_video_memory_output)}, + ) + + text = normalizer.normalize( + RawPayload( + id="text-1", + scope=scope, + modality=Modality.TEXT, + data=b"plain text", + ) + ) + video = json.loads( + normalizer.normalize( + RawPayload( + id="video-1", + scope=scope, + modality=Modality.VIDEO, + uri="file:///data/demo.mp4", + ) + ) + ) + + assert text == "plain text" + assert video["payload_id"] == "video-1" + assert video["asset_uri"] == "file:///data/demo.mp4" + assert video["clips"][0]["start_seconds"] == 1.25 + assert video["events"][0]["clip_ids"] == ["clip-1"] + + +def test_video_normalizer_uses_configured_temp_root(tmp_path, monkeypatch) -> None: + video_path = tmp_path / "demo.mp4" + video_path.write_bytes(b"video") + temp_root = tmp_path / "video-work" + observed: dict[str, Path] = {} + + def fake_run_pipeline(self, source_path: Path, run_root: Path): + del self + observed["source_path"] = source_path + observed["run_root"] = run_root + assert run_root.parent == temp_root + assert run_root.is_dir() + return [], [] + + monkeypatch.setattr(VideoNormalizer, "_run_pipeline", fake_run_pipeline) + normalizer = VideoNormalizer.from_config({"temp_root": str(temp_root)}) + normalizer.normalize( + RawPayload( + id="video-1", + scope=Scope(user="user-1"), + modality=Modality.VIDEO, + uri=video_path.as_uri(), + ) + ) + + assert observed["source_path"] == video_path + assert not observed["run_root"].exists() + + +def test_video_normalizer_passes_yaml_model_settings(tmp_path, monkeypatch) -> None: + video_path = tmp_path / "demo.mp4" + video_path.write_bytes(b"video") + observed: dict[str, object] = {} + + def fake_run_pipeline( + source_path: Path, + run_root: Path, + config: Any, + ): + observed.update( + source_path=source_path, + run_root=run_root, + whisper_model_dir=config.whisper_model_dir, + whisper_device=config.whisper_device, + whisper_language=config.whisper_language, + whisper_batch_size=config.whisper_batch_size, + vllm_base_url=config.vllm_base_url, + vllm_api_key=config.vllm_api_key, + llm_model=config.llm_model, + ) + return {"short_term": [], "medium_term": []} + + monkeypatch.setitem( + sys.modules, + "jiuwen_memory.common.normalizer.normalizer_impl.video_pipeline", + SimpleNamespace( + VideoPipelineConfig=lambda **kwargs: SimpleNamespace(**kwargs), + run_video_memory_pipeline_off=fake_run_pipeline, + ), + ) + normalizer = VideoNormalizer.from_config( + { + "whisper_model_dir": "/models/whisper", + "whisper_device": "npu:0", + "whisper_language": "zh", + "whisper_batch_size": 2, + "vllm_base_url": "http://127.0.0.1:8000/v1", + "vllm_api_key": "test-key", + "llm_model": "qwen-vl-test", + "temp_root": str(tmp_path / "video-work"), + } + ) + + normalizer.normalize( + RawPayload( + id="video-1", + scope=Scope(user="user-1"), + modality=Modality.VIDEO, + uri=video_path.as_uri(), + ) + ) + + assert observed["whisper_model_dir"] == "/models/whisper" + assert observed["whisper_device"] == "npu:0" + assert observed["whisper_language"] == "zh" + assert observed["whisper_batch_size"] == 2 + assert observed["vllm_base_url"] == "http://127.0.0.1:8000/v1" + assert observed["vllm_api_key"] == "test-key" + assert observed["llm_model"] == "qwen-vl-test" + + +class _RecordingRetriever(Retriever): + def __init__(self) -> None: + self.queries: list[RetrievalQuery] = [] + + def operator_type(self) -> RetrievalOperatorType: + return RetrievalOperatorType.RETRIEVER + + def health(self) -> None: + return None + + def retrieve(self, scope: Scope, query: RetrievalQuery) -> RetrievalResult: + del scope + self.queries.append(query) + return RetrievalResult() + + +def test_multimodal_retriever_supports_filter_expr_and_space_isolation() -> None: + scope = Scope(org="org-1", space="space-a", user="user-1") + other_space = Scope(org="org-1", space="space-b", user="user-1") + unit = MemoryUnit( + id="clip-unit", + scope=scope, + segments=[Segment(content="video clip", source=Modality.VIDEO)], + metadata={"modal_type": "multimodal", "memory_level": "clm"}, + ) + kv = InMemoryKVStore() + kv.insert(scope, memory_key(unit.id), dumps(unit)) + base = _RecordingRetriever() + retriever = MultimodalRetriever(base, kv) + + retriever.retrieve( + scope, + RetrievalQuery( + text="video", + filters=FilterClause("tags", FilterOp.CONTAINS, "keep"), + ), + ) + assert len(base.queries) == 3 + fields = [{clause.field for clause in iter_clauses(query.filters)} for query in base.queries] + assert {"tags", "source"} in fields + assert {"tags", "modal_type", "memory_level"} in fields + + base.queries.clear() + retriever.retrieve(other_space, RetrievalQuery(text="video")) + assert len(base.queries) == 1 + + +def test_multimodal_config_add_and_search_end_to_end(tmp_path, monkeypatch) -> None: + settings = yaml.safe_load( + (Path(__file__).parents[3] / "examples" / "config_multimodal.yml").read_text( + encoding="utf-8" + ) + )["memory_api"] + settings["normalizer"]["default"]["params"]["routes"]["video"]["params"][ + "temp_root" + ] = str(tmp_path / "video-work") + + def fake_run_pipeline(self, video_path: Path, run_root: Path): + del self, video_path, run_root + return _video_memory_output( + RawPayload(id="video-1", scope=Scope(user="user-1")) + ) + + monkeypatch.setattr(VideoNormalizer, "_run_pipeline", fake_run_pipeline) + video_path = tmp_path / "demo.mp4" + video_path.write_bytes(b"video") + kernel = build_kernel(config=Config.from_dict(settings)) + scope = Scope(org="org-1", user="user-1") + units = kernel.api.add( + video_path.as_uri(), + scope, + Modality.VIDEO, + identity=scope, + assets=[video_path.as_uri()], + metadata={"infer": "true", "pipeline": "video", "payload_id": "video-1"}, + ) + + assert {unit.metadata["memory_level"] for unit in units} == {"clm", "elm"} + result = kernel.api.search( + "deployment", + Context(scope), + identity=scope, + top_k=10, + with_trajectory=True, + ) + assert result.items + assert {item.unit_id for item in result.items}.issubset({unit.id for unit in units}) + assert { + step.detail.get("branch") + for step in result.trajectory + if step.detail.get("branch") + } >= {"native", "multimodal_clip", "multimodal_event"} + + +def test_video_add_and_prefixed_job_status_share_handler_route( + tmp_path, monkeypatch +) -> None: + settings = yaml.safe_load( + (Path(__file__).parents[3] / "examples" / "config_multimodal.yml").read_text( + encoding="utf-8" + ) + )["memory_api"] + settings["normalizer"]["default"]["params"]["routes"]["video"]["params"][ + "temp_root" + ] = str(tmp_path / "video-work") + monkeypatch.setattr( + VideoNormalizer, + "_run_pipeline", + lambda self, video_path, run_root: _video_memory_output( + RawPayload(id="video-1", scope=Scope(user="user-1")) + ), + ) + video_path = tmp_path / "demo.mp4" + video_path.write_bytes(b"video") + kernel = build_kernel(config=Config.from_dict(settings)) + controller = IngestJobController(max_workers=1, max_pending_jobs=1, kv=kernel.kv) + srv = type( + "ServerStub", + (), + {"api": kernel.api, "ingest_jobs": controller}, + )() + try: + status, submitted = handler.dispatch( + srv, + "add", + { + "tenant_id": "org-1", + "scope": "user-1", + "payload_id": "video-1", + "modality": "video", + "uri": video_path.as_uri(), + }, + ) + assert status == 200 + assert submitted["accepted"] is True + assert submitted["job_id"].startswith("ing_") + assert submitted["status"] in {"pending", "running", "succeeded"} + assert submitted["reused"] is False + + deadline = time.monotonic() + 2 + while time.monotonic() < deadline: + status, result = handler.dispatch( + srv, + "job", + { + "tenant_id": "org-1", + "scope": "user-1", + "job_id": submitted["job_id"], + }, + ) + if result.get("status") == "succeeded": + break + time.sleep(0.01) + assert status == 200 + assert result["status"] == "succeeded" + assert result["item_ids"] + assert {item["metadata"]["video_id"] for item in result["items"]} == { + "video-1" + } + + status, reused = handler.dispatch( + srv, + "add", + { + "tenant_id": "org-1", + "scope": "user-1", + "payload_id": "video-1", + "modality": "video", + "uri": video_path.as_uri(), + }, + ) + assert status == 200 + assert reused["job_id"] == submitted["job_id"] + assert reused["status"] == "succeeded" + assert reused["reused"] is True + finally: + controller.close() + + +def test_video_add_requires_uri() -> None: + srv = type("ServerStub", (), {})() + status, body = handler.dispatch( + srv, + "add", + { + "tenant_id": "org-1", + "scope": "user-1", + "modality": "video", + "content": "file:///data/demo.mp4", + }, + ) + + assert status == 400 + assert body["error"] == "ValidationError" + assert "uri" in body["message"] diff --git a/uv.lock b/uv.lock index 788e9097..72856c91 100644 --- a/uv.lock +++ b/uv.lock @@ -1363,6 +1363,11 @@ mcp = [ { name = "mcp" }, { name = "pyyaml" }, ] +multimodal = [ + { name = "soundfile" }, + { name = "torch" }, + { name = "transformers" }, +] nlp = [ { name = "hanlp" }, { name = "spacy" }, @@ -1394,11 +1399,14 @@ requires-dist = [ { name = "pyyaml", marker = "extra == 'mcp'", specifier = ">=6" }, { name = "redis", marker = "extra == 'deploy'", specifier = ">=5" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.8" }, + { name = "soundfile", marker = "extra == 'multimodal'", specifier = ">=0.12" }, { name = "spacy", marker = "extra == 'nlp'", specifier = ">=3.7" }, { name = "torch", marker = "extra == 'embed'", specifier = ">=2.0" }, + { name = "torch", marker = "extra == 'multimodal'", specifier = ">=2.0" }, { name = "transformers", marker = "extra == 'embed'", specifier = ">=4.39,<5" }, + { name = "transformers", marker = "extra == 'multimodal'", specifier = ">=4.39,<5" }, ] -provides-extras = ["dev", "nlp", "embed", "deploy", "mcp"] +provides-extras = ["dev", "nlp", "embed", "multimodal", "deploy", "mcp"] [package.metadata.requires-dev] dev = [ @@ -3752,6 +3760,28 @@ wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl", hash = "sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2" }, ] +[[package]] +name = "soundfile" +version = "0.14.0" +source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } +dependencies = [ + { name = "cffi" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://mirrors.aliyun.com/pypi/simple/" }, marker = "python_full_version < '3.12'" }, + { name = "numpy", version = "2.5.0", source = { registry = "https://mirrors.aliyun.com/pypi/simple/" }, marker = "python_full_version >= '3.12'" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/d2/db/949331952a6fb1c5b12e9de80fd08747966c2039d1a61db4764fbd3981c2/soundfile-0.14.0.tar.gz", hash = "sha256:ba1c1a2d618bca5c406647c83b89f07cc8810fa506a50622a6993ba130c1de11" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/b1/d1/5e338af9ca6ed0786cd5bb03f6d60de1c325728c1189014f3b59aae7403c/soundfile-0.14.0-py2.py3-none-any.whl", hash = "sha256:8ba81ae3a89fd5ab3bef8a8eb481fbbe794e806309675a89b4df48b8d31908a8" }, + { url = "https://mirrors.aliyun.com/pypi/packages/7e/72/c6b21e58d3113596e7e8de0a08d6f1d95173492cfbca0a4db14148cbba2a/soundfile-0.14.0-py2.py3-none-macosx_10_9_x86_64.whl", hash = "sha256:19be05428da76ed61a4cad29b8e4bcf43a3e5c100089d2ec81dc961eed1b0dd4" }, + { url = "https://mirrors.aliyun.com/pypi/packages/63/7a/dfdd6f8c748988427119f75eb860a3cedd858d1aea1fe28f39ad8559ef22/soundfile-0.14.0-py2.py3-none-macosx_11_0_arm64.whl", hash = "sha256:d828d35a059626da52f1415b5faee610aeab393319cb3fc4a9aef47b619fc14c" }, + { url = "https://mirrors.aliyun.com/pypi/packages/4a/f8/fc39fad6f879633461d27394cd1ddaf1f769ffa0597dca35872f51b16461/soundfile-0.14.0-py2.py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:e85724a90bc99a6e8062c0b4ddf725f53b2a3b70afd4da875e9d2cfc4e92f377" }, + { url = "https://mirrors.aliyun.com/pypi/packages/7b/a2/70fd4432b924684c372df8b0a45708c36c057ef3596c9eb53e0a806b980b/soundfile-0.14.0-py2.py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:1e38bac1853412871318e82a1ba69a8be677619b56025bbfcccdb41b6cafe82d" }, + { url = "https://mirrors.aliyun.com/pypi/packages/d9/34/c9e80783d83eab739a9531fdee03675d53e0bf1b2ccb4bb3af5844675046/soundfile-0.14.0-py2.py3-none-win32.whl", hash = "sha256:0a6ae43c50c71b4e020cc55382925cb89451c1ed1a0c3d0f5d802da269226849" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ed/97/b39c18ac1df45e755ca22b8b00e872929da5d107998a207a5e4ac831bfda/soundfile-0.14.0-py2.py3-none-win_amd64.whl", hash = "sha256:299491d3499460fb1b74bb4bd78b57ffc2d243a5fafa7b6ec1b264875c78453e" }, + { url = "https://mirrors.aliyun.com/pypi/packages/f4/83/55c65e61cf457805ce2ec157c1c6ae17715d0851aa2374422de0538838ca/soundfile-0.14.0-py2.py3-none-win_arm64.whl", hash = "sha256:e090704718e124e7c844695236f1fce8d18a5e761eaf7c82dfcd124620805f98" }, +] + [[package]] name = "spacy" version = "3.8.14"