Merge pull request #256 from plastic-labs/ben/better-dedup

feat: better dedup with configurability
This commit is contained in:
doria 2025-11-13 11:46:43 -05:00 committed by GitHub
commit c8d75c77df
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 111 additions and 115 deletions

View File

@ -85,6 +85,7 @@ LLM_ANTHROPIC_API_KEY=your-anthropic-api-key-here
# DERIVER_STALE_SESSION_TIMEOUT_MINUTES=5
# DERIVER_PROVIDER=google
# DERIVER_MODEL=gemini-2.0-flash-lite
# DERIVER_DEDUPLICATE=true
# DERIVER_MAX_OUTPUT_TOKENS=2500
# only applied when using Anthropic as provider
# DERIVER_THINKING_BUDGET_TOKENS=1024

View File

@ -64,6 +64,7 @@ POLLING_SLEEP_INTERVAL_SECONDS = 1.0
STALE_SESSION_TIMEOUT_MINUTES = 5
PROVIDER = "google"
MODEL = "gemini-2.0-flash-lite"
DEDUPLICATE = true
MAX_OUTPUT_TOKENS = 2500
THINKING_BUDGET_TOKENS = 1024 # only applied when using Anthropic

View File

@ -200,6 +200,9 @@ class DeriverSettings(HonchoSettings):
PROVIDER: SupportedProviders = "google"
MODEL: str = "gemini-2.5-flash-lite"
# Whether to deduplicate documents when creating them
DEDUPLICATE: bool = True
MAX_OUTPUT_TOKENS: Annotated[int, Field(default=10_000, gt=0, le=100_000)] = 10_000
# Thinking budget tokens are only applied when using Anthropic as provider
THINKING_BUDGET_TOKENS: Annotated[int, Field(default=1024, gt=0, le=5000)] = 1024

View File

@ -103,9 +103,10 @@ async def create_documents(
*,
observer: str,
observed: str,
deduplicate: bool = False,
) -> int:
"""
Create multiple documents with NO duplicate detection.
Create multiple documents with optional duplicate detection.
Args:
db: Database session
@ -120,6 +121,16 @@ async def create_documents(
honcho_documents: list[models.Document] = []
for doc in documents:
try:
# for each document, if deduplicate is True, perform a process
# that checks against existing documents and either rejects this document
# as a duplicate OR deletes an existing document that is a duplicate.
if deduplicate:
is_duplicate = await is_rejected_duplicate(
db, doc, workspace_name, observer=observer, observed=observed
)
if is_duplicate:
continue
metadata_dict = doc.metadata.model_dump(exclude_none=True)
honcho_documents.append(
models.Document(
@ -149,3 +160,68 @@ async def create_documents(
) from e
return len(honcho_documents)
async def is_rejected_duplicate(
db: AsyncSession,
doc: schemas.DocumentCreate,
workspace_name: str,
*,
observer: str,
observed: str,
) -> bool:
"""
Check if a document is a duplicate of an existing document.
Uses: 1) Cosine similarity (>=0.95), 2) Token diff for retention.
Returns True if both:
- the document is deemed a duplicate of an existing document
- the existing document is deemed a superior duplicate
If the document is not a duplicate, returns False.
If the document is a duplicate AND the new document is superior,
deletes the existing document and returns False.
"""
# Step 1: Find potential duplicates using cosine similarity
similar_docs = await query_documents(
db=db,
workspace_name=workspace_name,
query=doc.content,
observer=observer,
observed=observed,
max_distance=0.05,
top_k=1,
embedding=doc.embedding,
)
if not similar_docs:
return False
existing_doc = similar_docs[0]
# Step 2: Determine which has more information using token set difference
tokens_new = set(embedding_client.encoding.encode(doc.content))
tokens_existing = set(embedding_client.encoding.encode(existing_doc.content))
unique_new = len(tokens_new - tokens_existing)
unique_existing = len(tokens_existing - tokens_new)
score_new = len(tokens_new) + (unique_new * 10)
score_existing = len(tokens_existing) + (unique_existing * 10)
# If new document has more or equal information, keep it and delete existing
if score_new >= score_existing:
logger.warning(
f"[DUPLICATE DETECTION] Deleting existing in favor of new. new='{doc.content}', existing='{existing_doc.content}'."
)
await db.delete(existing_doc)
await db.flush() # Flush to make deletion visible in this transaction
return False # Don't reject the new document
# Existing document has more information, reject the new one
logger.warning(
f"[DUPLICATE DETECTION] Rejecting new in favor of existing. new='{doc.content}', existing='{existing_doc.content}'."
)
return True

View File

@ -159,13 +159,14 @@ class RepresentationManager:
)
)
# Use bulk creation with NO duplicate detection
# Use bulk creation with optional duplicate detection
new_documents = await crud.create_documents(
db,
documents_to_create,
self.workspace_name,
observer=self.observer,
observed=self.observed,
deduplicate=settings.DERIVER.DEDUPLICATE,
)
try:
@ -175,41 +176,6 @@ class RepresentationManager:
return new_documents
async def get_relevant_observations(
self,
query: str,
*,
top_k: int = 5,
max_distance: float = 0.3,
level: str | None = None,
conversation_context: str = "",
) -> Representation:
"""
Unified method to get relevant observations with flexible options.
Args:
query: The search query
top_k: Number of results to return
max_distance: Maximum distance for semantic similarity
level: Optional reasoning level to filter by
conversation_context: Additional conversation context
Returns:
Representation
"""
async with tracked_db("representation_manager.get_relevant_observations") as db:
documents = await self._get_observations_internal(
db,
query,
top_k,
max_distance,
level,
conversation_context,
)
# convert documents to representation
return Representation.from_documents(documents)
async def get_working_representation(
self,
*,
@ -330,7 +296,6 @@ class RepresentationManager:
top_k: int,
max_distance: float | None = None,
level: str | None = None,
conversation_context: str = "",
) -> list[models.Document]:
"""Query documents by semantic similarity."""
try:
@ -339,7 +304,6 @@ class RepresentationManager:
db,
query,
level,
conversation_context,
top_k,
max_distance,
)
@ -349,7 +313,7 @@ class RepresentationManager:
workspace_name=self.workspace_name,
observer=self.observer,
observed=self.observed,
query=self._build_truncated_query(query, conversation_context),
query=query,
max_distance=max_distance,
top_k=top_k,
)
@ -412,11 +376,10 @@ class RepresentationManager:
top_k: int,
max_distance: float,
level: str | None,
conversation_context: str,
) -> list[models.Document]:
"""Internal method that does the actual observation retrieval."""
return await self._query_documents_semantic(
db, query, top_k, max_distance, level, conversation_context
db, query, top_k, max_distance, level
)
async def _query_documents_for_level(
@ -424,7 +387,6 @@ class RepresentationManager:
db: AsyncSession,
query: str,
level: str,
conversation_context: str,
count: int,
max_distance: float | None = None,
) -> list[models.Document]:
@ -434,7 +396,7 @@ class RepresentationManager:
workspace_name=self.workspace_name,
observer=self.observer,
observed=self.observed,
query=self._build_truncated_query(query, conversation_context),
query=query,
max_distance=max_distance,
top_k=count * FILTER_OVERSAMPLING_FACTOR,
filters=self._build_filter_conditions(level),
@ -461,69 +423,6 @@ class RepresentationManager:
return conditions[0] if len(conditions) == 1 else {"AND": conditions}
def _build_truncated_query(
self,
query: str,
conversation_context: str = "",
max_tokens: int | None = None,
) -> str:
"""Build a query that fits within token limits with clear priorities.
Args:
query: The search query
conversation_context: Optional conversation context to include
max_tokens: Maximum tokens allowed (defaults to setting with buffer)
Returns:
Truncated query string that fits within token limits
"""
max_tokens = max_tokens or (settings.MAX_EMBEDDING_TOKENS - 100)
encoding = embedding_client.encoding
# Pre-calculate all token counts once
query_prefix = "Current message: "
context_prefix = "\nContext: "
prefix_tokens = len(encoding.encode(query_prefix))
context_prefix_tokens = len(encoding.encode(context_prefix))
query_tokens = encoding.encode(query)
# Simple case: query alone fits
if prefix_tokens + len(query_tokens) <= max_tokens:
if not conversation_context:
return f"{query_prefix}{query}"
# Try to add context
context_tokens = encoding.encode(conversation_context)
total_without_context = (
prefix_tokens + len(query_tokens) + context_prefix_tokens
)
if total_without_context + len(context_tokens) <= max_tokens:
return f"{query_prefix}{query}{context_prefix}{conversation_context}"
# Truncate context to fit
available_context_tokens = max_tokens - total_without_context
if available_context_tokens > 0:
truncated_context = encoding.decode(
context_tokens[-available_context_tokens:]
)
return f"{query_prefix}{query}{context_prefix}{truncated_context}"
else:
# No room left for context; keep full query intact
return f"{query_prefix}{query}"
# Query itself is too long - truncate it
available_query_tokens = max_tokens - prefix_tokens
if available_query_tokens > 0:
# Keep the end (recency) of the query
truncated_query = encoding.decode(query_tokens[-available_query_tokens:])
return f"{query_prefix}{truncated_query}"
# Pathological case - just return what we can
logger.warning("Token limit too restrictive: %s", max_tokens)
return encoding.decode(query_tokens[:max_tokens])
# Module-level functions for backward compatibility and convenience

View File

@ -65,11 +65,9 @@ Current understanding of {peer_id}:
You are an agent who critically analyzes messages from {peer_id} through rigorous logical reasoning to produce only conclusions about them that are CERTAIN.
TARGET USER TO ANALYZE
You are analyzing: {peer_id}
The conversation may include messages from multiple participants, but you MUST focus ONLY on deriving conclusions about {peer_id}. Only use other participants' messages as context for understanding {peer_id}.
IMPORTANT NAMING RULES
When you write a conclusion about {peer_id}, always start the sentence with their name (e.g. "Anthony is 25 years old").
@ -87,7 +85,7 @@ Here are strict definitions for the reasoning modes you are to employ:
- Current date and time (which is: {message_created_at})
- Timestamps from conversation history
- Follow strict literal necessity--if stated directly in message, extract a conclusion
- Latest message MUST be a premise, previous messages and timestamps may be used to contextualize
- New turn MUST be a premise, previous messages and timestamps may be used to contextualize
- Transforms a single message (premise) into ONE OR MULTIPLE conclusions
- Derive EVERYTHING that can be explicitly concluded
- Make sure EVERY conclusion is sufficiently contextualized, i.e. ensure each conclusion contains enough specific information about subjects and objects to make it self-contained and useful (e.g. instead of "Ann is nervous about the interview", use "Ann is nervous about the job interview at the pharmacy")
@ -128,7 +126,7 @@ Based on our definitions and examples, here's a summary of the logical reasoning
- Explicit/Deductive/Temporal/General Further Deductive: Can use certain conclusions and known facts to deduce additional certain conclusions
- Probabilistic Deductive Further Deductive: If a deductive conclusion includes probabilistic qualifiers (likely, potentially, typically, might, etc) it may NOT be used as a premise for further deductions
**INSTRUCTIONS:** Given the above, first think critically about what it means to do explicit and deductive reasoning, then consider how to apply that to the latest message, finally do explicit and deductive reasoning about the user to reach useful, contextually-rich conclusions.
**INSTRUCTIONS:** Given the above, first think critically about what it means to do explicit and deductive reasoning, then consider how to apply that to all new turns, finally do explicit and deductive reasoning about the user to reach useful, contextually-rich conclusions. You must extract observations from all new turns.
{peer_card_section}

View File

@ -870,7 +870,7 @@ Evaluate whether the actual response correctly answers the question based on the
return results
async def run_all_questions(
self, test_file: Path, batch_size: int = 10
self, test_file: Path, batch_size: int = 10, test_count: int | None = None
) -> tuple[list[TestResult], float]:
"""
Run all questions in a longmemeval test file.
@ -878,11 +878,20 @@ Evaluate whether the actual response correctly answers the question based on the
Args:
test_file: Path to the longmemeval JSON file
batch_size: Number of questions to run concurrently in each batch
test_count: Optional number of tests to run (runs first N tests)
Returns:
Tuple of (list of test results, total duration)
"""
questions = self.load_test_file(test_file)
# Limit to first N questions if test_count is specified
if test_count is not None and test_count > 0:
questions = questions[:test_count]
print(
f"limiting to first {len(questions)} {'question' if len(questions) == 1 else 'questions'} from {test_file}"
)
print(
f"found {len(questions)} {'question' if len(questions) == 1 else 'questions'} in {test_file}"
)
@ -1140,6 +1149,7 @@ Examples:
%(prog)s --test-file tests/bench/longmemeval_data/longmemeval_s.json # Run longmemeval tests
%(prog)s --test-file test.json --pool-size 4 # Use 4 Honcho instances
%(prog)s --test-file test.json --base-api-port 8000 --pool-size 4 # Custom base port with pool
%(prog)s --test-file test.json --test-count 50 # Run only first 50 tests
""",
)
@ -1208,6 +1218,12 @@ Examples:
help="Use get_context + judge LLM instead of dialectic .chat endpoint (default: False)",
)
parser.add_argument(
"--test-count",
type=int,
help="Number of tests to run from the test file (default: all tests)",
)
args = parser.parse_args()
# Validate arguments
@ -1223,6 +1239,10 @@ Examples:
print(f"Error: Pool size must be positive, got {args.pool_size}")
return 1
if args.test_count is not None and args.test_count <= 0:
print(f"Error: Test count must be positive, got {args.test_count}")
return 1
# Create test runner
runner = LongMemEvalRunner(
base_api_port=args.base_api_port,
@ -1237,7 +1257,7 @@ Examples:
try:
# Run all questions
results, total_elapsed = await runner.run_all_questions(
args.test_file, args.batch_size
args.test_file, args.batch_size, args.test_count
)
runner.print_summary(results, total_elapsed_seconds=total_elapsed)

View File

@ -310,6 +310,5 @@ def mock_representation_manager(monkeypatch: pytest.MonkeyPatch) -> AsyncMock:
mock_manager = AsyncMock(spec=RepresentationManager)
mock_manager.save_representation.return_value = 0
mock_manager.get_relevant_observations = AsyncMock(return_value=MagicMock())
return mock_manager

View File

@ -102,7 +102,6 @@ class TestDeriverProcessing:
await mock_representation_manager.save_representation(
Representation(explicit=[], deductive=[])
)
mock_representation_manager.get_relevant_observations.return_value = [] # type: ignore[attr-defined]
# Verify the methods were called
assert mock_representation_manager.save_representation.called # type: ignore[attr-defined]