fix(session): batch long-term extraction of a large commit without an auto-commit policy

openviking/session/extraction_batch.py

@@ -8,7 +8,10 @@
 from typing import Any, Iterable, List, Optional
 
 from openviking.message import Message
-from openviking.session.auto_commit_policy import AutoCommitPolicy
+from openviking.session.auto_commit_policy import (
+    DEFAULT_PENDING_TOKEN_THRESHOLD,
+    AutoCommitPolicy,
+)
 from openviking.session.retention import UserTurn, build_turns
 
 
@@ -152,10 +155,25 @@ def flush() -> None:
     return batches
 
 
+def long_term_extraction_batch_limits(
+    limits: ExtractionBatchLimits, messages: Iterable[Message]
+) -> ExtractionBatchLimits:
+    """Policy limits, or the default auto-commit token threshold for a large commit
+    without a policy (manual commits, imports), which would otherwise go to the
+    model in one call however long it is."""
+    messages = list(messages)
+    if limits.enabled or (
+        estimate_extraction_message_tokens(messages) <= DEFAULT_PENDING_TOKEN_THRESHOLD
+    ):
+        return limits
+    return ExtractionBatchLimits(max_message_tokens=DEFAULT_PENDING_TOKEN_THRESHOLD)
+
+
 __all__ = [
     "ExtractionBatchLimits",
     "ExtractionMessageBatch",
     "estimate_extraction_message_tokens",
+    "long_term_extraction_batch_limits",
     "plan_extraction_batches",
     "resolve_extraction_batch_limits",
 ]

openviking/session/session.py

@@ -41,6 +41,7 @@
     ExtractionBatchLimits,
     ExtractionMessageBatch,
     estimate_extraction_message_tokens,
+    long_term_extraction_batch_limits,
     plan_extraction_batches,
     resolve_extraction_batch_limits,
 )
@@ -2235,11 +2236,14 @@ async def _run_long_term_memory_extraction(
                                     event_search_tags=event_search_tags,
                                 )
 
-                            if extraction_batch_limits.enabled:
+                            long_term_limits = long_term_extraction_batch_limits(
+                                extraction_batch_limits, long_term_messages
+                            )
+                            if long_term_limits.enabled:
                                 extraction_tasks.append(
                                     self._extract_long_term_memories_with_batching(
                                         messages=long_term_messages,
-                                        limits=extraction_batch_limits,
+                                        limits=long_term_limits,
                                         archive_uri=archive_uri,
                                         extract_batch=_run_long_term_memory_extraction,
                                         record_batch=_run_recorded_memory_step,

tests/unit/session/test_extraction_batch.py

@@ -7,8 +7,11 @@
 import pytest
 
 from openviking.message import Message, TextPart
+from openviking.session.auto_commit_policy import DEFAULT_PENDING_TOKEN_THRESHOLD
 from openviking.session.extraction_batch import (
     ExtractionBatchLimits,
+    estimate_extraction_message_tokens,
+    long_term_extraction_batch_limits,
     plan_extraction_batches,
     resolve_extraction_batch_limits,
 )
@@ -64,6 +67,21 @@ def test_phase2_batches_split_an_oversized_batch_by_token_count():
     assert all(batch.estimated_tokens <= one_message_tokens for batch in batches)
 
 
+def test_long_term_extraction_without_policy_batches_only_a_large_commit():
+    small = [_message("u1"), _message("a1", role="assistant")]
+    large = [_message(f"m{i}", text="word " * 20_000) for i in range(30)]
+    assert estimate_extraction_message_tokens(large) > DEFAULT_PENDING_TOKEN_THRESHOLD
+    policy = ExtractionBatchLimits(max_messages=5)
+
+    assert long_term_extraction_batch_limits(ExtractionBatchLimits(), small).enabled is False
+    assert long_term_extraction_batch_limits(policy, large) == policy
+    fallback = long_term_extraction_batch_limits(ExtractionBatchLimits(), large)
+    assert fallback == ExtractionBatchLimits(max_message_tokens=DEFAULT_PENDING_TOKEN_THRESHOLD)
+    batches = plan_extraction_batches(large, fallback)
+    assert len(batches) > 1
+    assert all(batch.estimated_tokens <= DEFAULT_PENDING_TOKEN_THRESHOLD for batch in batches)
+
+
 def test_phase2_batches_keep_a_turn_together_when_it_fits():
     messages = [
         _message("u1", text="question"),