Add conversation context management

utils/context.py

@@ -0,0 +1,216 @@
+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,
+    }