From aade1479bf086f5cef1e9f61b1cad549a9e48867 Mon Sep 17 00:00:00 2001 From: arsam Date: Fri, 18 Sep 2026 22:41:13 +0330 Subject: [PATCH] Add conversation context management --- utils/context.py | 216 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 216 insertions(+) create mode 100644 utils/context.py diff --git a/utils/context.py b/utils/context.py new file mode 100644 index 0000000..692e32c --- /dev/null +++ b/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, + }