Merge pull request #256 from plastic-labs/ben/better-dedup
feat: better dedup with configurability
This commit is contained in:
commit
c8d75c77df
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
Loading…
Reference in New Issue