fix dialectic held connection (plastic-labs#477)

* fix: dialectic held connection

* fix: (agent) pre-compute embeddings for agent tools

* fix: (tests) refactor tests to use smaller test db connections

* fix: Embedding client to branch depending on vector store

* fix: reflect dedup-skipped observations in created counts and isolate DB sessions in extract_preferences

* fix: (tests) update tests to match changes

* fix: expunge docs + don't pass in db to query_documents

---------

Co-authored-by: Rajat Ahuja <rahuja445@gmail.com>

Author: VVoruganti

src/crud/__init__.py

@@ -9,13 +9,15 @@
     create_observations,
     delete_document,
     delete_document_by_id,
+    fetch_documents_by_ids,
     get_all_documents,
     get_child_observations,
     get_documents_by_ids,
     get_documents_with_filters,
     query_documents,
     query_documents_most_derived,
     query_documents_recent,
+    query_external_vector_document_ids,
 )
 from .message import (
     create_messages,
@@ -82,13 +84,15 @@
     # Document
     "create_documents",
     "create_observations",
+    "fetch_documents_by_ids",
     "get_all_documents",
     "get_child_observations",
     "get_documents_by_ids",
     "get_documents_with_filters",
     "query_documents",
     "query_documents_most_derived",
     "query_documents_recent",
+    "query_external_vector_document_ids",
     "delete_document",
     "delete_document_by_id",
     # Message

src/crud/document.py

@@ -15,6 +15,7 @@
 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.dependencies import tracked_db
 from src.embedding_client import embedding_client
 from src.exceptions import ResourceNotFoundException, ValidationException
 from src.utils.filter import apply_filter
@@ -190,70 +191,34 @@ async def query_documents_most_derived(
     return result.scalars().all()
 
 
-async def query_documents(
-    db: AsyncSession,
+def _uses_pgvector() -> bool:
+    """Check whether queries should go through pgvector (DB-only) path."""
+    return (
+        settings.VECTOR_STORE.TYPE == "pgvector" or not settings.VECTOR_STORE.MIGRATED
+    )
+
+
+async def query_external_vector_document_ids(
     workspace_name: str,
-    query: str,
-    *,
     observer: str,
     observed: str,
-    filters: dict[str, Any] | None = None,
-    max_distance: float | None = None,
+    embedding: list[float],
     top_k: int = 5,
-    embedding: list[float] | None = None,
-) -> Sequence[models.Document]:
-    """
-    Query documents using semantic similarity.
+    max_distance: float | None = None,
+    filters: dict[str, Any] | None = None,
+) -> list[str] | None:
+    """Query external vector store for document IDs sorted by 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)
+    No DB session needed — safe to call outside a tracked_db scope.
 
     Returns:
-        Sequence of matching documents
+        Ordered list of document IDs on the external-store path,
+        empty list when the external store has no results,
+        or None when the pgvector (DB-only) path should be used instead.
     """
-    # 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
+    if _uses_pgvector():
+        return None
 
-    # Query Postgres directly when using pgvector OR during migration (not yet migrated)
-    # This ensures we use pgvector as source of truth until migration is complete
-    if settings.VECTOR_STORE.TYPE == "pgvector" or not settings.VECTOR_STORE.MIGRATED:
-        stmt = (
-            select(models.Document)
-            .where(models.Document.workspace_name == workspace_name)
-            .where(models.Document.observer == observer)
-            .where(models.Document.observed == observed)
-            .where(models.Document.embedding.isnot(None))
-            .where(models.Document.deleted_at.is_(None))
-        )
-
-        if max_distance is not None:
-            stmt = stmt.where(
-                models.Document.embedding.cosine_distance(embedding) <= max_distance
-            )
-
-        stmt = apply_filter(stmt, models.Document, filters)
-        stmt = stmt.order_by(
-            models.Document.embedding.cosine_distance(embedding)
-        ).limit(top_k)
-
-        result = await db.execute(stmt)
-        return list(result.scalars().all())
-
-    # FALLBACK: Use external vector store (Turbopuffer, LanceDB)
     external_vector_store = get_external_vector_store()
     if external_vector_store is None:
         return []
@@ -262,18 +227,12 @@ async def query_documents(
         "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 external vector store for similar documents with filters applied
     vector_results = await external_vector_store.query(
         namespace,
         embedding,
@@ -285,10 +244,21 @@ async def query_documents(
     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]
+    return [result.id for result in vector_results]
+
+
+async def fetch_documents_by_ids(
+    db: AsyncSession,
+    workspace_name: str,
+    observer: str,
+    observed: str,
+    document_ids: list[str],
+    filters: dict[str, Any] | None = None,
+) -> list[models.Document]:
+    """Fetch documents by IDs, preserving input order. DB-only operation."""
+    if not document_ids:
+        return []
 
-    # Fetch documents from database
     stmt = (
         select(models.Document)
         .where(models.Document.workspace_name == workspace_name)
@@ -297,20 +267,153 @@ async def query_documents(
         .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 [documents[doc_id] for doc_id in document_ids if doc_id in documents]
 
-    return ordered_docs
+
+async def _query_documents_pgvector(
+    db: AsyncSession,
+    workspace_name: str,
+    observer: str,
+    observed: str,
+    embedding: list[float],
+    filters: dict[str, Any] | None,
+    max_distance: float | None,
+    top_k: int,
+) -> list[models.Document]:
+    """pgvector similarity search — pure DB operation."""
+    stmt = (
+        select(models.Document)
+        .where(models.Document.workspace_name == workspace_name)
+        .where(models.Document.observer == observer)
+        .where(models.Document.observed == observed)
+        .where(models.Document.embedding.isnot(None))
+        .where(models.Document.deleted_at.is_(None))
+    )
+
+    if max_distance is not None:
+        stmt = stmt.where(
+            models.Document.embedding.cosine_distance(embedding) <= max_distance
+        )
+
+    stmt = apply_filter(stmt, models.Document, filters)
+    stmt = stmt.order_by(models.Document.embedding.cosine_distance(embedding)).limit(
+        top_k
+    )
+
+    result = await db.execute(stmt)
+    return list(result.scalars().all())
+
+
+async def query_documents(
+    db: AsyncSession | None,
+    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.
+
+    When *db* is provided the caller owns the session lifetime.  When *db* is
+    ``None`` the function opens (and closes) its own short-lived session so that
+    no DB connection is held during external vector-store calls.
+
+    Args:
+        db: Database session, or None to let the function manage its own
+        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
+
+    if _uses_pgvector():
+        # pgvector path — pure DB, open a short session if none provided
+        if db is not None:
+            return await _query_documents_pgvector(
+                db,
+                workspace_name,
+                observer,
+                observed,
+                embedding,
+                filters,
+                max_distance,
+                top_k,
+            )
+        async with tracked_db("query_documents.pgvector") as managed_db:
+            docs = await _query_documents_pgvector(
+                managed_db,
+                workspace_name,
+                observer,
+                observed,
+                embedding,
+                filters,
+                max_distance,
+                top_k,
+            )
+            for doc in docs:
+                managed_db.expunge(doc)
+            return docs
+
+    # External vector store — network call first, DB only for the ID fetch
+    document_ids = await query_external_vector_document_ids(
+        workspace_name=workspace_name,
+        observer=observer,
+        observed=observed,
+        embedding=embedding,
+        top_k=top_k,
+        max_distance=max_distance,
+        filters=filters,
+    )
+
+    if not document_ids:
+        return []
+
+    if db is not None:
+        return await fetch_documents_by_ids(
+            db=db,
+            workspace_name=workspace_name,
+            observer=observer,
+            observed=observed,
+            document_ids=document_ids,
+            filters=filters,
+        )
+    async with tracked_db("query_documents.fetch") as managed_db:
+        docs = await fetch_documents_by_ids(
+            db=managed_db,
+            workspace_name=workspace_name,
+            observer=observer,
+            observed=observed,
+            document_ids=document_ids,
+            filters=filters,
+        )
+        for doc in docs:
+            managed_db.expunge(doc)
+        return docs
 
 
 async def create_documents(
@@ -321,7 +424,7 @@ async def create_documents(
     observer: str,
     observed: str,
     deduplicate: bool = False,
-) -> int:
+) -> list[schemas.DocumentCreate]:
     """
     Create multiple documents with optional duplicate detection.
 
@@ -333,9 +436,11 @@ async def create_documents(
         observed: Name of the observed peer
 
     Returns:
-        Count of new documents
+        List of DocumentCreate schemas that were actually inserted (excludes
+        duplicates and failures).
     """
     honcho_documents: list[models.Document] = []
+    accepted_documents: list[schemas.DocumentCreate] = []
     # Store (document_model, embedding) pairs - IDs aren't available until after commit
     docs_with_embeddings: list[tuple[models.Document, list[float]]] = []
 
@@ -391,6 +496,7 @@ async def create_documents(
             if doc.embedding:
                 new_doc.sync_state = "pending"
             honcho_documents.append(new_doc)
+            accepted_documents.append(doc)
 
             # Track embedding for vector store (ID will be available after commit)
             if doc.embedding:
@@ -489,7 +595,7 @@ async def create_documents(
             "Failed to create documents due to integrity constraint violation"
         ) from e
 
-    return len(honcho_documents)
+    return accepted_documents
 
 
 async def delete_document(