diff --git a/openviking/session/extraction_batch.py b/openviking/session/extraction_batch.py index 9694e89cf2..b60be79f37 100644 --- a/openviking/session/extraction_batch.py +++ b/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", ] diff --git a/openviking/session/session.py b/openviking/session/session.py index 084115951a..bea9c57345 100644 --- a/openviking/session/session.py +++ b/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, diff --git a/tests/unit/session/test_extraction_batch.py b/tests/unit/session/test_extraction_batch.py index f50e225933..c9ff63ccf0 100644 --- a/tests/unit/session/test_extraction_batch.py +++ b/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"),