247 lines
8.5 KiB
Python
247 lines
8.5 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
|
||
|
|
|
||
|
|
|
||
|
|
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)
|
||
|
|
|
||
|
|
|
||
|
|
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)
|
||
|
|
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)
|
||
|
|
|
||
|
|
|
||
|
|
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)
|
||
|
|
|
||
|
|
|
||
|
|
async def index_company(ctx: dict[str, Any], company_id: str) -> None:
|
||
|
|
"""Index a company: 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(company_id)
|
||
|
|
result = await db.execute(
|
||
|
|
text("SELECT tenant_id FROM companies WHERE id = :cid"),
|
||
|
|
{"cid": eid},
|
||
|
|
)
|
||
|
|
row = result.mappings().first()
|
||
|
|
if not row:
|
||
|
|
return
|
||
|
|
await index_entity("company", eid, row["tenant_id"], db)
|
||
|
|
except Exception:
|
||
|
|
logger.exception("Failed to index company %s", company_id)
|
||
|
|
|
||
|
|
|
||
|
|
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)
|
||
|
|
|
||
|
|
|
||
|
|
async def reindex(ctx: dict[str, Any], entity_type: str) -> None:
|
||
|
|
"""Reindex all entities of a given type with pagination."""
|
||
|
|
from sqlalchemy import text
|
||
|
|
from app.plugins.builtins.unified_search.embedding import index_entity
|
||
|
|
|
||
|
|
table_map = {
|
||
|
|
"contact": "contacts",
|
||
|
|
"company": "companies",
|
||
|
|
"mail": "mails",
|
||
|
|
"file": "files",
|
||
|
|
"event": "calendar_entries",
|
||
|
|
}
|
||
|
|
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:
|
||
|
|
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,
|
||
|
|
)
|
||
|
|
except Exception:
|
||
|
|
logger.exception("Reindex failed for %s/%s", entity_type, row["id"])
|
||
|
|
offset += BATCH_SIZE
|
||
|
|
|
||
|
|
logger.info("Reindex complete for %s", entity_type)
|
||
|
|
|
||
|
|
|
||
|
|
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
|
||
|
|
|
||
|
|
table_map = {
|
||
|
|
"contact": "contacts",
|
||
|
|
"company": "companies",
|
||
|
|
"mail": "mails",
|
||
|
|
"file": "files",
|
||
|
|
"event": "calendar_entries",
|
||
|
|
}
|
||
|
|
|
||
|
|
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"])
|
||
|
|
except Exception:
|
||
|
|
logger.exception("Batch query failed for %s", etype)
|
||
|
|
|
||
|
|
logger.info("Embedding batch job complete")
|