Agent Zero
b15a62bec6
- Removed: app/routes/companies.py, app/services/company_service.py, app/models/company.py, app/schemas/company.py - Updated: main.py, routes/__init__.py, models/__init__.py (removed company imports) - Updated: action_mapper.py (company intents → contact intents, /api/v1/companies → /api/v1/contacts) - Updated: workflows/engine.py (company.created → contact.created) - Updated: worker.py (removed index_company) - Updated: roles.py, permission_registry.py, permissions.py, deps.py (companies: → contacts:) - Updated: permission_registry.py CORE_FIELD_DEFINITIONS (old field names → unified contact fields) - Updated: address model/schema/service (entity_type company → contact only) - Updated: ai_copilot_service.py (_exec_companies removed, _exec_contacts extended) - Updated: ai_proactive services/jobs/context_tools (Company → Contact, company_contacts → contact_persons) - Updated: unified_search jobs.py (removed index_company), query_understanding.py - Updated: calendar/models.py, addresses.py docstrings - Updated: conftest.py (Company → Contact, removed companies/company_contacts from TRUNCATE) - Updated: test_unified_search.py (index_company → index_contact)
224 lines
7.7 KiB
Python
224 lines
7.7 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, 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)
|
|
|
|
|
|
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_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",
|
|
"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",
|
|
"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")
|