atlas-iac/services/hermes/scripts/migrate_telegram_api_sessions.py
2026-08-17 14:39:21 +00:00

375 lines
13 KiB
Python

#!/usr/bin/env python3
"""Bound Telegram context and backfill durable API-session metadata."""
from __future__ import annotations
import argparse
import json
import re
import sqlite3
from pathlib import Path
from typing import Any
MAX_RESPONSE_BYTES = 32 << 20
MAX_HISTORY_ITEMS = 80
MAX_HISTORY_BYTES = 192_000
MAX_RECENT_TURNS = 12
MAX_ITEM_CHARS = 8_000
MAX_SUMMARY_CHARS = 12_000
SUMMARY_PREFIX = "Telegram continuity summary (older bounded context):"
ATTACHMENT_NOTE = (
"[Earlier Telegram image omitted from bounded context; attach it again "
"in Hermes WebUI if the original pixels are required.]"
)
_CONVERSATION = re.compile(r"^telegram(?:-topic-[a-z0-9-]{1,160})?$")
_MEDIA_PATH = re.compile(
r'(?:MEDIA:\s*)?(?:/opt/data/(?:cache/images|workspace)|/workspace)/[^\s\]\[)}`"]+',
re.IGNORECASE,
)
_DATA_IMAGE = re.compile(r"data:image/[^;,\s]+;base64,[A-Za-z0-9+/=]+", re.IGNORECASE)
_MESSAGE_IDENTITY_COLUMNS = (
"session_id",
"role",
"content",
"tool_call_id",
"tool_calls",
"tool_name",
"timestamp",
"token_count",
"finish_reason",
"reasoning",
"reasoning_content",
"reasoning_details",
"codex_reasoning_items",
"codex_message_items",
"platform_message_id",
"observed",
"active",
"compacted",
)
def _text(value: Any) -> str:
"""Return bounded context text without reusable transport-only images."""
if isinstance(value, str):
text = value
elif isinstance(value, list):
parts: list[str] = []
for part in value:
if not isinstance(part, dict):
continue
kind = str(part.get("type") or "").lower()
if "image" in kind or "image_url" in part:
parts.append(ATTACHMENT_NOTE)
continue
candidate = part.get("text", part.get("content", ""))
if isinstance(candidate, str):
parts.append(candidate)
text = "\n".join(parts)
else:
try:
text = json.dumps(value, ensure_ascii=False, sort_keys=True)
except (TypeError, ValueError):
text = str(value)
text = _DATA_IMAGE.sub(ATTACHMENT_NOTE, text)
text = _MEDIA_PATH.sub("[private attachment]", text)
return text[:MAX_ITEM_CHARS]
def _normalized_item(raw: Any) -> dict[str, Any] | None:
if not isinstance(raw, dict):
return None
role = str(raw.get("role") or "").strip().lower()
if role not in {"assistant", "system", "tool", "user"}:
return None
content = _text(raw.get("content", ""))
if not content and role != "assistant":
return None
item: dict[str, Any] = {"role": role, "content": content}
if raw.get("_db_persisted") is True:
item["_db_persisted"] = True
for key in ("name", "tool_call_id", "tool_name"):
value = raw.get(key)
if isinstance(value, str) and value:
item[key] = value[:256]
calls = raw.get("tool_calls")
if isinstance(calls, list):
bounded_calls = []
for call in calls[:16]:
if not isinstance(call, dict):
continue
function = (
call.get("function") if isinstance(call.get("function"), dict) else {}
)
bounded = {
"id": str(call.get("id") or "")[:256],
"type": str(call.get("type") or "function")[:32],
"function": {
"name": str(function.get("name") or "")[:256],
"arguments": _text(function.get("arguments", ""))[:2_000],
},
}
bounded_calls.append(bounded)
if bounded_calls:
item["tool_calls"] = bounded_calls
return item
def _fingerprint(item: dict[str, Any]) -> str:
try:
visible = {key: value for key, value in item.items() if not key.startswith("_")}
return json.dumps(
visible, ensure_ascii=False, sort_keys=True, separators=(",", ":")
)
except (TypeError, ValueError):
return repr(item)
def _deduplicate(items: list[dict[str, Any]]) -> list[dict[str, Any]]:
"""Drop replay duplicates while preserving intentional repeated prose."""
result: list[dict[str, Any]] = []
stable_ids: set[tuple[str, str]] = set()
for item in items:
stable = ""
if isinstance(item.get("tool_call_id"), str):
stable = item["tool_call_id"]
elif isinstance(item.get("id"), str):
stable = item["id"]
if stable:
key = (item["role"], stable)
if key in stable_ids:
continue
stable_ids.add(key)
if result and _fingerprint(result[-1]) == _fingerprint(item):
continue
result.append(item)
# Response replays can append one exact transcript to itself. Remove only
# complete adjacent blocks, not ordinary repeated user phrases.
changed = True
while changed and len(result) > 1:
changed = False
for width in range(len(result) // 2, 0, -1):
if [_fingerprint(x) for x in result[-2 * width : -width]] == [
_fingerprint(x) for x in result[-width:]
]:
del result[-width:]
changed = True
break
return result
def _summary(older: list[dict[str, Any]]) -> dict[str, str] | None:
lines: list[str] = []
for item in older:
role = item["role"]
content = str(item.get("content") or "").strip()
if role == "system" and content.startswith(SUMMARY_PREFIX):
candidates = content[len(SUMMARY_PREFIX) :].strip().splitlines()
elif role in {"user", "assistant"} and content:
candidates = [f"- {role.title()}: {content[:600]}"]
else:
continue
for line in candidates:
clean = " ".join(line.split())
if clean and clean not in lines:
lines.append(clean)
content = SUMMARY_PREFIX
for line in lines[-40:]:
candidate = content + "\n" + line
if len(candidate) > MAX_SUMMARY_CHARS:
break
content = candidate
return (
{
"role": "system",
"content": content,
"_db_persisted": True,
"_compressed_summary": True,
}
if lines
else None
)
def compact_telegram_history(history: Any) -> list[dict[str, Any]]:
"""Return a deterministic summary plus bounded recent Telegram turns."""
if not isinstance(history, list):
return []
normalized = _deduplicate(
[item for raw in history if (item := _normalized_item(raw)) is not None]
)
recent: list[dict[str, Any]] = []
recent_bytes = 2
user_turns = 0
split = len(normalized)
for index in range(len(normalized) - 1, -1, -1):
item = normalized[index]
encoded = len(_fingerprint(item).encode("utf-8")) + 1
next_turns = user_turns + (item["role"] == "user")
if recent and (
len(recent) >= MAX_HISTORY_ITEMS
or recent_bytes + encoded > MAX_HISTORY_BYTES
or next_turns > MAX_RECENT_TURNS
):
break
recent.insert(0, item)
recent_bytes += encoded
user_turns = next_turns
split = index
summary = _summary(normalized[:split])
result = ([summary] if summary else []) + recent
while (
len(json.dumps(result, ensure_ascii=False).encode("utf-8"))
> (MAX_HISTORY_BYTES + MAX_SUMMARY_CHARS)
and len(recent) > 1
):
recent.pop(0)
result = ([summary] if summary else []) + recent
return result
def telegram_sessions(response_store: Path) -> dict[str, str]:
"""Return safe session ID to stable Telegram conversation mappings."""
if not response_store.exists():
return {}
connection = sqlite3.connect(f"file:{response_store}?mode=ro", uri=True)
try:
rows = connection.execute(
"""SELECT c.name, r.data
FROM conversations AS c
JOIN responses AS r ON r.response_id = c.response_id
WHERE (c.name = 'telegram' OR c.name LIKE 'telegram-topic-%')
AND LENGTH(r.data) <= ?
ORDER BY c.name, r.response_id""",
(MAX_RESPONSE_BYTES,),
).fetchall()
except sqlite3.DatabaseError:
return {}
finally:
connection.close()
result: dict[str, str] = {}
for conversation, raw in rows:
name = str(conversation or "")
if not _CONVERSATION.fullmatch(name):
continue
try:
payload = json.loads(raw)
session_id = str(payload.get("session_id") or "").strip()
except (AttributeError, TypeError, ValueError, json.JSONDecodeError):
continue
if (
session_id
and len(session_id) <= 256
and not re.search(r"[\r\n\x00/\\]", session_id)
):
result.setdefault(session_id, name)
return result
def deduplicate_stored_messages(
connection: sqlite3.Connection, session_ids: list[str]
) -> int:
"""Remove only byte-for-byte replay rows for known Telegram API sessions."""
columns = {str(row[1]) for row in connection.execute("PRAGMA table_info(messages)")}
required = {"id", "session_id", "role", "content", "timestamp"}
if not required.issubset(columns):
return 0
identity = [column for column in _MESSAGE_IDENTITY_COLUMNS if column in columns]
partition = ", ".join(identity)
changed = 0
for session_id in session_ids:
query = f"""DELETE FROM messages WHERE id IN (
SELECT id FROM (
SELECT id, ROW_NUMBER() OVER (
PARTITION BY {partition} ORDER BY id
) AS replay_number
FROM messages
WHERE session_id = ? AND EXISTS (
SELECT 1 FROM sessions
WHERE sessions.id = messages.session_id
AND sessions.source = 'api_server'
)
) WHERE replay_number > 1
)""" # noqa: S608 -- identifiers come only from the fixed allowlist
cursor = connection.execute(query, (session_id,))
changed += max(cursor.rowcount, 0)
return changed
def migrate(state_database: Path, response_store: Path) -> int:
"""Annotate known sessions without creating or rewriting conversations."""
mappings = telegram_sessions(response_store)
if not mappings or not state_database.exists():
return 0
connection = sqlite3.connect(state_database)
changed = 0
try:
columns = {
str(row[1]) for row in connection.execute("PRAGMA table_info(sessions)")
}
required = {
"id",
"source",
"session_key",
"chat_type",
"display_name",
"origin_json",
"title",
}
if not required.issubset(columns):
return 0
for session_id, conversation in mappings.items():
origin = json.dumps(
{"platform": "telegram", "session_key": conversation},
separators=(",", ":"),
)
default_title = (
"Telegram · General" if conversation == "telegram" else "Telegram"
)
cursor = connection.execute(
"""UPDATE sessions
SET session_key = ?, chat_type = 'private',
display_name = 'Telegram', origin_json = ?,
title = CASE
WHEN title IS NULL OR TRIM(title) = '' OR
LOWER(TRIM(title)) = 'unassigned'
THEN ? ELSE title END
WHERE id = ? AND source = 'api_server'
AND (COALESCE(session_key, '') != ? OR
COALESCE(chat_type, '') != 'private' OR
COALESCE(display_name, '') != 'Telegram' OR
COALESCE(origin_json, '') != ? OR title IS NULL OR
TRIM(title) = '' OR LOWER(TRIM(title)) = 'unassigned')""",
(
conversation,
origin,
default_title,
session_id,
conversation,
origin,
),
)
changed += max(cursor.rowcount, 0)
deduplicate_stored_messages(connection, list(mappings))
connection.commit()
except sqlite3.DatabaseError:
connection.rollback()
return 0
finally:
connection.close()
return changed
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("state_database", type=Path)
parser.add_argument("response_store", type=Path)
args = parser.parse_args()
migrate(args.state_database, args.response_store)
return 0
if __name__ == "__main__":
raise SystemExit(main())