agentic.email/backend/app/services/thread_resolver.py
2026-10-06 03:14:53 +05:30

267 lines
8 KiB
Python

import logging
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from bson import ObjectId
from app.services.message_ids import (
canonical_message_id,
folded_message_id,
parse_message_id_list,
)
from app.services.subject_normalize import normalize_subject
logger = logging.getLogger(__name__)
FALLBACK_WINDOW = timedelta(days=7)
@dataclass(frozen=True)
class ThreadResolution:
thread_id: ObjectId | None
parent_message_id: ObjectId | None
thread_relation: str
thread_resolution_method: str
@dataclass(frozen=True)
class KnownMessage:
id: ObjectId
thread_id: ObjectId
internet_message_id: str
account_id: str
sent_at: datetime | None
received_at: datetime | None
created_at: datetime
participants: frozenset[str]
normalized_subject: str
def _message_timestamp(message: KnownMessage) -> datetime:
return message.sent_at or message.received_at or message.created_at
def resolve_thread_from_known(
*,
account_id: str,
in_reply_to: str | None,
references: list[str],
subject: str,
participants: set[str],
known_by_internet_id: dict[str, KnownMessage],
known_threads: dict[ObjectId, list[KnownMessage]] | None = None,
now: datetime | None = None,
) -> ThreadResolution:
scoped: dict[str, KnownMessage] = {}
for key, value in known_by_internet_id.items():
if value.account_id != account_id:
continue
fold_key = folded_message_id(key) or folded_message_id(value.internet_message_id)
if fold_key:
scoped[fold_key] = value
in_reply = canonical_message_id(in_reply_to)
in_reply_folded = folded_message_id(in_reply)
if in_reply_folded and in_reply_folded in scoped:
parent = scoped[in_reply_folded]
return ThreadResolution(
thread_id=parent.thread_id,
parent_message_id=parent.id,
thread_relation="reply",
thread_resolution_method="in_reply_to",
)
reference_ids = [
folded_message_id(value)
for value in references
if folded_message_id(value)
]
for reference in reversed(reference_ids):
if reference in scoped:
parent = scoped[reference]
return ThreadResolution(
thread_id=parent.thread_id,
parent_message_id=parent.id,
thread_relation="reply",
thread_resolution_method="references",
)
fallback = _conservative_fallback(
account_id=account_id,
subject=subject,
participants=participants,
known_by_internet_id=scoped,
known_threads=known_threads or {},
now=now or datetime.now(timezone.utc),
)
if fallback is not None:
return fallback
return ThreadResolution(
thread_id=None,
parent_message_id=None,
thread_relation="new_thread",
thread_resolution_method="new_thread",
)
def _conservative_fallback(
*,
account_id: str,
subject: str,
participants: set[str],
known_by_internet_id: dict[str, KnownMessage],
known_threads: dict[ObjectId, list[KnownMessage]],
now: datetime,
) -> ThreadResolution | None:
normalized = normalize_subject(subject)
if not normalized or not participants:
return None
candidates: dict[ObjectId, KnownMessage] = {}
for message in known_by_internet_id.values():
if message.normalized_subject != normalized:
continue
if not participants.intersection(message.participants):
continue
timestamp = _message_timestamp(message)
if timestamp < now - FALLBACK_WINDOW:
continue
candidates[message.thread_id] = message
if len(candidates) != 1:
return None
parent = next(iter(candidates.values()))
return ThreadResolution(
thread_id=parent.thread_id,
parent_message_id=parent.id,
thread_relation="unknown",
thread_resolution_method="fallback",
)
def participants_from_parsed(parsed: dict) -> set[str]:
addresses: set[str] = set()
for key in ("sender",):
sender = parsed.get(key)
if isinstance(sender, dict) and sender.get("address"):
addresses.add(sender["address"].lower())
for key in ("to", "cc", "reply_to"):
for item in parsed.get(key, []) or []:
if isinstance(item, dict) and item.get("address"):
addresses.add(item["address"].lower())
return addresses
async def lookup_known_messages(
database,
account_id: str,
internet_ids: list[str],
) -> dict[str, KnownMessage]:
folded_ids = [
value
for value in {folded_message_id(item) for item in internet_ids}
if value
]
if not folded_ids:
return {}
pipeline = [
{
"$match": {
"$and": [
{
"$or": [
{"message_id_folded": {"$in": folded_ids}},
{"message_id": {"$in": folded_ids}},
]
},
{
"$or": [
{"message_scope": "EXTERNAL"},
{"message_scope": {"$exists": False}},
]
},
]
}
},
{
"$lookup": {
"from": "mailbox_entries",
"let": {"message_id": "$_id"},
"pipeline": [
{
"$match": {
"$expr": {
"$and": [
{"$eq": ["$message_id", "$$message_id"]},
{"$eq": ["$account_id", account_id]},
]
}
}
},
{"$limit": 1},
],
"as": "entry",
}
},
{"$match": {"entry.0": {"$exists": True}}},
]
known: dict[str, KnownMessage] = {}
cursor = await database.messages.aggregate(pipeline)
async for document in cursor:
internet_id = canonical_message_id(
document.get("internet_message_id") or document.get("message_id")
)
internet_key = folded_message_id(internet_id)
if not internet_id or not internet_key:
continue
entry = document["entry"][0]
thread_id = entry.get("thread_id") or document.get("thread_id")
if thread_id is None:
continue
participants = participants_from_parsed(document)
known[internet_key] = KnownMessage(
id=document["_id"],
thread_id=thread_id,
internet_message_id=internet_id,
account_id=account_id,
sent_at=document.get("sent_at"),
received_at=document.get("received_at"),
created_at=document.get("created_at"),
participants=frozenset(participants),
normalized_subject=normalize_subject(document.get("subject")),
)
return known
async def resolve_thread(
database,
*,
account_id: str,
in_reply_to: str | None,
references: list[str],
subject: str,
participants: set[str],
) -> ThreadResolution:
search_ids = [
*([in_reply_to] if in_reply_to else []),
*references,
]
known = await lookup_known_messages(database, account_id, search_ids)
resolution = resolve_thread_from_known(
account_id=account_id,
in_reply_to=in_reply_to,
references=references,
subject=subject,
participants=participants,
known_by_internet_id=known,
)
logger.info(
"thread_resolved account_id=%s method=%s relation=%s parent_found=%s",
account_id,
resolution.thread_resolution_method,
resolution.thread_relation,
resolution.parent_message_id is not None,
)
return resolution