927 lines
33 KiB
Python
927 lines
33 KiB
Python
import datetime
|
|
from collections.abc import Sequence
|
|
from logging import getLogger
|
|
from typing import Any
|
|
|
|
from sqlalchemy import delete, select, update
|
|
from sqlalchemy.exc import IntegrityError
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from sqlalchemy.sql import Select
|
|
from sqlalchemy.sql.functions import func
|
|
from tenacity import (
|
|
AsyncRetrying,
|
|
retry_if_exception_type,
|
|
retry_if_result,
|
|
stop_after_attempt,
|
|
wait_exponential,
|
|
)
|
|
|
|
from src import models, schemas
|
|
from src.config import settings
|
|
from src.crud.collection import get_or_create_collection
|
|
from src.crud.peer import get_peer
|
|
from src.crud.session import get_session
|
|
from src.embedding_client import embedding_client
|
|
from src.exceptions import ResourceNotFoundException, ValidationException
|
|
from src.utils.filter import apply_filter
|
|
from src.vector_store import VectorRecord, VectorStore, get_vector_store
|
|
|
|
logger = getLogger(__name__)
|
|
|
|
|
|
def get_all_documents(
|
|
workspace_name: str,
|
|
*,
|
|
observer: str,
|
|
observed: str,
|
|
filters: dict[str, Any] | None = None,
|
|
reverse: bool = False,
|
|
limit: int | None = None,
|
|
) -> Select[tuple[models.Document]]:
|
|
"""
|
|
Get all documents in a collection.
|
|
|
|
Returns a Select query for pagination support via apaginate().
|
|
Results are ordered by created_at timestamp.
|
|
|
|
Args:
|
|
workspace_name: Name of the workspace
|
|
observer: Name of the observing peer
|
|
observed: Name of the observed peer
|
|
filters: Optional filters to apply
|
|
reverse: Whether to reverse the order (oldest first)
|
|
|
|
Returns:
|
|
Select query for documents
|
|
"""
|
|
stmt = (
|
|
select(models.Document)
|
|
.where(models.Document.workspace_name == workspace_name)
|
|
.where(models.Document.observer == observer)
|
|
.where(models.Document.observed == observed)
|
|
.where(models.Document.deleted_at.is_(None)) # Exclude soft-deleted
|
|
)
|
|
|
|
# Apply additional filters if provided
|
|
stmt = apply_filter(stmt, models.Document, filters)
|
|
|
|
# Order by created_at (newest first by default)
|
|
if reverse:
|
|
stmt = stmt.order_by(models.Document.created_at.asc())
|
|
else:
|
|
stmt = stmt.order_by(models.Document.created_at.desc())
|
|
|
|
if limit is not None:
|
|
stmt = stmt.limit(limit)
|
|
|
|
return stmt
|
|
|
|
|
|
def get_documents_with_filters(
|
|
workspace_name: str,
|
|
*,
|
|
filters: dict[str, Any] | None = None,
|
|
reverse: bool = False,
|
|
) -> Select[tuple[models.Document]]:
|
|
"""
|
|
Get all documents using custom filters.
|
|
|
|
Returns a Select query for pagination support via apaginate().
|
|
Results are ordered by created_at timestamp.
|
|
|
|
Args:
|
|
workspace_name: Name of the workspace
|
|
filters: Optional filters to apply
|
|
reverse: Whether to reverse the order (oldest first)
|
|
|
|
Returns:
|
|
Select query for documents
|
|
"""
|
|
stmt = (
|
|
select(models.Document)
|
|
.where(models.Document.workspace_name == workspace_name)
|
|
.where(models.Document.deleted_at.is_(None)) # Exclude soft-deleted
|
|
)
|
|
|
|
# Apply additional filters if provided
|
|
stmt = apply_filter(stmt, models.Document, filters)
|
|
|
|
# Order by created_at (newest first by default)
|
|
if reverse:
|
|
stmt = stmt.order_by(models.Document.created_at.asc())
|
|
else:
|
|
stmt = stmt.order_by(models.Document.created_at.desc())
|
|
|
|
return stmt
|
|
|
|
|
|
async def query_documents(
|
|
db: AsyncSession,
|
|
workspace_name: str,
|
|
query: str,
|
|
*,
|
|
observer: str,
|
|
observed: str,
|
|
filters: dict[str, Any] | None = None,
|
|
max_distance: float | None = None,
|
|
top_k: int = 5,
|
|
embedding: list[float] | None = None,
|
|
) -> Sequence[models.Document]:
|
|
"""
|
|
Query documents using semantic similarity.
|
|
|
|
Args:
|
|
db: Database session
|
|
workspace_name: Name of the workspace
|
|
query: Search query text
|
|
observer: Name of the observing peer
|
|
observed: Name of the observed peer
|
|
filters: Optional filters to apply at vector store level (supports: level, session_name)
|
|
max_distance: Maximum cosine distance for results
|
|
top_k: Number of results to return
|
|
embedding: Optional pre-computed embedding for the query (avoids extra API call if possible)
|
|
|
|
Returns:
|
|
Sequence of matching documents
|
|
"""
|
|
# Use provided embedding or generate one
|
|
if embedding is None:
|
|
try:
|
|
embedding = await embedding_client.embed(query)
|
|
except ValueError as e:
|
|
raise ValidationException(
|
|
f"Query exceeds maximum token limit of {settings.MAX_EMBEDDING_TOKENS}."
|
|
) from e
|
|
|
|
# Get vector store and namespace for this collection
|
|
vector_store = get_vector_store()
|
|
namespace = vector_store.get_vector_namespace(
|
|
"document", workspace_name, observer, observed
|
|
)
|
|
|
|
# Build vector store filters
|
|
# Convert filter dict to vector store format (handles level, session_name, etc.)
|
|
vector_filters: dict[str, Any] = {}
|
|
if filters:
|
|
# Direct pass-through for simple equality filters
|
|
# The filters dict can contain: level, session_name, or other document fields
|
|
# We can push level and session_name to vector store since they're in metadata
|
|
for key in ["level", "session_name"]:
|
|
if key in filters:
|
|
vector_filters[key] = filters[key]
|
|
|
|
# Query vector store for similar documents with filters applied
|
|
vector_results = await vector_store.query(
|
|
namespace,
|
|
embedding,
|
|
top_k=top_k,
|
|
max_distance=max_distance,
|
|
filters=vector_filters if vector_filters else None,
|
|
)
|
|
|
|
if not vector_results:
|
|
return []
|
|
|
|
# Get document IDs from vector results (vector ID = document ID for documents)
|
|
document_ids = [result.id for result in vector_results]
|
|
|
|
# Fetch documents from database
|
|
stmt = (
|
|
select(models.Document)
|
|
.where(models.Document.workspace_name == workspace_name)
|
|
.where(models.Document.observer == observer)
|
|
.where(models.Document.observed == observed)
|
|
.where(models.Document.deleted_at.is_(None))
|
|
.where(models.Document.id.in_(document_ids))
|
|
)
|
|
# Re-apply all filters at the database layer to catch any constraints
|
|
# that aren't supported by the vector store metadata.
|
|
stmt = apply_filter(stmt, models.Document, filters)
|
|
|
|
result = await db.execute(stmt)
|
|
documents = {doc.id: doc for doc in result.scalars().all()}
|
|
|
|
# Return documents in order of similarity (preserving vector store order)
|
|
ordered_docs: list[models.Document] = []
|
|
for vr in vector_results:
|
|
if vr.id in documents:
|
|
ordered_docs.append(documents[vr.id])
|
|
|
|
return ordered_docs
|
|
|
|
|
|
async def create_documents(
|
|
db: AsyncSession,
|
|
documents: list[schemas.DocumentCreate],
|
|
workspace_name: str,
|
|
*,
|
|
observer: str,
|
|
observed: str,
|
|
deduplicate: bool = False,
|
|
) -> int:
|
|
"""
|
|
Create multiple documents with optional duplicate detection.
|
|
|
|
Args:
|
|
db: Database session
|
|
documents: List of document creation schemas
|
|
workspace_name: Name of the workspace
|
|
observer: Name of the observing peer
|
|
observed: Name of the observed peer
|
|
|
|
Returns:
|
|
Count of new documents
|
|
"""
|
|
honcho_documents: list[models.Document] = []
|
|
# Store (document_model, embedding) pairs - IDs aren't available until after commit
|
|
docs_with_embeddings: list[tuple[models.Document, list[float]]] = []
|
|
|
|
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)
|
|
|
|
# Check if pgvector is being used (primary or secondary)
|
|
# If so, write embeddings to ORM since pgvector relies on postgres
|
|
pgvector_in_use = (
|
|
settings.VECTOR_STORE.PRIMARY_TYPE == "pgvector"
|
|
or settings.VECTOR_STORE.SECONDARY_TYPE == "pgvector"
|
|
)
|
|
|
|
if pgvector_in_use and doc.embedding:
|
|
# pgvector in use: write embedding to ORM (postgres)
|
|
new_doc = models.Document(
|
|
workspace_name=workspace_name,
|
|
observer=observer,
|
|
observed=observed,
|
|
content=doc.content,
|
|
level=doc.level,
|
|
times_derived=doc.times_derived,
|
|
internal_metadata=metadata_dict,
|
|
session_name=doc.session_name,
|
|
embedding=doc.embedding,
|
|
)
|
|
else:
|
|
# pgvector not in use or no embedding: don't write embedding to postgres
|
|
new_doc = models.Document(
|
|
workspace_name=workspace_name,
|
|
observer=observer,
|
|
observed=observed,
|
|
content=doc.content,
|
|
level=doc.level,
|
|
times_derived=doc.times_derived,
|
|
internal_metadata=metadata_dict,
|
|
session_name=doc.session_name,
|
|
)
|
|
|
|
if doc.embedding:
|
|
new_doc.sync_state = "pending"
|
|
honcho_documents.append(new_doc)
|
|
|
|
# Track embedding for vector store (ID will be available after commit)
|
|
if doc.embedding:
|
|
docs_with_embeddings.append((new_doc, doc.embedding))
|
|
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Error adding new document to {workspace_name}/{doc.session_name}/{observer}/{observed}: {e}"
|
|
)
|
|
continue
|
|
|
|
try:
|
|
db.add_all(honcho_documents)
|
|
# NOTE
|
|
# If the process crashes after this commit but before vector upsert completes,
|
|
# documents will be left in sync_state='pending' with NULL embeddings.
|
|
# The reconciliation job will automatically re-embed and sync these documents,
|
|
await db.commit()
|
|
|
|
# Store embeddings in vector store after documents are committed (IDs now available)
|
|
if docs_with_embeddings:
|
|
vector_store = get_vector_store()
|
|
namespace = vector_store.get_vector_namespace(
|
|
"document",
|
|
workspace_name,
|
|
observer,
|
|
observed,
|
|
)
|
|
|
|
# Build vector records with metadata for filtering
|
|
vector_records: list[VectorRecord] = []
|
|
doc_ids: list[str] = []
|
|
for doc, embedding in docs_with_embeddings:
|
|
doc_ids.append(doc.id)
|
|
vector_records.append(
|
|
VectorRecord(
|
|
id=doc.id,
|
|
embedding=embedding,
|
|
metadata={
|
|
"workspace_name": workspace_name,
|
|
"observer": observer,
|
|
"observed": observed,
|
|
"session_name": doc.session_name,
|
|
"level": doc.level,
|
|
},
|
|
)
|
|
)
|
|
|
|
# Retry vector upsert with exponential backoff (3 attempts)
|
|
try:
|
|
result = None
|
|
async for attempt in AsyncRetrying(
|
|
stop=stop_after_attempt(3),
|
|
wait=wait_exponential(multiplier=0.5, min=0.5, max=2.0),
|
|
retry=retry_if_exception_type(Exception)
|
|
| retry_if_result(
|
|
lambda res: res is not None and res.secondary_ok is False
|
|
),
|
|
reraise=True,
|
|
):
|
|
with attempt:
|
|
result = await vector_store.upsert_many(
|
|
namespace, vector_records
|
|
)
|
|
|
|
if result is not None and result.secondary_ok is False:
|
|
# Partial success: primary has data but secondary doesn't
|
|
# Keep as "pending" for reconciliation to sync secondary
|
|
logger.warning(
|
|
f"Partial sync for namespace {namespace}: {result.secondary_error}"
|
|
)
|
|
await db.execute(
|
|
update(models.Document)
|
|
.where(models.Document.id.in_(doc_ids))
|
|
.values(
|
|
sync_attempts=models.Document.sync_attempts + 1,
|
|
last_sync_at=func.now(),
|
|
)
|
|
)
|
|
await db.commit()
|
|
else:
|
|
# Success: both primary and secondary stores have the data
|
|
await db.execute(
|
|
update(models.Document)
|
|
.where(models.Document.id.in_(doc_ids))
|
|
.values(
|
|
sync_state="synced",
|
|
last_sync_at=func.now(),
|
|
sync_attempts=0,
|
|
)
|
|
)
|
|
await db.commit()
|
|
|
|
except Exception as e:
|
|
# Total failure: primary write failed
|
|
# Keep as "pending" for reconciliation to retry
|
|
logger.error(f"Failed to upsert vectors after 3 retries: {e}")
|
|
await db.execute(
|
|
update(models.Document)
|
|
.where(models.Document.id.in_(doc_ids))
|
|
.values(
|
|
sync_attempts=models.Document.sync_attempts + 1,
|
|
last_sync_at=func.now(),
|
|
)
|
|
)
|
|
await db.commit()
|
|
|
|
except IntegrityError as e:
|
|
await db.rollback()
|
|
raise ValidationException(
|
|
"Failed to create documents due to integrity constraint violation"
|
|
) from e
|
|
|
|
return len(honcho_documents)
|
|
|
|
|
|
async def delete_document(
|
|
db: AsyncSession,
|
|
workspace_name: str,
|
|
document_id: str,
|
|
*,
|
|
observer: str,
|
|
observed: str,
|
|
session_name: str | None = None,
|
|
) -> None:
|
|
"""
|
|
Delete a single document by ID using hybrid sync/soft delete pattern.
|
|
|
|
Tries to delete from vector store first, then hard deletes from DB.
|
|
If vector store delete fails, soft deletes (sets deleted_at) and lets
|
|
cleanup job handle vector deletion later.
|
|
|
|
Args:
|
|
db: Database session
|
|
workspace_name: Name of the workspace
|
|
document_id: ID of the document to delete
|
|
observer: Name of the observing peer (for authorization)
|
|
observed: Name of the observed peer (for authorization)
|
|
session_name: Optional session name to verify document belongs to session
|
|
|
|
Raises:
|
|
ResourceNotFoundException: If document not found or doesn't match criteria
|
|
"""
|
|
# Build base query conditions
|
|
conditions = [
|
|
models.Document.id == document_id,
|
|
models.Document.workspace_name == workspace_name,
|
|
models.Document.observer == observer,
|
|
models.Document.observed == observed,
|
|
models.Document.deleted_at.is_(None), # Only delete non-deleted docs
|
|
]
|
|
if session_name is not None:
|
|
conditions.append(models.Document.session_name == session_name)
|
|
|
|
# Check document exists first
|
|
check_stmt = select(models.Document).where(*conditions)
|
|
result = await db.execute(check_stmt)
|
|
doc = result.scalar_one_or_none()
|
|
|
|
if doc is None:
|
|
raise ResourceNotFoundException(
|
|
f"Document {document_id} not found or does not belong to the specified collection/session"
|
|
)
|
|
|
|
# Try to delete from vector store first
|
|
vector_store = get_vector_store()
|
|
namespace = vector_store.get_vector_namespace(
|
|
"document", workspace_name, observer, observed
|
|
)
|
|
vector_deleted = False
|
|
|
|
try:
|
|
await vector_store.delete_many(namespace, [document_id])
|
|
vector_deleted = True
|
|
except Exception as e:
|
|
logger.warning(f"Failed to delete vector for document {document_id}: {e}")
|
|
|
|
if vector_deleted:
|
|
# Happy path: hard delete from DB
|
|
delete_stmt = delete(models.Document).where(models.Document.id == document_id)
|
|
await db.execute(delete_stmt)
|
|
else:
|
|
# Fallback: soft delete, let cleanup job handle vector
|
|
update_stmt = (
|
|
update(models.Document)
|
|
.where(models.Document.id == document_id)
|
|
.values(deleted_at=func.now())
|
|
)
|
|
await db.execute(update_stmt)
|
|
|
|
await db.commit()
|
|
|
|
|
|
async def delete_document_by_id(
|
|
db: AsyncSession,
|
|
workspace_name: str,
|
|
document_id: str,
|
|
) -> None:
|
|
"""
|
|
Delete a single document by ID and workspace using hybrid sync/soft delete pattern.
|
|
|
|
Tries to delete from vector store first, then hard deletes from DB.
|
|
If vector store delete fails, soft deletes (sets deleted_at) and lets
|
|
cleanup job handle vector deletion later.
|
|
|
|
Args:
|
|
db: Database session
|
|
workspace_name: Name of the workspace
|
|
document_id: ID of the document to delete
|
|
|
|
Raises:
|
|
ResourceNotFoundException: If document not found or doesn't belong to the workspace
|
|
"""
|
|
# Fetch document to get observer/observed for namespace
|
|
stmt = select(models.Document).where(
|
|
models.Document.id == document_id,
|
|
models.Document.workspace_name == workspace_name,
|
|
models.Document.deleted_at.is_(None), # Only delete non-deleted docs
|
|
)
|
|
result = await db.execute(stmt)
|
|
doc = result.scalar_one_or_none()
|
|
|
|
if doc is None:
|
|
raise ResourceNotFoundException(
|
|
f"Document {document_id} not found or does not belong to workspace {workspace_name}"
|
|
)
|
|
|
|
# Try to delete from vector store first
|
|
vector_store = get_vector_store()
|
|
namespace = vector_store.get_vector_namespace(
|
|
"document",
|
|
workspace_name,
|
|
doc.observer,
|
|
doc.observed,
|
|
)
|
|
vector_deleted = False
|
|
|
|
try:
|
|
await vector_store.delete_many(namespace, [document_id])
|
|
vector_deleted = True
|
|
except Exception as e:
|
|
logger.warning(f"Failed to delete vector for document {document_id}: {e}")
|
|
|
|
if vector_deleted:
|
|
# Happy path: hard delete from DB
|
|
delete_stmt = delete(models.Document).where(models.Document.id == document_id)
|
|
await db.execute(delete_stmt)
|
|
else:
|
|
# Fallback: soft delete, let cleanup job handle vector
|
|
update_stmt = (
|
|
update(models.Document)
|
|
.where(models.Document.id == document_id)
|
|
.values(deleted_at=func.now())
|
|
)
|
|
await db.execute(update_stmt)
|
|
|
|
await db.commit()
|
|
|
|
|
|
async def create_observations(
|
|
db: AsyncSession,
|
|
observations: list[schemas.ObservationCreate],
|
|
workspace_name: str,
|
|
) -> list[models.Document]:
|
|
"""
|
|
Create multiple observations (documents) from user input.
|
|
|
|
This function validates all referenced resources, generates embeddings
|
|
in batch, and creates the documents.
|
|
|
|
Args:
|
|
db: Database session
|
|
observations: List of observation creation schemas
|
|
workspace_name: Name of the workspace
|
|
|
|
Returns:
|
|
List of created Document objects
|
|
|
|
Raises:
|
|
ResourceNotFoundException: If any session or peer is not found
|
|
ValidationException: If embedding generation fails or integrity constraint is violated
|
|
"""
|
|
if not observations:
|
|
return []
|
|
|
|
# Collect unique sessions and peer pairs to validate
|
|
sessions_to_validate: set[str] = set()
|
|
peers_to_validate: set[str] = set()
|
|
collection_pairs: set[tuple[str, str]] = set()
|
|
|
|
for obs in observations:
|
|
sessions_to_validate.add(obs.session_id)
|
|
peers_to_validate.add(obs.observer_id)
|
|
peers_to_validate.add(obs.observed_id)
|
|
collection_pairs.add((obs.observer_id, obs.observed_id))
|
|
|
|
# Validate all sessions exist
|
|
for session_name in sessions_to_validate:
|
|
await get_session(db, session_name, workspace_name)
|
|
|
|
# Validate all peers exist
|
|
for peer_name in peers_to_validate:
|
|
await get_peer(db, workspace_name, schemas.PeerCreate(name=peer_name))
|
|
|
|
# Get or create all collections
|
|
for observer, observed in collection_pairs:
|
|
await get_or_create_collection(
|
|
db, workspace_name, observer=observer, observed=observed
|
|
)
|
|
|
|
# Generate embeddings in batch
|
|
contents = [obs.content for obs in observations]
|
|
try:
|
|
embeddings = await embedding_client.simple_batch_embed(contents)
|
|
except ValueError as e:
|
|
raise ValidationException(str(e)) from e
|
|
|
|
# Create document objects and track embeddings for vector store
|
|
honcho_documents: list[models.Document] = []
|
|
# Group observations by collection (observer, observed) for vector store upserts
|
|
collection_embeddings: dict[
|
|
tuple[str, str], list[tuple[models.Document, list[float]]]
|
|
] = {}
|
|
|
|
# Check if pgvector is being used (primary or secondary)
|
|
# If so, write embeddings to ORM since pgvector relies on postgres
|
|
pgvector_in_use = (
|
|
settings.VECTOR_STORE.PRIMARY_TYPE == "pgvector"
|
|
or settings.VECTOR_STORE.SECONDARY_TYPE == "pgvector"
|
|
)
|
|
|
|
for obs, embedding in zip(observations, embeddings, strict=True):
|
|
if pgvector_in_use:
|
|
# pgvector in use: write embedding to ORM (postgres)
|
|
doc = models.Document(
|
|
workspace_name=workspace_name,
|
|
observer=obs.observer_id,
|
|
observed=obs.observed_id,
|
|
content=obs.content,
|
|
level="explicit", # Manually created observations are always explicit
|
|
times_derived=1,
|
|
internal_metadata={}, # No message_ids since not derived from messages
|
|
session_name=obs.session_id,
|
|
embedding=embedding,
|
|
)
|
|
else:
|
|
# pgvector not in use: don't write embedding to postgres
|
|
doc = models.Document(
|
|
workspace_name=workspace_name,
|
|
observer=obs.observer_id,
|
|
observed=obs.observed_id,
|
|
content=obs.content,
|
|
level="explicit", # Manually created observations are always explicit
|
|
times_derived=1,
|
|
internal_metadata={}, # No message_ids since not derived from messages
|
|
session_name=obs.session_id,
|
|
)
|
|
doc.sync_state = "pending"
|
|
honcho_documents.append(doc)
|
|
|
|
# Track embedding for vector store (grouped by collection)
|
|
collection_key = (obs.observer_id, obs.observed_id)
|
|
if collection_key not in collection_embeddings:
|
|
collection_embeddings[collection_key] = []
|
|
collection_embeddings[collection_key].append((doc, embedding))
|
|
|
|
try:
|
|
db.add_all(honcho_documents)
|
|
await db.commit()
|
|
# Refresh all documents to get generated IDs and timestamps
|
|
for doc in honcho_documents:
|
|
await db.refresh(doc)
|
|
|
|
# Store embeddings in vector store after documents are committed (IDs now available)
|
|
vector_store = get_vector_store()
|
|
for (observer, observed), docs_with_embeddings in collection_embeddings.items():
|
|
namespace = vector_store.get_vector_namespace(
|
|
"document",
|
|
workspace_name,
|
|
observer,
|
|
observed,
|
|
)
|
|
|
|
# Build vector records with metadata for filtering
|
|
vector_records: list[VectorRecord] = []
|
|
doc_ids: list[str] = []
|
|
for doc, embedding in docs_with_embeddings:
|
|
doc_ids.append(doc.id)
|
|
vector_records.append(
|
|
VectorRecord(
|
|
id=doc.id,
|
|
embedding=embedding,
|
|
metadata={
|
|
"workspace_name": workspace_name,
|
|
"observer": observer,
|
|
"observed": observed,
|
|
"session_name": doc.session_name,
|
|
"level": doc.level,
|
|
},
|
|
)
|
|
)
|
|
|
|
# Retry vector upsert with exponential backoff (3 attempts)
|
|
try:
|
|
result = None
|
|
async for attempt in AsyncRetrying(
|
|
stop=stop_after_attempt(3),
|
|
wait=wait_exponential(multiplier=0.5, min=0.5, max=2.0),
|
|
retry=retry_if_exception_type(Exception)
|
|
| retry_if_result(
|
|
lambda res: res is not None and res.secondary_ok is False
|
|
),
|
|
reraise=True,
|
|
):
|
|
with attempt:
|
|
result = await vector_store.upsert_many(
|
|
namespace, vector_records
|
|
)
|
|
|
|
if result is not None and result.secondary_ok is False:
|
|
# Partial success: primary has data but secondary doesn't
|
|
# Keep as "pending" for reconciliation to sync secondary
|
|
logger.warning(
|
|
f"Partial sync for namespace {namespace}: {result.secondary_error}"
|
|
)
|
|
await db.execute(
|
|
update(models.Document)
|
|
.where(models.Document.id.in_(doc_ids))
|
|
.values(
|
|
sync_attempts=models.Document.sync_attempts + 1,
|
|
last_sync_at=func.now(),
|
|
)
|
|
)
|
|
await db.commit()
|
|
else:
|
|
# Success: both primary and secondary stores have the data
|
|
await db.execute(
|
|
update(models.Document)
|
|
.where(models.Document.id.in_(doc_ids))
|
|
.values(
|
|
sync_state="synced",
|
|
last_sync_at=func.now(),
|
|
sync_attempts=0,
|
|
)
|
|
)
|
|
await db.commit()
|
|
|
|
except Exception as e:
|
|
# Total failure: primary write failed
|
|
# Keep as "pending" for reconciliation to retry
|
|
logger.error(
|
|
f"Failed to upsert vectors for {namespace} after 3 retries: {e}"
|
|
)
|
|
await db.execute(
|
|
update(models.Document)
|
|
.where(models.Document.id.in_(doc_ids))
|
|
.values(
|
|
sync_attempts=models.Document.sync_attempts + 1,
|
|
last_sync_at=func.now(),
|
|
)
|
|
)
|
|
await db.commit()
|
|
|
|
except IntegrityError as e:
|
|
await db.rollback()
|
|
raise ValidationException(
|
|
"Failed to create observations due to integrity constraint violation"
|
|
) from e
|
|
|
|
logger.debug(
|
|
"Created %d observations in workspace %s",
|
|
len(honcho_documents),
|
|
workspace_name,
|
|
)
|
|
return 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}'."
|
|
)
|
|
vector_store = get_vector_store()
|
|
namespace = vector_store.get_vector_namespace(
|
|
"document",
|
|
workspace_name,
|
|
observer,
|
|
observed,
|
|
)
|
|
vector_deleted = False
|
|
try:
|
|
await vector_store.delete_many(namespace, [existing_doc.id])
|
|
vector_deleted = True
|
|
except Exception:
|
|
existing_doc.deleted_at = datetime.datetime.now(datetime.timezone.utc)
|
|
await db.flush()
|
|
|
|
if vector_deleted:
|
|
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
|
|
|
|
|
|
async def cleanup_soft_deleted_documents(
|
|
db: AsyncSession,
|
|
vector_store: VectorStore,
|
|
batch_size: int = 100,
|
|
older_than_minutes: int = 5,
|
|
) -> int:
|
|
"""
|
|
Cleanup soft-deleted documents by removing their vectors and database records.
|
|
|
|
This function implements a two-phase cleanup process for documents that have been
|
|
soft-deleted (deleted_at is not NULL)
|
|
|
|
Args:
|
|
db: Database session for executing queries
|
|
vector_store: Vector store instance for deleting vectors
|
|
batch_size: Maximum number of documents to process per call (default 100)
|
|
older_than_minutes: Only process documents soft-deleted more than this many
|
|
minutes ago (default 5).
|
|
|
|
Returns:
|
|
Count of documents cleaned up (only those where vector deletion succeeded).
|
|
"""
|
|
cutoff = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(
|
|
minutes=older_than_minutes
|
|
)
|
|
|
|
# Find soft-deleted documents ready for cleanup
|
|
# Use FOR UPDATE SKIP LOCKED to prevent multiple deriver instances from
|
|
# processing the same documents simultaneously
|
|
stmt = (
|
|
select(models.Document)
|
|
.where(models.Document.deleted_at.is_not(None))
|
|
.where(models.Document.deleted_at < cutoff)
|
|
.limit(batch_size)
|
|
.with_for_update(skip_locked=True)
|
|
)
|
|
result = await db.execute(stmt)
|
|
documents = list(result.scalars().all())
|
|
|
|
if not documents:
|
|
return 0
|
|
|
|
# Group by namespace for batch vector deletion
|
|
by_namespace: dict[str, list[str]] = {}
|
|
for doc in documents:
|
|
namespace = vector_store.get_vector_namespace(
|
|
"document",
|
|
doc.workspace_name,
|
|
doc.observer,
|
|
doc.observed,
|
|
)
|
|
by_namespace.setdefault(namespace, []).append(doc.id)
|
|
|
|
# Delete from vector store (per namespace) and track successful deletions
|
|
successfully_deleted_ids: set[str] = set()
|
|
for namespace, ids in by_namespace.items():
|
|
try:
|
|
await vector_store.delete_many(namespace, ids)
|
|
# Only add to successfully_deleted_ids if vector deletion succeeded
|
|
successfully_deleted_ids.update(ids)
|
|
except Exception as e:
|
|
# Log but continue - vectors may already be deleted or namespace may not exist
|
|
logger.warning(f"Failed to delete vectors from {namespace}: {e}")
|
|
|
|
# Only hard delete documents where vector deletion succeeded
|
|
if successfully_deleted_ids:
|
|
await db.execute(
|
|
delete(models.Document).where(
|
|
models.Document.id.in_(successfully_deleted_ids)
|
|
)
|
|
)
|
|
await db.commit()
|
|
logger.debug(
|
|
f"Cleaned up {len(successfully_deleted_ids)} soft-deleted documents"
|
|
)
|
|
return len(successfully_deleted_ids)
|
|
|
|
# No documents were successfully deleted from vector store
|
|
return 0
|