From 69d4ff723e6af634dd17969a647edf1318de2a47 Mon Sep 17 00:00:00 2001 From: Dan Haggerty Date: Mon, 2 Nov 2026 16:08:18 -0500 Subject: [PATCH 1/4] Reuse identical sentiment inputs without dropping message occurrences --- .../user_embedding/embedding.py | 28 ++++++++- .../compass_signals/user_sentiment/signal.py | 44 ++++++++++---- tests/test_user_embedding.py | 60 +++++++++++++++---- 3 files changed, 109 insertions(+), 23 deletions(-) diff --git a/src/compass/compass_signals/user_embedding/embedding.py b/src/compass/compass_signals/user_embedding/embedding.py index f21a0f10..59e6d5ba 100644 --- a/src/compass/compass_signals/user_embedding/embedding.py +++ b/src/compass/compass_signals/user_embedding/embedding.py @@ -8,18 +8,22 @@ import os from collections.abc import Iterator from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait +from contextvars import copy_context from dataclasses import dataclass from functools import cached_property from importlib.resources import files from importlib.util import find_spec from itertools import islice -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast import numpy as np from litellm import EmbeddingResponse +from opentelemetry import trace from pydantic import BaseModel, ConfigDict, Field from tokenizers import Tokenizer +tracer = trace.get_tracer(__name__) + class _NoRetries(int): """Zero retries, including in LiteLLM adapters that use ``value or default``.""" @@ -265,7 +269,12 @@ def generate_remote(self, texts: list[str]) -> list[UserEmbedding]: while True: # Keep only one batch per worker in flight; refill as each finishes. for batch in islice(batches, config.max_concurrency - len(active)): - future = executor.submit(self._remote_batch, batch, config, api_key) + future = cast( + Future[list[UserEmbedding]], + executor.submit( + copy_context().run, self._remote_batch, batch, config, api_key + ), + ) active[future] = batch.index if not active: break @@ -276,6 +285,21 @@ def generate_remote(self, texts: list[str]) -> list[UserEmbedding]: def _remote_batch( self, batch: _RemoteBatch, config: LiteLLMEmbeddingConfig, api_key: str | None + ) -> list[UserEmbedding]: + with tracer.start_as_current_span( + "sentiment.embedding_batch", + attributes={ + "openinference.span.kind": "CHAIN", + "input_offset": batch.index, + "input_count": len(batch.inputs), + "token_count": sum(batch.counts), + "model": config.model, + }, + ): + return self._encode_remote_batch(batch, config, api_key) + + def _encode_remote_batch( + self, batch: _RemoteBatch, config: LiteLLMEmbeddingConfig, api_key: str | None ) -> list[UserEmbedding]: response = _request_embeddings( model=config.model, diff --git a/src/compass/compass_signals/user_sentiment/signal.py b/src/compass/compass_signals/user_sentiment/signal.py index 4f8da6c8..0cfc5175 100644 --- a/src/compass/compass_signals/user_sentiment/signal.py +++ b/src/compass/compass_signals/user_sentiment/signal.py @@ -10,6 +10,7 @@ from nooa import Agent from nooa.unifiedllm import UnifiedLLM +from opentelemetry import trace from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, computed_field from compass.compass_signals.compass_signals import CompassSignalResult @@ -25,6 +26,8 @@ ) from compass.traces import TraceSnapshot +tracer = trace.get_tracer(__name__) + USER_SENTIMENT = """ Identify explicit user anger or criticism of the agent system, its responses, behavior, tools, or work in the supplied user messages. Each message is screened @@ -115,21 +118,37 @@ 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[:20000] + for messages in user_messages.values() + for text in messages + if text[:20000].strip() ) ) + with tracer.start_as_current_span( + "sentiment.embeddings", + attributes={ + "openinference.span.kind": "CHAIN", + "message_count": sum(len(messages) for messages in user_messages.values()), + "unique_input_count": len(unique_texts), + }, + ): + embeddings = embedding_generator.generate_batch(unique_texts) + with tracer.start_as_current_span( + "sentiment.classification", attributes={"openinference.span.kind": "CHAIN"} + ): + 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)) + classified[text[:20000]].model_copy(deep=True) if text[:20000].strip() else ScreeningResult(label="no_user_messages") for text in messages @@ -182,7 +201,10 @@ async def analyze(self, snapshot: TraceSnapshot) -> CompassSignalResult: async def _analyze(self, snapshot: TraceSnapshot) -> CompassSignalResult: messages = {} if len(snapshot): - messages = await UserMessageExtractor(llm=self.llm).extract(snapshot) + with tracer.start_as_current_span( + "sentiment.extraction", attributes={"openinference.span.kind": "CHAIN"} + ): + messages = await UserMessageExtractor(llm=self.llm).extract(snapshot) messages = TypeAdapter(dict[str, list[str]]).validate_python(messages, strict=True) if messages.keys() != snapshot.traces_by_id.keys(): raise ValueError("User message extraction must cover exactly the supplied trace IDs") diff --git a/tests/test_user_embedding.py b/tests/test_user_embedding.py index 99d5de08..f2697a95 100644 --- a/tests/test_user_embedding.py +++ b/tests/test_user_embedding.py @@ -4,6 +4,10 @@ import numpy as np import pytest from litellm import EmbeddingResponse +from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from tokenizers import Tokenizer, models, pre_tokenizers from compass.compass_signals.user_embedding.embedding import ( @@ -30,15 +34,23 @@ 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_and_trace_context(monkeypatch): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + tracer = provider.get_tracer("test") + monkeypatch.setattr("compass.compass_signals.user_embedding.embedding.tracer", tracer) + monkeypatch.setattr("compass.compass_signals.user_sentiment.signal.tracer", tracer) generator = UserEmbeddingGenerator( litellm=LiteLLMEmbeddingConfig(model="openai/qwen", batch_size=4) ) generator.remote_tokenizer = Tokenizer(models.WordLevel({"[UNK]": 0}, unk_token="[UNK]")) calls = [] + request_parents = [] def embed(**kwargs): calls.append(kwargs["input"]) + request_parents.append(trace.get_current_span().get_span_context().span_id) # Reversed response order and non-unit vectors exercise both requirements. return EmbeddingResponse( data=[ @@ -50,19 +62,47 @@ def embed(**kwargs): monkeypatch.setattr( "compass.compass_signals.user_embedding.embedding._request_embeddings", embed ) - messages = {"a": ["message"] * 8, "b": ["", "last"], "c": []} - 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 + long = "x" * 20000 + messages: dict[str, list[str]] = { + "a": ["message"] * 8 + [" spaced ", "spaced", long + "a"], + "b": ["", "last", long + "b", "message"], + "c": [], + "d": [" " * 20000 + "ignored"], + } + with tracer.start_as_current_span("signal.user_sentiment") as parent: + results = screen_user_messages( + messages, generator, ComplaintClassifier(generator.projection) + ) + 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__ + spans = exporter.get_finished_spans() + embeddings = next(span for span in spans if span.name == "sentiment.embeddings") + batches = [span for span in spans if span.name == "sentiment.embedding_batch"] + assert embeddings.parent is not None + assert embeddings.attributes is not None + assert embeddings.parent.span_id == parent.get_span_context().span_id + assert embeddings.attributes["unique_input_count"] == 5 + assert len(batches) == 2 + assert all( + span.parent is not None and span.parent.span_id == embeddings.context.span_id + for span in batches + ) + assert set(request_parents) == {span.context.span_id for span in batches} + provider.shutdown() def test_remote_rejects_oversized_input_before_request_and_propagates_failures(monkeypatch): From 9daa47f9a6c3ac9fd0996aa110354708caf04b16 Mon Sep 17 00:00:00 2001 From: Dan Haggerty Date: Wed, 4 Nov 2026 09:58:51 -0500 Subject: [PATCH 2/4] Separate sentiment timing instrumentation from embedding reuse --- .../user_embedding/embedding.py | 28 ++-------------- tests/test_user_embedding.py | 33 ++----------------- 2 files changed, 4 insertions(+), 57 deletions(-) diff --git a/src/compass/compass_signals/user_embedding/embedding.py b/src/compass/compass_signals/user_embedding/embedding.py index 59e6d5ba..f21a0f10 100644 --- a/src/compass/compass_signals/user_embedding/embedding.py +++ b/src/compass/compass_signals/user_embedding/embedding.py @@ -8,22 +8,18 @@ import os from collections.abc import Iterator from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait -from contextvars import copy_context from dataclasses import dataclass from functools import cached_property from importlib.resources import files from importlib.util import find_spec from itertools import islice -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING import numpy as np from litellm import EmbeddingResponse -from opentelemetry import trace from pydantic import BaseModel, ConfigDict, Field from tokenizers import Tokenizer -tracer = trace.get_tracer(__name__) - class _NoRetries(int): """Zero retries, including in LiteLLM adapters that use ``value or default``.""" @@ -269,12 +265,7 @@ def generate_remote(self, texts: list[str]) -> list[UserEmbedding]: while True: # Keep only one batch per worker in flight; refill as each finishes. for batch in islice(batches, config.max_concurrency - len(active)): - future = cast( - Future[list[UserEmbedding]], - executor.submit( - copy_context().run, self._remote_batch, batch, config, api_key - ), - ) + future = executor.submit(self._remote_batch, batch, config, api_key) active[future] = batch.index if not active: break @@ -285,21 +276,6 @@ def generate_remote(self, texts: list[str]) -> list[UserEmbedding]: def _remote_batch( self, batch: _RemoteBatch, config: LiteLLMEmbeddingConfig, api_key: str | None - ) -> list[UserEmbedding]: - with tracer.start_as_current_span( - "sentiment.embedding_batch", - attributes={ - "openinference.span.kind": "CHAIN", - "input_offset": batch.index, - "input_count": len(batch.inputs), - "token_count": sum(batch.counts), - "model": config.model, - }, - ): - return self._encode_remote_batch(batch, config, api_key) - - def _encode_remote_batch( - self, batch: _RemoteBatch, config: LiteLLMEmbeddingConfig, api_key: str | None ) -> list[UserEmbedding]: response = _request_embeddings( model=config.model, diff --git a/tests/test_user_embedding.py b/tests/test_user_embedding.py index f2697a95..065261b0 100644 --- a/tests/test_user_embedding.py +++ b/tests/test_user_embedding.py @@ -4,10 +4,6 @@ import numpy as np import pytest from litellm import EmbeddingResponse -from opentelemetry import trace -from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.sdk.trace.export import SimpleSpanProcessor -from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from tokenizers import Tokenizer, models, pre_tokenizers from compass.compass_signals.user_embedding.embedding import ( @@ -34,23 +30,15 @@ def embed(text): assert generator.generate_batch([]) == [] -def test_screening_reuses_exact_inputs_preserving_occurrences_and_trace_context(monkeypatch): - exporter = InMemorySpanExporter() - provider = TracerProvider() - provider.add_span_processor(SimpleSpanProcessor(exporter)) - tracer = provider.get_tracer("test") - monkeypatch.setattr("compass.compass_signals.user_embedding.embedding.tracer", tracer) - monkeypatch.setattr("compass.compass_signals.user_sentiment.signal.tracer", tracer) +def test_screening_reuses_exact_inputs_preserving_occurrences(monkeypatch): generator = UserEmbeddingGenerator( litellm=LiteLLMEmbeddingConfig(model="openai/qwen", batch_size=4) ) generator.remote_tokenizer = Tokenizer(models.WordLevel({"[UNK]": 0}, unk_token="[UNK]")) calls = [] - request_parents = [] def embed(**kwargs): calls.append(kwargs["input"]) - request_parents.append(trace.get_current_span().get_span_context().span_id) # Reversed response order and non-unit vectors exercise both requirements. return EmbeddingResponse( data=[ @@ -69,10 +57,7 @@ def embed(**kwargs): "c": [], "d": [" " * 20000 + "ignored"], } - with tracer.start_as_current_span("signal.user_sentiment") as parent: - results = screen_user_messages( - messages, generator, ComplaintClassifier(generator.projection) - ) + results = screen_user_messages(messages, generator, ComplaintClassifier(generator.projection)) assert sorted(map(len, calls)) == [1, 4] prefix = generator.projection.metadata["encoder"]["prefix"] assert {text for call in calls for text in call} == { @@ -89,20 +74,6 @@ def embed(**kwargs): 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__ - spans = exporter.get_finished_spans() - embeddings = next(span for span in spans if span.name == "sentiment.embeddings") - batches = [span for span in spans if span.name == "sentiment.embedding_batch"] - assert embeddings.parent is not None - assert embeddings.attributes is not None - assert embeddings.parent.span_id == parent.get_span_context().span_id - assert embeddings.attributes["unique_input_count"] == 5 - assert len(batches) == 2 - assert all( - span.parent is not None and span.parent.span_id == embeddings.context.span_id - for span in batches - ) - assert set(request_parents) == {span.context.span_id for span in batches} - provider.shutdown() def test_remote_rejects_oversized_input_before_request_and_propagates_failures(monkeypatch): From 595400e2ee2079e3e6b1119f370b4a606c9ef564 Mon Sep 17 00:00:00 2001 From: Dan Haggerty Date: Wed, 4 Nov 2026 10:00:37 -0500 Subject: [PATCH 3/4] Move remaining sentiment stage spans to instrumentation change --- .../compass_signals/user_sentiment/signal.py | 29 ++++--------------- 1 file changed, 6 insertions(+), 23 deletions(-) diff --git a/src/compass/compass_signals/user_sentiment/signal.py b/src/compass/compass_signals/user_sentiment/signal.py index 0cfc5175..4c67b875 100644 --- a/src/compass/compass_signals/user_sentiment/signal.py +++ b/src/compass/compass_signals/user_sentiment/signal.py @@ -10,7 +10,6 @@ from nooa import Agent from nooa.unifiedllm import UnifiedLLM -from opentelemetry import trace from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, computed_field from compass.compass_signals.compass_signals import CompassSignalResult @@ -26,8 +25,6 @@ ) from compass.traces import TraceSnapshot -tracer = trace.get_tracer(__name__) - USER_SENTIMENT = """ Identify explicit user anger or criticism of the agent system, its responses, behavior, tools, or work in the supplied user messages. Each message is screened @@ -129,22 +126,11 @@ def screen_user_messages( if text[:20000].strip() ) ) - with tracer.start_as_current_span( - "sentiment.embeddings", - attributes={ - "openinference.span.kind": "CHAIN", - "message_count": sum(len(messages) for messages in user_messages.values()), - "unique_input_count": len(unique_texts), - }, - ): - embeddings = embedding_generator.generate_batch(unique_texts) - with tracer.start_as_current_span( - "sentiment.classification", attributes={"openinference.span.kind": "CHAIN"} - ): - classified = { - text: classifier.classify(embedding) - for text, embedding in zip(unique_texts, embeddings, strict=True) - } + 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=[ @@ -201,10 +187,7 @@ async def analyze(self, snapshot: TraceSnapshot) -> CompassSignalResult: async def _analyze(self, snapshot: TraceSnapshot) -> CompassSignalResult: messages = {} if len(snapshot): - with tracer.start_as_current_span( - "sentiment.extraction", attributes={"openinference.span.kind": "CHAIN"} - ): - messages = await UserMessageExtractor(llm=self.llm).extract(snapshot) + messages = await UserMessageExtractor(llm=self.llm).extract(snapshot) messages = TypeAdapter(dict[str, list[str]]).validate_python(messages, strict=True) if messages.keys() != snapshot.traces_by_id.keys(): raise ValueError("User message extraction must cover exactly the supplied trace IDs") From 2bbf62cacab43440608631576c36b3c8dbba5114 Mon Sep 17 00:00:00 2001 From: Dan Haggerty Date: Wed, 4 Nov 2026 14:39:11 -0500 Subject: [PATCH 4/4] Share the user message character limit across embedding and screening --- src/compass/compass_signals/user_embedding/embedding.py | 8 +++++--- src/compass/compass_signals/user_sentiment/signal.py | 9 +++++---- 2 files changed, 10 insertions(+), 7 deletions(-) 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 4c67b875..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, @@ -120,10 +121,10 @@ def screen_user_messages( # but retain whitespace: stripping would change the text sent to the model. unique_texts = list( dict.fromkeys( - text[:20000] + text[:MAX_USER_MESSAGE_CHARS] for messages in user_messages.values() for text in messages - if text[:20000].strip() + if text[:MAX_USER_MESSAGE_CHARS].strip() ) ) embeddings = embedding_generator.generate_batch(unique_texts) @@ -134,8 +135,8 @@ def screen_user_messages( return { trace_id: TraceScreeningResult( message_results=[ - classified[text[:20000]].model_copy(deep=True) - 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 ]