Agent Zero
abbe7a18fc
- P0: hooks.py 3-tuple fix, trigger_dispatcher Contract, contacts/plugin unregister_actions_by_owner - P0: 5 test files — check_permission mocks removed, hardcoded DB credential → env var - P1: attachment_service DmsFile via Contract helper, restore_registry/history_hooks dedup - P1: mail/plugin restore unregister, mcp_client datetime.now(UTC), saved_views/filters patterns - P1: ProtectedRoute fail-closed, 13 test assertion fixes (bcrypt, DB-URLs, SECRET_KEYs) - P2: deprecated notifications → post_system_message (3 files), forgejo Base, report_generator lazy import - P2: webhooks permissions, deps.py/roles.py plugin perms removed, import_export default - P2: address/tags/entity_links patterns removed, worker.py Contract-Umgehungen fixed - P2: 28 frontend TODOs (hardcoded constants, deprecated notification API) - P3: dead code, duplicates, deprecated imports, private attr, __import__ inline - P3: 8 frontend TODOs (LucideIcons, inline styles, XSS, i18n) - ruff: 838 → 0 (612 auto-fix + 246 manual + 27 F821 regression fix) - F821: 30 → 0 (AutomationDefinition, DmsFile, user_id, Path, Any, String) - Contract-Umgehungen: 2 neue gefunden (worker.py:169, worker.py:280) und gefixt
531 lines
19 KiB
Python
531 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.embedding import generate_embedding
|
|
from app.plugins.builtins.unified_search.text_extraction import extract_text_from_file
|
|
|
|
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.chunking import chunk_text
|
|
from app.plugins.builtins.unified_search.embedding import generate_embedding
|
|
from app.plugins.builtins.unified_search.text_extraction import extract_text_from_file
|
|
|
|
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 # noqa: E402
|
|
|
|
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)
|