diff --git a/src/compass/compass_signals/user_embedding/embedding.py b/src/compass/compass_signals/user_embedding/embedding.py index f21a0f10..8556fd7f 100644 --- a/src/compass/compass_signals/user_embedding/embedding.py +++ b/src/compass/compass_signals/user_embedding/embedding.py @@ -20,6 +20,8 @@ from pydantic import BaseModel, ConfigDict, Field from tokenizers import Tokenizer +MAX_USER_MESSAGE_CHARS = 20_000 + class _NoRetries(int): """Zero retries, including in LiteLLM adapters that use ``value or default``.""" @@ -230,8 +232,8 @@ def generate_local(self, texts: list[str]) -> list[UserEmbedding]: """Encode locally, preserving the per-message character limit.""" return [ UserEmbedding( - data=self.projection.compress(self.embed(text[:20000])), - token_count=self.token_count(text[:20000]), + data=self.projection.compress(self.embed(text[:MAX_USER_MESSAGE_CHARS])), + token_count=self.token_count(text[:MAX_USER_MESSAGE_CHARS]), ) for text in texts ] @@ -240,7 +242,7 @@ def _remote_batches(self, texts: list[str], size: int) -> Iterator[_RemoteBatch] # Prepare only the next batch, on the coordinator thread. for start in range(0, len(texts), size): inputs = [ - self.projection.metadata["encoder"]["prefix"] + text[:20000] + self.projection.metadata["encoder"]["prefix"] + text[:MAX_USER_MESSAGE_CHARS] for text in texts[start : start + size] ] counts = [len(tokens.ids) for tokens in self.remote_tokenizer.encode_batch(inputs)] diff --git a/src/compass/compass_signals/user_sentiment/signal.py b/src/compass/compass_signals/user_sentiment/signal.py index 4f8da6c8..2b7c4822 100644 --- a/src/compass/compass_signals/user_sentiment/signal.py +++ b/src/compass/compass_signals/user_sentiment/signal.py @@ -15,6 +15,7 @@ from compass.compass_signals.compass_signals import CompassSignalResult from compass.compass_signals.issue_detector import IssueDetector from compass.compass_signals.user_embedding.embedding import ( + MAX_USER_MESSAGE_CHARS, LiteLLMEmbeddingConfig, UserEmbeddingGenerator, validate_embedding_dependencies, @@ -115,22 +116,27 @@ def screen_user_messages( embedding_generator: UserEmbeddingGenerator, classifier: ComplaintClassifier, ) -> dict[str, TraceScreeningResult]: - """Classify every message separately; later turns cannot clear an earlier flag.""" - embeddings = iter( - embedding_generator.generate_batch( - [ - text - for messages in user_messages.values() - for text in messages - if text[:20000].strip() - ] + """Reuse identical encoder inputs while preserving every message occurrence.""" + # The encoder prefix is fixed for this generator. Match its character limit, + # but retain whitespace: stripping would change the text sent to the model. + unique_texts = list( + dict.fromkeys( + text[:MAX_USER_MESSAGE_CHARS] + for messages in user_messages.values() + for text in messages + if text[:MAX_USER_MESSAGE_CHARS].strip() ) ) + embeddings = embedding_generator.generate_batch(unique_texts) + classified = { + text: classifier.classify(embedding) + for text, embedding in zip(unique_texts, embeddings, strict=True) + } return { trace_id: TraceScreeningResult( message_results=[ - classifier.classify(next(embeddings)) - if text[:20000].strip() + classified[text[:MAX_USER_MESSAGE_CHARS]].model_copy(deep=True) + if text[:MAX_USER_MESSAGE_CHARS].strip() else ScreeningResult(label="no_user_messages") for text in messages ] diff --git a/tests/test_user_embedding.py b/tests/test_user_embedding.py index 99d5de08..065261b0 100644 --- a/tests/test_user_embedding.py +++ b/tests/test_user_embedding.py @@ -30,7 +30,7 @@ def embed(text): assert generator.generate_batch([]) == [] -def test_remote_batches_preserve_trace_order_and_skip_empty_messages(monkeypatch): +def test_screening_reuses_exact_inputs_preserving_occurrences(monkeypatch): generator = UserEmbeddingGenerator( litellm=LiteLLMEmbeddingConfig(model="openai/qwen", batch_size=4) ) @@ -50,18 +50,29 @@ def embed(**kwargs): monkeypatch.setattr( "compass.compass_signals.user_embedding.embedding._request_embeddings", embed ) - messages = {"a": ["message"] * 8, "b": ["", "last"], "c": []} + long = "x" * 20000 + messages: dict[str, list[str]] = { + "a": ["message"] * 8 + [" spaced ", "spaced", long + "a"], + "b": ["", "last", long + "b", "message"], + "c": [], + "d": [" " * 20000 + "ignored"], + } results = screen_user_messages(messages, generator, ComplaintClassifier(generator.projection)) - assert sorted(map(len, calls)) == [1, 4, 4] - assert all( - text.startswith(generator.projection.metadata["encoder"]["prefix"]) - for call in calls - for text in call - ) - assert len(results["a"].message_results) == 8 + assert sorted(map(len, calls)) == [1, 4] + prefix = generator.projection.metadata["encoder"]["prefix"] + assert {text for call in calls for text in call} == { + prefix + text for text in ["message", " spaced ", "spaced", long, "last"] + } + assert len(results["a"].message_results) == 11 assert results["b"].message_results[0].label == "no_user_messages" assert results["b"].message_results[1].token_count == 1 assert results["c"].message_results == [] + assert results["d"].message_results[0].label == "no_user_messages" + assert results["a"].message_results[-1] == results["b"].message_results[2] + assert results["a"].message_results[0] == results["b"].message_results[-1] + results["a"].message_results[0].scores.clear() + assert results["b"].message_results[-1].scores # Occurrences don't share mutable state. + assert messages["a"][-1] == long + "a" # Evidence remains verbatim. assert "model" not in generator.__dict__