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, }