""" Export app messages to JSONL.GZ format. Outputs: conversation_id, message_id, query, answer, inputs (raw JSON), retriever_resources (from message_metadata), feedback (user feedbacks array). Uses (created_at, id) cursor pagination and batch-loads feedbacks to avoid N+1. Does NOT touch Message.inputs / Message.user_feedback properties. """ import datetime import gzip import json import logging import tempfile from collections import defaultdict from collections.abc import Generator, Iterable from pathlib import Path, PurePosixPath from typing import Any, BinaryIO, cast import orjson import sqlalchemy as sa from pydantic import BaseModel, ConfigDict, Field from sqlalchemy import select, tuple_ from sqlalchemy.orm import Session from extensions.ext_database import db from extensions.ext_storage import storage from models.model import Message, MessageFeedback logger = logging.getLogger(__name__) MAX_FILENAME_BASE_LENGTH = 1024 FORBIDDEN_FILENAME_SUFFIXES = (".jsonl.gz", ".jsonl", ".gz") class AppMessageExportFeedback(BaseModel): id: str app_id: str conversation_id: str message_id: str rating: str content: str | None = None from_source: str from_end_user_id: str | None = None from_account_id: str | None = None created_at: str updated_at: str model_config = ConfigDict(extra="forbid") class AppMessageExportRecord(BaseModel): conversation_id: str message_id: str query: str answer: str inputs: dict[str, Any] retriever_resources: list[Any] = Field(default_factory=list) feedback: list[AppMessageExportFeedback] = Field(default_factory=list) model_config = ConfigDict(extra="forbid") class AppMessageExportStats(BaseModel): batches: int = 0 total_messages: int = 0 messages_with_feedback: int = 0 total_feedbacks: int = 0 model_config = ConfigDict(extra="forbid") class AppMessageExportService: @staticmethod def validate_export_filename(filename: str) -> str: normalized = filename.strip() if not normalized: raise ValueError("--filename must not be empty.") normalized_lower = normalized.lower() if normalized_lower.endswith(FORBIDDEN_FILENAME_SUFFIXES): raise ValueError("--filename must not include .jsonl.gz/.jsonl/.gz suffix; pass base filename only.") if normalized.startswith("/"): raise ValueError("--filename must be a relative path; absolute paths are not allowed.") if "\\" in normalized: raise ValueError("--filename must use '/' as path separator; '\\' is not allowed.") if "//" in normalized: raise ValueError("--filename must not contain empty path segments ('//').") if len(normalized) > MAX_FILENAME_BASE_LENGTH: raise ValueError(f"--filename is too long; max length is {MAX_FILENAME_BASE_LENGTH}.") for ch in normalized: if ch == "\x00" or ord(ch) < 32 or ord(ch) == 127: raise ValueError("--filename must not contain control characters or NUL.") parts = PurePosixPath(normalized).parts if not parts: raise ValueError("--filename must include a file name.") if any(part in (".", "..") for part in parts): raise ValueError("--filename must not contain '.' or '..' path segments.") return normalized @property def output_gz_name(self) -> str: return f"{self._filename_base}.jsonl.gz" @property def output_jsonl_name(self) -> str: return f"{self._filename_base}.jsonl" def __init__( self, app_id: str, end_before: datetime.datetime, filename: str, *, start_from: datetime.datetime | None = None, batch_size: int = 1000, use_cloud_storage: bool = False, dry_run: bool = False, ) -> None: if start_from and start_from >= end_before: raise ValueError(f"start_from ({start_from}) must be before end_before ({end_before})") self._app_id = app_id self._end_before = end_before self._start_from = start_from self._filename_base = self.validate_export_filename(filename) self._batch_size = batch_size self._use_cloud_storage = use_cloud_storage self._dry_run = dry_run def run(self) -> AppMessageExportStats: stats = AppMessageExportStats() logger.info( "export_app_messages: app_id=%s, start_from=%s, end_before=%s, dry_run=%s, cloud=%s, output_gz=%s", self._app_id, self._start_from, self._end_before, self._dry_run, self._use_cloud_storage, self.output_gz_name, ) if self._dry_run: for _ in self._iter_records_with_stats(stats): pass self._finalize_stats(stats) return stats if self._use_cloud_storage: self._export_to_cloud(stats) else: self._export_to_local(stats) self._finalize_stats(stats) return stats def iter_records(self) -> Generator[AppMessageExportRecord, None, None]: for batch in self._iter_record_batches(): yield from batch @staticmethod def write_jsonl_gz(records: Iterable[AppMessageExportRecord], fileobj: BinaryIO) -> None: with gzip.GzipFile(fileobj=fileobj, mode="wb") as gz: for record in records: gz.write(orjson.dumps(record.model_dump(mode="json")) + b"\n") def _export_to_local(self, stats: AppMessageExportStats) -> None: output_path = Path.cwd() / self.output_gz_name output_path.parent.mkdir(parents=True, exist_ok=True) with output_path.open("wb") as output_file: self.write_jsonl_gz(self._iter_records_with_stats(stats), output_file) def _export_to_cloud(self, stats: AppMessageExportStats) -> None: with tempfile.SpooledTemporaryFile(max_size=64 * 1024 * 1024) as tmp: self.write_jsonl_gz(self._iter_records_with_stats(stats), cast(BinaryIO, tmp)) tmp.seek(0) data = tmp.read() storage.save(self.output_gz_name, data) logger.info("export_app_messages: uploaded %d bytes to cloud key=%s", len(data), self.output_gz_name) def _iter_records_with_stats(self, stats: AppMessageExportStats) -> Generator[AppMessageExportRecord, None, None]: for record in self.iter_records(): self._update_stats(stats, record) yield record @staticmethod def _update_stats(stats: AppMessageExportStats, record: AppMessageExportRecord) -> None: stats.total_messages += 1 if record.feedback: stats.messages_with_feedback += 1 stats.total_feedbacks += len(record.feedback) def _finalize_stats(self, stats: AppMessageExportStats) -> None: if stats.total_messages == 0: stats.batches = 0 return stats.batches = (stats.total_messages + self._batch_size - 1) // self._batch_size def _iter_record_batches(self) -> Generator[list[AppMessageExportRecord], None, None]: cursor: tuple[datetime.datetime, str] | None = None while True: rows, cursor = self._fetch_batch(cursor) if not rows: break message_ids = [str(row.id) for row in rows] feedbacks_map = self._fetch_feedbacks(message_ids) yield [self._build_record(row, feedbacks_map) for row in rows] def _fetch_batch( self, cursor: tuple[datetime.datetime, str] | None ) -> tuple[list[Any], tuple[datetime.datetime, str] | None]: with Session(db.engine, expire_on_commit=False) as session: stmt = ( select( Message.id, Message.conversation_id, Message.query, Message.answer, Message._inputs, Message.message_metadata, Message.created_at, ) .where( Message.app_id == self._app_id, Message.created_at < self._end_before, ) .order_by(Message.created_at, Message.id) .limit(self._batch_size) ) if self._start_from: stmt = stmt.where(Message.created_at >= self._start_from) if cursor: stmt = stmt.where( tuple_(Message.created_at, Message.id) > tuple_( sa.literal(cursor[0], type_=sa.DateTime()), sa.literal(cursor[1], type_=Message.id.type), ) ) rows = list(session.execute(stmt).all()) if not rows: return [], cursor last = rows[-1] return rows, (last.created_at, last.id) def _fetch_feedbacks(self, message_ids: list[str]) -> dict[str, list[AppMessageExportFeedback]]: if not message_ids: return {} with Session(db.engine, expire_on_commit=False) as session: stmt = ( select(MessageFeedback) .where( MessageFeedback.message_id.in_(message_ids), MessageFeedback.from_source == "user", ) .order_by(MessageFeedback.message_id, MessageFeedback.created_at) ) feedbacks = list(session.scalars(stmt).all()) result: dict[str, list[AppMessageExportFeedback]] = defaultdict(list) for feedback in feedbacks: result[str(feedback.message_id)].append(AppMessageExportFeedback.model_validate(feedback.to_dict())) return result @staticmethod def _build_record(row: Any, feedbacks_map: dict[str, list[AppMessageExportFeedback]]) -> AppMessageExportRecord: retriever_resources: list[Any] = [] if row.message_metadata: try: metadata = json.loads(row.message_metadata) value = metadata.get("retriever_resources", []) if isinstance(value, list): retriever_resources = value except (json.JSONDecodeError, TypeError): pass message_id = str(row.id) return AppMessageExportRecord( conversation_id=str(row.conversation_id), message_id=message_id, query=row.query, answer=row.answer, inputs=row._inputs if isinstance(row._inputs, dict) else {}, retriever_resources=retriever_resources, feedback=feedbacks_map.get(message_id, []), )