Agent Zero
3d9b76cea4
Check Cross-Plugin Imports / check (push) Has been cancelled
- SPIKE-E: FTS+Vector+Permission benchmark on 10k records (all <30ms) - E-PROV: supports_fts/vector/rag/graph capability flags on all providers - E-FTS/VEC: All 11 providers refactored to BaseSearchProvider with permission filtering - E-PERM: Over-fetch strategy for vector+permission (15x faster than ANY() filter) - E-FUSE: rrf_fusion_multi() for N-way RRF over FTS+Vector+RAG+Graph - E-LLM: Query understanding cleaned up to use central llm_complete() - E-CHUNK: Document chunking module + document_chunks table with HNSW index - E-EMB: Chunk embedding ARQ jobs (index_file_chunks, reindex_chunks) - E-RAG: RAG retrieval via FileSearchProvider.search_rag() - E-GRAPH: GraphRAG BFS traversal via GraphRAGSearchProvider.search_graph() - E-IX-EVT: Auto-indexing via outbox events + delete/cleanup handlers - E-IX-RE: Batch reindex with progress tracking + reindex_all job - E-DATA-LIFE: Lifecycle module (remove/rebuild/restore/correct) + API endpoints - E-K-MEM: AgentMemorySearchProvider - E-P-AI: AIChatSearchProvider - E-P-WF: WorkflowSearchProvider - E-P-COMM: ConversationSearchProvider verified (already on BaseSearchProvider) - E-API: Filter params (date_from/to, tags, sort) + /facets endpoint - E-TOOL: unified_search AI tool registered in ToolRegistry - E-MCP: Search tool in MCP server with normal RBAC/tenant checks - E-UI-CMD: CommandPalette (Cmd+K) with debounced search + recent searches - E-UI-FAC: SearchFacets, SearchResultCard, SavedSearches components - E-TEST: 40 new tests in test_unified_search_phase_e.py (105 total green) - E-DOC: api-documentation.md, plugin-development-guide.md, test-strategy.md updated 105 tests passing, TypeScript clean.
523 lines
19 KiB
Python
523 lines
19 KiB
Python
"""ARQ background jobs for the Unified Search plugin."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import uuid
|
|
from typing import Any
|
|
|
|
from app.core.db import get_session_factory
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
BATCH_SIZE = 100
|
|
|
|
# Entity type -> table name mapping (shared by multiple jobs)
|
|
_TABLE_MAP: dict[str, str] = {
|
|
"contact": "contacts",
|
|
"mail": "mails",
|
|
"file": "files",
|
|
"event": "calendar_entries",
|
|
}
|
|
|
|
|
|
def _parse_id(id_str: str) -> uuid.UUID:
|
|
"""Parse a string to UUID."""
|
|
if isinstance(id_str, uuid.UUID):
|
|
return id_str
|
|
return uuid.UUID(str(id_str))
|
|
|
|
|
|
async def index_mails(ctx: dict[str, Any], mail_ids: list[str]) -> None:
|
|
"""Index mails: generate and store embeddings."""
|
|
from app.plugins.builtins.unified_search.embedding import index_entity
|
|
|
|
factory = get_session_factory()
|
|
async with factory() as db:
|
|
for mail_id in mail_ids:
|
|
try:
|
|
eid = _parse_id(mail_id)
|
|
# tenant_id is derived from the mail itself
|
|
from sqlalchemy import text
|
|
result = await db.execute(
|
|
text("SELECT tenant_id FROM mails WHERE id = :mid"),
|
|
{"mid": eid},
|
|
)
|
|
row = result.mappings().first()
|
|
if not row:
|
|
continue
|
|
tenant_id = row["tenant_id"]
|
|
await index_entity("mail", eid, tenant_id, db)
|
|
except Exception:
|
|
logger.exception("Failed to index mail %s", mail_id)
|
|
try:
|
|
await db.rollback()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
async def index_file(ctx: dict[str, Any], file_id: str) -> None:
|
|
"""Index a file: extract text, store content_text, generate embedding."""
|
|
from sqlalchemy import text
|
|
from app.plugins.builtins.unified_search.text_extraction import extract_text_from_file
|
|
from app.plugins.builtins.unified_search.embedding import generate_embedding
|
|
|
|
factory = get_session_factory()
|
|
async with factory() as db:
|
|
try:
|
|
eid = _parse_id(file_id)
|
|
result = await db.execute(
|
|
text("SELECT tenant_id, storage_path, mime_type FROM files WHERE id = :fid"),
|
|
{"fid": eid},
|
|
)
|
|
row = result.mappings().first()
|
|
if not row:
|
|
logger.warning("File not found: %s", file_id)
|
|
return
|
|
|
|
tenant_id = row["tenant_id"]
|
|
storage_path = row["storage_path"]
|
|
mime_type = row["mime_type"]
|
|
|
|
# Extract text
|
|
content_text = await extract_text_from_file(storage_path, mime_type)
|
|
|
|
# Store content_text
|
|
await db.execute(
|
|
text("UPDATE files SET content_text = :ct WHERE id = :fid"),
|
|
{"ct": content_text, "fid": eid},
|
|
)
|
|
await db.commit()
|
|
|
|
# Generate embedding from extracted text + filename
|
|
result_name = await db.execute(
|
|
text("SELECT name FROM files WHERE id = :fid"),
|
|
{"fid": eid},
|
|
)
|
|
name_row = result_name.mappings().first()
|
|
name = name_row["name"] if name_row else ""
|
|
embedding_text = f"{name} {content_text[:5000]}"
|
|
|
|
if embedding_text.strip():
|
|
embedding = await generate_embedding(embedding_text, db=db, tenant_id=tenant_id)
|
|
if embedding:
|
|
await db.execute(
|
|
text("UPDATE files SET embedding = cast(:emb AS vector) WHERE id = :fid"),
|
|
{"emb": str(embedding), "fid": eid},
|
|
)
|
|
await db.commit()
|
|
|
|
logger.info("Indexed file %s", file_id)
|
|
except Exception:
|
|
logger.exception("Failed to index file %s", file_id)
|
|
try:
|
|
await db.rollback()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
async def index_contact(ctx: dict[str, Any], contact_id: str) -> None:
|
|
"""Index a contact: generate and store embedding."""
|
|
from sqlalchemy import text
|
|
from app.plugins.builtins.unified_search.embedding import index_entity
|
|
|
|
factory = get_session_factory()
|
|
async with factory() as db:
|
|
try:
|
|
eid = _parse_id(contact_id)
|
|
result = await db.execute(
|
|
text("SELECT tenant_id FROM contacts WHERE id = :cid"),
|
|
{"cid": eid},
|
|
)
|
|
row = result.mappings().first()
|
|
if not row:
|
|
return
|
|
await index_entity("contact", eid, row["tenant_id"], db)
|
|
except Exception:
|
|
logger.exception("Failed to index contact %s", contact_id)
|
|
try:
|
|
await db.rollback()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
async def index_event(ctx: dict[str, Any], event_id: str) -> None:
|
|
"""Index a calendar event: generate and store embedding."""
|
|
from sqlalchemy import text
|
|
from app.plugins.builtins.unified_search.embedding import index_entity
|
|
|
|
factory = get_session_factory()
|
|
async with factory() as db:
|
|
try:
|
|
eid = _parse_id(event_id)
|
|
result = await db.execute(
|
|
text("SELECT tenant_id FROM calendar_entries WHERE id = :eid"),
|
|
{"eid": eid},
|
|
)
|
|
row = result.mappings().first()
|
|
if not row:
|
|
return
|
|
await index_entity("event", eid, row["tenant_id"], db)
|
|
except Exception:
|
|
logger.exception("Failed to index event %s", event_id)
|
|
try:
|
|
await db.rollback()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
async def reindex(ctx: dict[str, Any], entity_type: str) -> None:
|
|
"""Reindex all entities of a given type with pagination.
|
|
|
|
Clears existing embeddings before re-indexing, tracks progress,
|
|
and continues on individual failures.
|
|
"""
|
|
from sqlalchemy import text
|
|
from app.plugins.builtins.unified_search.embedding import index_entity
|
|
|
|
table = _TABLE_MAP.get(entity_type)
|
|
if not table:
|
|
logger.warning("Unknown entity_type for reindex: %s", entity_type)
|
|
return
|
|
|
|
factory = get_session_factory()
|
|
async with factory() as db:
|
|
# Clear existing embeddings before re-indexing
|
|
try:
|
|
await db.execute(
|
|
text(f"UPDATE {table} SET embedding = NULL WHERE deleted_at IS NULL"),
|
|
)
|
|
await db.commit()
|
|
except Exception:
|
|
logger.exception("Failed to clear embeddings for %s", entity_type)
|
|
try:
|
|
await db.rollback()
|
|
except Exception:
|
|
pass
|
|
|
|
total_indexed = 0
|
|
total_failed = 0
|
|
offset = 0
|
|
while True:
|
|
result = await db.execute(
|
|
text(
|
|
f"SELECT id, tenant_id FROM {table} "
|
|
f"WHERE deleted_at IS NULL ORDER BY created_at LIMIT :lim OFFSET :off"
|
|
),
|
|
{"lim": BATCH_SIZE, "off": offset},
|
|
)
|
|
rows = result.mappings().all()
|
|
if not rows:
|
|
break
|
|
for row in rows:
|
|
try:
|
|
await index_entity(
|
|
entity_type,
|
|
row["id"],
|
|
row["tenant_id"],
|
|
db,
|
|
)
|
|
total_indexed += 1
|
|
except Exception:
|
|
logger.exception("Reindex failed for %s/%s", entity_type, row["id"])
|
|
total_failed += 1
|
|
try:
|
|
await db.rollback()
|
|
except Exception:
|
|
pass
|
|
# Progress tracking every BATCH_SIZE entities
|
|
logger.info(
|
|
"Reindex progress for %s: %d indexed, %d failed (offset=%d)",
|
|
entity_type, total_indexed, total_failed, offset,
|
|
)
|
|
offset += BATCH_SIZE
|
|
|
|
logger.info(
|
|
"Reindex complete for %s: %d indexed, %d failed",
|
|
entity_type, total_indexed, total_failed,
|
|
)
|
|
|
|
|
|
async def reindex_all(ctx: dict[str, Any]) -> None:
|
|
"""Reindex all entity types in sequence, including file chunks."""
|
|
entity_types = ["contact", "mail", "file", "event"]
|
|
for etype in entity_types:
|
|
try:
|
|
await reindex(ctx, etype)
|
|
except Exception:
|
|
logger.exception("reindex_all: reindex failed for %s", etype)
|
|
# For files, also re-index chunks
|
|
if etype == "file":
|
|
try:
|
|
from sqlalchemy import text
|
|
factory = get_session_factory()
|
|
async with factory() as db:
|
|
result = await db.execute(
|
|
text("SELECT id FROM files WHERE deleted_at IS NULL"),
|
|
)
|
|
rows = result.mappings().all()
|
|
for row in rows:
|
|
try:
|
|
await index_file_chunks(ctx, str(row["id"]))
|
|
except Exception:
|
|
logger.exception("reindex_all: chunk re-index failed for file %s", row["id"])
|
|
except Exception:
|
|
logger.exception("reindex_all: chunk re-index batch failed")
|
|
logger.info("reindex_all complete")
|
|
|
|
|
|
async def delete_entity_index(ctx: dict[str, Any], entity_type: str, entity_id: str) -> None:
|
|
"""Delete an entity's embedding (set embedding=NULL).
|
|
|
|
Logs the action to SearchIndexLog on success.
|
|
"""
|
|
from sqlalchemy import text
|
|
from app.plugins.builtins.unified_search.models import SearchIndexLog
|
|
|
|
table = _TABLE_MAP.get(entity_type)
|
|
if not table:
|
|
logger.warning("delete_entity_index: unknown entity_type=%s", entity_type)
|
|
return
|
|
|
|
factory = get_session_factory()
|
|
async with factory() as db:
|
|
try:
|
|
eid = _parse_id(entity_id)
|
|
# Get tenant_id from the entity
|
|
result = await db.execute(
|
|
text(f"SELECT tenant_id FROM {table} WHERE id = :eid"),
|
|
{"eid": eid},
|
|
)
|
|
row = result.mappings().first()
|
|
tenant_id = row["tenant_id"] if row else None
|
|
|
|
await db.execute(
|
|
text(
|
|
f"UPDATE {table} SET embedding = NULL, indexed_at = NULL "
|
|
f"WHERE id = :eid"
|
|
),
|
|
{"eid": eid},
|
|
)
|
|
await db.commit()
|
|
|
|
# Log to SearchIndexLog
|
|
if tenant_id:
|
|
log_entry = SearchIndexLog(
|
|
tenant_id=tenant_id,
|
|
entity_type=entity_type,
|
|
entity_id=eid,
|
|
action="delete",
|
|
status="success",
|
|
)
|
|
db.add(log_entry)
|
|
await db.commit()
|
|
|
|
logger.info("Deleted index for %s/%s", entity_type, entity_id)
|
|
except Exception:
|
|
logger.exception("Failed to delete index for %s/%s", entity_type, entity_id)
|
|
try:
|
|
await db.rollback()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
async def delete_file_chunks(ctx: dict[str, Any], file_id: str) -> None:
|
|
"""Delete all document_chunks for a file."""
|
|
from sqlalchemy import text
|
|
|
|
factory = get_session_factory()
|
|
async with factory() as db:
|
|
try:
|
|
eid = _parse_id(file_id)
|
|
await db.execute(
|
|
text("DELETE FROM document_chunks WHERE file_id = :fid"),
|
|
{"fid": eid},
|
|
)
|
|
await db.commit()
|
|
logger.info("Deleted chunks for file %s", file_id)
|
|
except Exception:
|
|
logger.exception("Failed to delete chunks for file %s", file_id)
|
|
try:
|
|
await db.rollback()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
async def retry_failed_index(ctx: dict[str, Any], entity_type: str, entity_id: str) -> None:
|
|
"""Retry a failed index operation."""
|
|
from sqlalchemy import text
|
|
from app.plugins.builtins.unified_search.embedding import index_entity
|
|
|
|
table = _TABLE_MAP.get(entity_type)
|
|
if not table:
|
|
logger.warning("retry_failed_index: unknown entity_type=%s", entity_type)
|
|
return
|
|
|
|
factory = get_session_factory()
|
|
async with factory() as db:
|
|
try:
|
|
eid = _parse_id(entity_id)
|
|
result = await db.execute(
|
|
text(f"SELECT tenant_id FROM {table} WHERE id = :eid"),
|
|
{"eid": eid},
|
|
)
|
|
row = result.mappings().first()
|
|
if not row:
|
|
logger.warning("retry_failed_index: entity not found %s/%s", entity_type, entity_id)
|
|
return
|
|
await index_entity(entity_type, eid, row["tenant_id"], db)
|
|
except Exception:
|
|
logger.exception("retry_failed_index failed for %s/%s", entity_type, entity_id)
|
|
try:
|
|
await db.rollback()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
async def embedding_batch(ctx: dict[str, Any]) -> None:
|
|
"""Periodic job: find entities without embeddings and index them."""
|
|
from sqlalchemy import text
|
|
from app.plugins.builtins.unified_search.embedding import index_entity
|
|
|
|
factory = get_session_factory()
|
|
async with factory() as db:
|
|
for etype, table in _TABLE_MAP.items():
|
|
try:
|
|
result = await db.execute(
|
|
text(
|
|
f"SELECT id, tenant_id FROM {table} "
|
|
f"WHERE deleted_at IS NULL AND embedding IS NULL "
|
|
f"LIMIT :lim"
|
|
),
|
|
{"lim": BATCH_SIZE},
|
|
)
|
|
rows = result.mappings().all()
|
|
for row in rows:
|
|
try:
|
|
await index_entity(etype, row["id"], row["tenant_id"], db)
|
|
except Exception:
|
|
logger.exception("Batch index failed for %s/%s", etype, row["id"])
|
|
try:
|
|
await db.rollback()
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
logger.exception("Batch query failed for %s", etype)
|
|
try:
|
|
await db.rollback()
|
|
except Exception:
|
|
pass
|
|
|
|
logger.info("Embedding batch job complete")
|
|
|
|
|
|
async def index_file_chunks(ctx: dict[str, Any], file_id: str) -> None:
|
|
"""Extract text from a file, chunk it, generate embeddings, and store in document_chunks."""
|
|
from sqlalchemy import text
|
|
from app.plugins.builtins.unified_search.text_extraction import extract_text_from_file
|
|
from app.plugins.builtins.unified_search.chunking import chunk_text
|
|
from app.plugins.builtins.unified_search.embedding import generate_embedding
|
|
|
|
factory = get_session_factory()
|
|
async with factory() as db:
|
|
try:
|
|
eid = _parse_id(file_id)
|
|
result = await db.execute(
|
|
text("SELECT tenant_id, storage_path, mime_type, name FROM files WHERE id = :fid"),
|
|
{"fid": eid},
|
|
)
|
|
row = result.mappings().first()
|
|
if not row:
|
|
logger.warning("File not found for chunking: %s", file_id)
|
|
return
|
|
|
|
tenant_id = row["tenant_id"]
|
|
storage_path = row["storage_path"]
|
|
mime_type = row["mime_type"]
|
|
name = row.get("name", "") or ""
|
|
|
|
# Extract text from file
|
|
content_text = await extract_text_from_file(storage_path, mime_type)
|
|
if not content_text.strip():
|
|
logger.debug("No text extracted from file %s, skipping chunking", file_id)
|
|
return
|
|
|
|
# Store content_text on the file record
|
|
await db.execute(
|
|
text("UPDATE files SET content_text = :ct WHERE id = :fid"),
|
|
{"ct": content_text, "fid": eid},
|
|
)
|
|
await db.commit()
|
|
|
|
# Chunk the text (include filename for context)
|
|
full_text = f"{name}\n{content_text}"
|
|
chunks = chunk_text(full_text, chunk_size=1000, overlap=200)
|
|
if not chunks:
|
|
logger.debug("No chunks generated for file %s", file_id)
|
|
return
|
|
|
|
# Delete existing chunks for this file (idempotent re-index)
|
|
await db.execute(
|
|
text("DELETE FROM document_chunks WHERE file_id = :fid"),
|
|
{"fid": eid},
|
|
)
|
|
|
|
# Generate embeddings and insert chunks in batches
|
|
for chunk in chunks:
|
|
try:
|
|
embedding = await generate_embedding(
|
|
chunk["chunk_text"], db=db, tenant_id=tenant_id
|
|
)
|
|
if embedding:
|
|
await db.execute(
|
|
text(
|
|
"INSERT INTO document_chunks "
|
|
"(tenant_id, file_id, chunk_index, chunk_text, chunk_hash, embedding) "
|
|
"VALUES (:tid, :fid, :idx, :ctext, :chash, cast(:emb AS vector))"
|
|
),
|
|
{
|
|
"tid": tenant_id,
|
|
"fid": eid,
|
|
"idx": chunk["chunk_index"],
|
|
"ctext": chunk["chunk_text"],
|
|
"chash": chunk["chunk_hash"],
|
|
"emb": str(embedding),
|
|
},
|
|
)
|
|
else:
|
|
logger.warning("Empty embedding for chunk %d of file %s", chunk["chunk_index"], file_id)
|
|
except Exception:
|
|
logger.exception("Failed to embed chunk %d for file %s", chunk["chunk_index"], file_id)
|
|
|
|
await db.commit()
|
|
logger.info("Indexed %d chunks for file %s", len(chunks), file_id)
|
|
|
|
except Exception:
|
|
logger.exception("Failed to index file chunks for %s", file_id)
|
|
try:
|
|
await db.rollback()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
async def reindex_chunks(ctx: dict[str, Any], file_id: str) -> None:
|
|
"""Re-chunk and re-embed a file — delegates to index_file_chunks."""
|
|
await index_file_chunks(ctx, file_id)
|
|
|
|
|
|
# ── Register all job functions with the job registry ──────────────────────────
|
|
from app.core.job_registry import register_job
|
|
|
|
register_job("index_mails", index_mails)
|
|
register_job("index_file", index_file)
|
|
register_job("index_contact", index_contact)
|
|
register_job("index_event", index_event)
|
|
register_job("reindex", reindex)
|
|
register_job("reindex_all", reindex_all)
|
|
register_job("embedding_batch", embedding_batch)
|
|
register_job("delete_entity_index", delete_entity_index)
|
|
register_job("delete_file_chunks", delete_file_chunks)
|
|
register_job("retry_failed_index", retry_failed_index)
|
|
register_job("index_file_chunks", index_file_chunks)
|
|
register_job("reindex_chunks", reindex_chunks)
|