ai-tg/utils/context.py

import time
from collections import defaultdict
from config import (
    MAX_HISTORY_MESSAGES,
    MAX_TOKENS_DEFAULT,
    DEFAULT_SYSTEM_PROMPT,
    IMAGE_DAILY_LIMIT,
    MAX_INPUT_TOKENS,
)

# chat_id → list of {"role": ..., "content": ...}
_histories: dict[int, list] = defaultdict(list)

# user_id → settings dict
_user_settings: dict[int, dict] = {}

_system_prompt: str = DEFAULT_SYSTEM_PROMPT
_start_time: float = time.time()
_msg_count: int = 0
_event_counts: dict[str, int] = defaultdict(int)
_last_error: str = ""


# ── History ──────────────────────────────────────────────────

def get_history(chat_id: int) -> list:
    return list(_histories[chat_id])


def add_message(chat_id: int, role: str, content):
    _histories[chat_id].append({"role": role, "content": content})
    if len(_histories[chat_id]) > MAX_HISTORY_MESSAGES:
        _histories[chat_id] = _histories[chat_id][-MAX_HISTORY_MESSAGES:]


def clear_history(chat_id: int):
    _histories[chat_id].clear()


def clear_all_histories():
    _histories.clear()


def pop_last_message(chat_id: int):
    """Remove the most recent message from the history (for error rollback)."""
    if _histories[chat_id]:
        _histories[chat_id].pop()


def history_summary(chat_id: int, n: int = 6) -> str:
    msgs = _histories[chat_id][-n:]
    lines = []
    for m in msgs:
        content = m["content"]
        if isinstance(content, list):
            text = "[image/media]"
        else:
            text = str(content)[:200]
        lines.append(f"{m['role'].capitalize()}: {text}")
    return "\n".join(lines) if lines else "(empty)"


def estimate_tokens(text) -> int:
    if text is None:
        return 0
    if not isinstance(text, str):
        text = str(text)
    # Conservative enough for cost guarding without provider-specific tokenizers.
    return max(1, len(text) // 3)


def estimate_input_tokens(messages: list[dict], system_prompt: str = "") -> int:
    total = estimate_tokens(system_prompt)
    for msg in messages:
        total += estimate_tokens(msg.get("role", ""))
        total += estimate_tokens(msg.get("content", ""))
    return total


def reset_chat_if_input_too_large(chat_id: int, system_prompt: str = "", keep_last: bool = True) -> bool:
    messages = get_history(chat_id)
    if estimate_input_tokens(messages, system_prompt) <= MAX_INPUT_TOKENS:
        return False

    last = messages[-1:] if keep_last else []
    clear_history(chat_id)
    for msg in last:
        add_message(chat_id, msg.get("role", "user"), msg.get("content", ""))
    record_event("history_auto_resets")
    return True


# ── User settings ─────────────────────────────────────────────

def get_user_settings(user_id: int) -> dict:
    if user_id not in _user_settings:
        _user_settings[user_id] = {
            "max_tokens": MAX_TOKENS_DEFAULT,
            "search_enabled": True,
            "lang": "auto",
        }
    return _user_settings[user_id]


def set_user_setting(user_id: int, key: str, value):
    get_user_settings(user_id)[key] = value
    _autosave()


# ── System prompt ─────────────────────────────────────────────

def get_system_prompt() -> str:
    return _system_prompt


def set_system_prompt(prompt: str):
    global _system_prompt
    _system_prompt = prompt
    _autosave()


def reset_system_prompt():
    global _system_prompt
    _system_prompt = DEFAULT_SYSTEM_PROMPT
    _autosave()


def export_settings() -> dict:
    return {
        "system_prompt": _system_prompt,
        "user_settings": {
            str(uid): dict(settings)
            for uid, settings in _user_settings.items()
        },
    }


def import_settings(data: dict):
    global _system_prompt, _user_settings
    if isinstance(data.get("system_prompt"), str):
        _system_prompt = data["system_prompt"]

    raw_settings = data.get("user_settings", {})
    if isinstance(raw_settings, dict):
        loaded: dict[int, dict] = {}
        for uid, settings in raw_settings.items():
            try:
                user_id = int(uid)
            except (TypeError, ValueError):
                continue
            if isinstance(settings, dict):
                loaded[user_id] = dict(settings)
        _user_settings = loaded


# ── Stats ─────────────────────────────────────────────────────

def increment_msg():
    global _msg_count
    _msg_count += 1


def record_event(name: str, amount: int = 1):
    _event_counts[name] += amount


def record_error(message: str):
    global _last_error
    _event_counts["errors"] += 1
    _last_error = str(message)[:300]


def _autosave():
    try:
        from services import settings as settings_store
        settings_store.save()
    except Exception:
        pass


# ── Image rate limiting (cost guard) ─────────────────────────

_image_log: dict[int, list[float]] = defaultdict(list)  # user_id → timestamps


def can_generate_image(user_id: int) -> bool:
    cutoff = time.time() - 86400
    _image_log[user_id] = [t for t in _image_log[user_id] if t > cutoff]
    return len(_image_log[user_id]) < IMAGE_DAILY_LIMIT


def record_image(user_id: int):
    _image_log[user_id].append(time.time())
    record_event("images_generated")


def images_used_today(user_id: int) -> int:
    cutoff = time.time() - 86400
    return len([t for t in _image_log[user_id] if t > cutoff])


# ── Stats ─────────────────────────────────────────────────────

def get_stats() -> dict:
    elapsed = int(time.time() - _start_time)
    h, remainder = divmod(elapsed, 3600)
    m, s = divmod(remainder, 60)
    return {
        "uptime": f"{h}h {m}m {s}s",
        "uptime_seconds": elapsed,
        "messages_processed": _msg_count,
        "active_chats": len([k for k, v in _histories.items() if v]),
        "stored_history_messages": sum(len(v) for v in _histories.values()),
        "events": dict(_event_counts),
        "last_error": _last_error,
    }