from __future__ import annotations from dataclasses import dataclass, field from enum import StrEnum from typing import Any class ContextSource(StrEnum): CONTACT = "contact" LOCATION = "location" REACTION = "reaction" POLL = "poll" BUTTON = "button" FORWARD = "forward" MENTION = "mention" MEDIA = "media" UNKNOWN = "unknown" class ContextType(StrEnum): STRUCTURED_DATA = "structured_data" USER_GENERATED = "user_generated" SYSTEM = "system" @dataclass class StructuredContextEntry: label: str source: ContextSource = ContextSource.UNKNOWN type: ContextType = ContextType.STRUCTURED_DATA data: dict[str, Any] = field(default_factory=dict) @dataclass class UntrustedStructuredContext: entries: list[StructuredContextEntry] = field(default_factory=list) @classmethod def from_contact( cls, display_name: str = "", vcard: str = "", contacts: list[dict[str, str]] | None = None, ) -> UntrustedStructuredContext: entry = StructuredContextEntry( label="contact", source=ContextSource.CONTACT, data={ "display_name": display_name, "vcard": vcard, "contacts": contacts or [], }, ) return cls(entries=[entry]) @classmethod def from_location( cls, latitude: float, longitude: float, name: str = "", address: str = "" ) -> UntrustedStructuredContext: entry = StructuredContextEntry( label="location", source=ContextSource.LOCATION, data={ "latitude": latitude, "longitude": longitude, "name": name, "address": address, }, ) return cls(entries=[entry]) @classmethod def from_reaction(cls, emoji: str, target_msg_id: str, target_jid: str) -> UntrustedStructuredContext: entry = StructuredContextEntry( label="reaction", source=ContextSource.REACTION, data={ "emoji": emoji, "target_msg_id": target_msg_id, "target_jid": target_jid, }, ) return cls(entries=[entry]) @classmethod def from_poll( cls, poll_name: str, options: list[str], selected_options: list[str] | None = None ) -> UntrustedStructuredContext: entry = StructuredContextEntry( label="poll", source=ContextSource.POLL, data={ "name": poll_name, "options": options, "selected_options": selected_options or [], }, ) return cls(entries=[entry]) @classmethod def from_button(cls, button_id: str, display_text: str) -> UntrustedStructuredContext: entry = StructuredContextEntry( label="button", source=ContextSource.BUTTON, data={ "button_id": button_id, "display_text": display_text, }, ) return cls(entries=[entry]) @classmethod def from_forward(cls, forwarding_score: int = 0, forwarded_from: str = "") -> UntrustedStructuredContext: entry = StructuredContextEntry( label="forward", source=ContextSource.FORWARD, data={ "is_forwarded": True, "forwarding_score": forwarding_score, "forwarded_from": forwarded_from, }, ) return cls(entries=[entry]) @classmethod def from_mentions(cls, mentioned_jids: list[str]) -> UntrustedStructuredContext: entry = StructuredContextEntry( label="mention", source=ContextSource.MENTION, data={"mentioned_jids": mentioned_jids}, ) return cls(entries=[entry]) def add_entry(self, entry: StructuredContextEntry) -> None: self.entries.append(entry) def get_by_label(self, label: str) -> StructuredContextEntry | None: for entry in self.entries: if entry.label == label: return entry return None def get_by_source(self, source: ContextSource) -> list[StructuredContextEntry]: return [e for e in self.entries if e.source == source] def to_metadata(self) -> dict[str, Any]: result: dict[str, Any] = {} for entry in self.entries: result[f"structured_{entry.label}"] = { "source": entry.source.value, "type": entry.type.value, "data": entry.data, } return result def is_empty(self) -> bool: return len(self.entries) == 0