Update app.py

app.py

@@ -3,25 +3,35 @@ from pathlib import Path
 
 import aiosqlite
 import gradio as gr
+import tiktoken
 from openai import AsyncOpenAI
 
 # =========================
 # Конфигурация
 # =========================
 
-# Каталог и файл базы данных
 DB_DIR = Path("./data")
 DB_PATH = DB_DIR / "chats.db"
 
-# Названия 5 чатов
 CHAT_IDS = ["chat1", "chat2", "chat3", "chat4", "chat5"]
 
+# Ограничение входного контекста
+INPUT_TOKEN_LIMIT = 4096
+
+# Ограничение ответа модели
+OUTPUT_TOKEN_LIMIT = 1024
+
+MODEL_NAME = "gpt-4o-mini"
+
 
 def ensure_db_dir():
     DB_DIR.mkdir(parents=True, exist_ok=True)
 
 
-# Предустановленные учетные данные
+# =========================
+# Переменные окружения
+# =========================
+
 authorized_users_str = os.getenv("AUTHORIZED_USERS")
 if not authorized_users_str:
     raise ValueError("Переменная окружения AUTHORIZED_USERS не установлена")
@@ -33,14 +43,84 @@ if not proctor_links:
 PROCTOR_LINKS = dict(user.split(":", 1) for user in proctor_links.split(","))
 
 SYSTEM_PROMPT = os.getenv("SYSTEM_PROMPT", "Ты полезный ассистент.")
+
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 if not OPENAI_API_KEY:
     raise ValueError("Переменная окружения OPENAI_API_KEY не установлена")
 
-# Инициализация AsyncOpenAI клиента один раз
+
+# =========================
+# OpenAI client
+# =========================
+
 client = AsyncOpenAI(api_key=OPENAI_API_KEY)
 
 
+# =========================
+# Подсчет токенов и обрезка истории
+# =========================
+
+def get_encoder(model_name: str = MODEL_NAME):
+    try:
+        return tiktoken.encoding_for_model(model_name)
+    except Exception:
+        return tiktoken.get_encoding("cl100k_base")
+
+
+def count_message_tokens(messages, model_name: str = MODEL_NAME) -> int:
+    enc = get_encoder(model_name)
+    total = 0
+
+    # Приблизительный расчет для chat messages
+    for msg in messages:
+        total += 4
+        total += len(enc.encode(msg.get("role", "")))
+        total += len(enc.encode(msg.get("content", "")))
+    total += 2
+
+    return total
+
+
+def trim_messages_to_limit(
+    system_prompt: str,
+    history: list,
+    new_user_message: str,
+    model_name: str = MODEL_NAME,
+    input_limit: int = INPUT_TOKEN_LIMIT,
+):
+    base_messages = [{"role": "system", "content": system_prompt}]
+    candidate_history = history[:] if history else []
+
+    final_messages = base_messages + candidate_history + [
+        {"role": "user", "content": new_user_message}
+    ]
+
+    while len(final_messages) > 2 and count_message_tokens(final_messages, model_name) > input_limit:
+        if candidate_history:
+            candidate_history.pop(0)
+        else:
+            break
+
+        final_messages = base_messages + candidate_history + [
+            {"role": "user", "content": new_user_message}
+        ]
+
+    # Если даже без истории сообщение слишком длинное — режем само сообщение
+    if count_message_tokens(final_messages, model_name) > input_limit:
+        enc = get_encoder(model_name)
+
+        system_only_tokens = count_message_tokens(base_messages, model_name)
+        allowed_for_user = max(1, input_limit - system_only_tokens - 20)
+
+        trimmed_user_message = enc.decode(enc.encode(new_user_message)[:allowed_for_user])
+
+        final_messages = base_messages + [
+            {"role": "user", "content": trimmed_user_message}
+        ]
+
+    return final_messages
+
+
 # =========================
 # Работа с БД
 # =========================
@@ -92,13 +172,16 @@ async def load_history(chat_id, username):
 async def clear_history(chat_id, username):
     ensure_db_dir()
     async with aiosqlite.connect(DB_PATH) as db:
-        await db.execute(f"DELETE FROM {chat_id} WHERE username = ?", (username,))
+        await db.execute(
+            f"DELETE FROM {chat_id} WHERE username = ?",
+            (username,),
+        )
         await db.commit()
     return []
 
 
 # =========================
-# Основная логика чата
+# Основная логика
 # =========================
 
 async def send_message(message, history, username, chat_id):
@@ -116,10 +199,8 @@ async def send_message(message, history, username, chat_id):
     history = history or []
     clean_message = message.strip()
 
-    # Сохраняем сообщение пользователя
     await save_message(chat_id, username, "user", clean_message)
 
-    # Спец-команда
     if clean_message.lower() == "proctorlink":
         if username in PROCTOR_LINKS:
             assistant_response = PROCTOR_LINKS[username]
@@ -134,17 +215,20 @@ async def send_message(message, history, username, chat_id):
         ]
         return "", new_history
 
-    # Формируем сообщения для модели с учетом истории
-    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
-    messages.extend(history)
-    messages.append({"role": "user", "content": clean_message})
+    messages = trim_messages_to_limit(
+        system_prompt=SYSTEM_PROMPT,
+        history=history,
+        new_user_message=clean_message,
+        model_name=MODEL_NAME,
+        input_limit=INPUT_TOKEN_LIMIT,
+    )
 
     try:
         response = await client.chat.completions.create(
-            model="gpt-4o",
+            model=MODEL_NAME,
             messages=messages,
             temperature=0.7,
-            max_tokens=4000,
+            max_tokens=OUTPUT_TOKEN_LIMIT,
             timeout=120,
         )
         assistant_response = response.choices[0].message.content or ""
@@ -152,7 +236,6 @@ async def send_message(message, history, username, chat_id):
         print(f"OpenAI API error: {e}")
         assistant_response = "Error: Could not connect to the AI service."
 
-    # Сохраняем ответ ассистента
     await save_message(chat_id, username, "assistant", assistant_response)
 
     new_history = history + [
@@ -262,34 +345,34 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
         load_history_btn = gr.Button("Загрузить историю", elem_id="load_history_btn")
         clear_history_btn = gr.Button("Очистить историю", elem_id="clear_history_btn")
 
-    # Инициализация базы данных при запуске приложения
     demo.load(init_db, inputs=None, outputs=None)
-
-    # Установка имени пользователя
     demo.load(set_username, inputs=None, outputs=username_state)
 
-    # Обновление выбранного чата
     chat_selector.change(update_chat_id, inputs=chat_selector, outputs=chat_id_state)
 
-    # Автозагрузка истории
-    chat_id_state.change(load_chat_history, inputs=[username_state, chat_id_state], outputs=chatbot)
-    username_state.change(load_chat_history, inputs=[username_state, chat_id_state], outputs=chatbot)
+    chat_id_state.change(
+        load_chat_history,
+        inputs=[username_state, chat_id_state],
+        outputs=chatbot,
+    )
+    username_state.change(
+        load_chat_history,
+        inputs=[username_state, chat_id_state],
+        outputs=chatbot,
+    )
 
-    # Ручная загрузка истории
     load_history_btn.click(
         load_chat_history,
         inputs=[username_state, chat_id_state],
         outputs=chatbot,
     )
 
-    # Очистка истории
     clear_history_btn.click(
         handle_clear_history,
         inputs=[username_state, chat_id_state],
         outputs=chatbot,
     )
 
-    # Отправка сообщения
     submit_btn.click(
         send_message,
         inputs=[msg, chatbot, username_state, chat_id_state],