Reuse identical sentiment embedding inputs

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)]

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
             ]

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__