feat(fai): implement incremental document and query index sync (#4710)

tsbhangu · claude · web-flow · commit 1f8e2d0b2dab · 2025-11-04T08:36:58.000-07:00
Co-authored-by: Claude &lt;noreply@anthropic.com&gt;
diff --git a/servers/fai/src/fai/routes/document.py b/servers/fai/src/fai/routes/document.py
@@ -40,7 +40,11 @@
     get_query_index_name,
 )
 from fai.utils.turbopuffer.sync import (
+    delete_documents_from_query_index,
+    delete_documents_from_tpuf,
     sync_document_db_to_tpuf,
+    sync_documents_to_query_index,
+    sync_documents_to_tpuf,
     sync_index_to_target,
 )
 
@@ -81,8 +85,10 @@ async def create_document(
             created_document_ids.append(document_id)
         await db.commit()
         await db.refresh(new_db_document)
-        await sync_document_db_to_tpuf(domain, db)
-        await sync_index_to_target(domain, get_document_index_name(), get_query_index_name())
+        await sync_documents_to_tpuf(domain, created_document_ids, db)
+        await sync_documents_to_query_index(
+            domain, created_document_ids, get_document_index_name(), get_query_index_name()
+        )
         LOGGER.info(f"Indexed document {new_db_document.id} for domain: {domain}")
         return JSONResponse(
             jsonable_encoder([CreateDocumentResponse(document_id=document_id) for document_id in created_document_ids])
@@ -132,8 +138,10 @@ async def batch_create_document(
                 LOGGER.info(f"Created document {document_id} for domain: {domain}")
 
         await db.commit()
-        await sync_document_db_to_tpuf(domain, db)
-        await sync_index_to_target(domain, get_document_index_name(), get_query_index_name())
+        await sync_documents_to_tpuf(domain, created_document_ids, db)
+        await sync_documents_to_query_index(
+            domain, created_document_ids, get_document_index_name(), get_query_index_name()
+        )
 
         return JSONResponse(
             jsonable_encoder([CreateDocumentResponse(document_id=document_id) for document_id in created_document_ids])
@@ -181,8 +189,10 @@ async def update_document(
             db.add(db_document)
             await db.commit()
             await db.refresh(db_document)
-            await sync_document_db_to_tpuf(domain, db)
-            await sync_index_to_target(domain, get_document_index_name(), get_query_index_name())
+            await sync_documents_to_tpuf(domain, [document_id], db)
+            await sync_documents_to_query_index(
+                domain, [document_id], get_document_index_name(), get_query_index_name()
+            )
             LOGGER.info(f"Updated document {document_id} for domain: {domain}")
             return JSONResponse(jsonable_encoder(UpdateDocumentResponse(document=db_document.to_api())))
         return JSONResponse(status_code=404, content=jsonable_encoder({"message": "Document not found"}))
@@ -212,8 +222,10 @@ async def delete_document_by_id(
         if db_document:
             await db.delete(db_document)
             await db.commit()
-            await sync_document_db_to_tpuf(domain, db)
-            await sync_index_to_target(domain, get_document_index_name(), get_query_index_name())
+            await delete_documents_from_tpuf(domain, [body.document_id])
+            await delete_documents_from_query_index(
+                domain, [body.document_id], get_document_index_name(), get_query_index_name()
+            )
             LOGGER.info(f"Deleted document {body.document_id} for domain: {domain}")
             return JSONResponse(jsonable_encoder(DeleteDocumentResponse(success=True)))
         return JSONResponse(jsonable_encoder(DeleteDocumentResponse(success=False)))
@@ -237,6 +249,7 @@ async def batch_delete_document(
 ) -> JSONResponse:
     try:
         deleted_count = 0
+        deleted_document_ids = []
 
         for document in body:
             db_document = await db.execute(
@@ -245,13 +258,16 @@ async def batch_delete_document(
             db_document = db_document.scalar_one_or_none()
             if db_document:
                 await db.delete(db_document)
+                deleted_document_ids.append(document.document_id)
                 deleted_count += 1
                 LOGGER.info(f"Deleted document {document.document_id} for domain: {domain}")
 
         if deleted_count > 0:
             await db.commit()
-            await sync_document_db_to_tpuf(domain, db)
-            await sync_index_to_target(domain, get_document_index_name(), get_query_index_name())
+            await delete_documents_from_tpuf(domain, deleted_document_ids)
+            await delete_documents_from_query_index(
+                domain, deleted_document_ids, get_document_index_name(), get_query_index_name()
+            )
             return JSONResponse(jsonable_encoder(DeleteDocumentResponse(success=True)))
         else:
             return JSONResponse(jsonable_encoder(DeleteDocumentResponse(success=False)))
diff --git a/servers/fai/src/fai/utils/turbopuffer/sync.py b/servers/fai/src/fai/utils/turbopuffer/sync.py
@@ -39,6 +39,56 @@ def prefixed_id(namespace: str, original_id: str, max_len: int = 64) -> str:
     return f"{short_ns}:{hashed}"
 
 
+async def sync_documents_to_tpuf(domain: str, document_ids: list[str], db: AsyncSession) -> None:
+    if not document_ids:
+        LOGGER.info("No document IDs provided for sync, skipping")
+        return
+
+    documents = await db.execute(select(DocumentDb).where(DocumentDb.domain == domain, DocumentDb.id.in_(document_ids)))
+    documents = documents.scalars().all()
+
+    if not documents:
+        LOGGER.warning(f"No documents found for IDs {document_ids} in domain {domain}")
+        return
+
+    async with AsyncOpenAI(api_key=VARIABLES.OPENAI_API_KEY) as openai_client:
+        async with AsyncTurbopuffer(
+            region=CONFIG.TURBOPUFFER_DEFAULT_REGION,
+            api_key=VARIABLES.TURBOPUFFER_API_KEY,
+        ) as tpuf_client:
+            target_namespace_id = get_tpuf_namespace(domain, get_document_index_name())
+            target_ns = tpuf_client.namespace(target_namespace_id)
+
+            tbuf_records = []
+            for document in documents:
+                tbuf_records.append(await document.to_tpuf_record(openai_client))
+
+            await target_ns.write(
+                upsert_rows=[jsonable_encoder(record) for record in tbuf_records],
+                distance_metric="cosine_distance",
+                schema=get_data_index_tpuf_schema(),
+            )
+            LOGGER.info(f"Upserted {len(documents)} documents to {target_namespace_id}")
+
+
+async def delete_documents_from_tpuf(domain: str, document_ids: list[str]) -> None:
+    if not document_ids:
+        LOGGER.info("No document IDs provided for deletion, skipping")
+        return
+
+    async with AsyncTurbopuffer(
+        region=CONFIG.TURBOPUFFER_DEFAULT_REGION,
+        api_key=VARIABLES.TURBOPUFFER_API_KEY,
+    ) as tpuf_client:
+        target_namespace_id = get_tpuf_namespace(domain, get_document_index_name())
+        target_ns = tpuf_client.namespace(target_namespace_id)
+
+        for document_id in document_ids:
+            await target_ns.write(delete_by_filter=["id", "Eq", document_id])
+
+        LOGGER.info(f"Deleted {len(document_ids)} documents from {target_namespace_id}")
+
+
 async def sync_document_db_to_tpuf(domain: str, db: AsyncSession) -> None:
     documents = await db.execute(select(DocumentDb).where(DocumentDb.domain == domain))
     documents = documents.scalars().all()
@@ -118,6 +168,68 @@ async def sync_slack_context_db_to_tpuf(domain: str, db: AsyncSession) -> None:
             LOGGER.info(f"Wrote {len(slack_contexts)} slack contexts to {target_namespace_id}")
 
 
+async def sync_documents_to_query_index(
+    domain: str, document_ids: list[str], source_index_name: str, target_index_name: str
+) -> None:
+    if not document_ids:
+        LOGGER.info("No document IDs provided for query index sync, skipping")
+        return
+
+    source_namespace_id = get_tpuf_namespace(domain, source_index_name)
+    target_namespace_id = get_tpuf_namespace(domain, target_index_name)
+
+    async with AsyncTurbopuffer(
+        region=CONFIG.TURBOPUFFER_DEFAULT_REGION,
+        api_key=VARIABLES.TURBOPUFFER_API_KEY,
+    ) as tpuf_client:
+        source_ns = tpuf_client.namespace(source_namespace_id)
+        target_ns = tpuf_client.namespace(target_namespace_id)
+
+        for document_id in document_ids:
+            prefixed_doc_id = prefixed_id(source_namespace_id, document_id)
+            await target_ns.write(delete_by_filter=["id", "Eq", prefixed_doc_id])
+
+        prefixed_rows = []
+        for document_id in document_ids:
+            result = await source_ns.query(filters=("id", "Eq", document_id), top_k=1, include_attributes=True)
+
+            if result.rows:
+                row = result.rows[0]
+                new_row = Row.from_dict(row.model_dump())
+                new_row.id = prefixed_id(source_namespace_id, document_id)
+                new_row.source = source_index_name
+                prefixed_rows.append(new_row)
+
+        if prefixed_rows:
+            await target_ns.write(
+                upsert_rows=prefixed_rows, distance_metric="cosine_distance", schema=get_query_index_tpuf_schema()
+            )
+            LOGGER.info(f"Synced {len(prefixed_rows)} documents to query index {target_namespace_id}")
+
+
+async def delete_documents_from_query_index(
+    domain: str, document_ids: list[str], source_index_name: str, target_index_name: str
+) -> None:
+    if not document_ids:
+        LOGGER.info("No document IDs provided for query index deletion, skipping")
+        return
+
+    source_namespace_id = get_tpuf_namespace(domain, source_index_name)
+    target_namespace_id = get_tpuf_namespace(domain, target_index_name)
+
+    async with AsyncTurbopuffer(
+        region=CONFIG.TURBOPUFFER_DEFAULT_REGION,
+        api_key=VARIABLES.TURBOPUFFER_API_KEY,
+    ) as tpuf_client:
+        target_ns = tpuf_client.namespace(target_namespace_id)
+
+        for document_id in document_ids:
+            prefixed_doc_id = prefixed_id(source_namespace_id, document_id)
+            await target_ns.write(delete_by_filter=["id", "Eq", prefixed_doc_id])
+
+        LOGGER.info(f"Deleted {len(document_ids)} documents from query index {target_namespace_id}")
+
+
 async def sync_index_to_target(domain: str, source_index_name: str, target_index_name: str) -> None:
     source_namespace_id = get_tpuf_namespace(domain, source_index_name)
     target_namespace_id = get_tpuf_namespace(domain, target_index_name)
