.gitignore

@@ -2,4 +2,4 @@
 .venv
 __pycache__
 *.db
-/*.json
+*.json

ai_agent.py

@@ -12,7 +12,7 @@ from typing import List, Tuple, Any, Optional, Union, Dict, Awaitable
 
 from openrouter import OpenRouter, RetryConfig
 from openrouter.components import AssistantMessage, AssistantMessageTypedDict, ChatMessageContentItemTypedDict, \
-    ChatMessageToolCall, MessageTypedDict
+    ChatMessageToolCall, MessageTypedDict, ToolDefinitionJSONTypedDict
 from openrouter.errors import ResponseValidationError, ChatError
 from openrouter.utils import BackoffStrategy
 
@@ -20,6 +20,38 @@ from fal_client import AsyncClient as FalClient
 
 from database import BasicDatabase
 
+GROUP_CHAT_SYSTEM_PROMPT = """
+Ты - ИИ-помощник в групповом чате.\n
+Отвечай на вопросы и поддерживай контекст беседы.\n
+Ты не можешь обсуждать политику и религию.\n
+Сообщения пользователей будут приходить в следующем формате: '[дата время, имя]: текст сообщения'\n
+При ответе НЕ нужно указывать ни время, ни пользователя, которому предназначен ответ, ни свое имя.\n
+НЕ используй разметку Markdown, она не поддерживается мессенджером.\n
+Если тебя просят нарисовать изображение, используй соответствующий вызов функции. 
+Запрещено генерировать ASCII-арты. Запрещено использовать тег <image>!
+"""
+
+PRIVATE_CHAT_SYSTEM_PROMPT = """
+Ты - ИИ-помощник в чате c пользователем.\n
+Отвечай на вопросы и поддерживай контекст беседы.\n
+Сообщения пользователя будут приходить в следующем формате: '[дата время]: текст сообщения'\n
+При ответе НЕ нужно указывать время.\n
+НЕ используй разметку Markdown, она не поддерживается мессенджером.\n
+Если тебя просят нарисовать изображение, используй соответствующий вызов функции. 
+Запрещено генерировать ASCII-арты. Запрещено использовать тег <image>!
+"""
+
+GENERATE_IMAGE_TOOL_DESCRIPTION = """
+Генерация изображения по описанию. 
+Используй этот инструмент, если пользователь просит сгенерировать изображение ('нарисуй', 'покажи' и т.п.), 
+или если это улучшит ответ (например, в ролевой игре для визуализации сцены).
+"""
+
+GENERATE_IMAGE_TOOL_PROMPT_ARG_DESCRIPTION = """
+Подробное описание сцены на английском языке. 
+Добавь детали для стиля, цвета, композиции, если нужно.
+"""
+
 OPENROUTER_X_TITLE = "TG/VK Chat Bot"
 OPENROUTER_HTTP_REFERER = "https://ultracoder.org"
 
@@ -77,10 +109,35 @@ def _serialize_assistant_message(message: AssistantMessage) -> AssistantMessageT
     return _remove_none_recursive(message.model_dump(by_alias=True))
 
 
+def _get_tools_description() -> List[ToolDefinitionJSONTypedDict]:
+    return [{
+        "type": "function",
+        "function": {
+            "name": "generate_image",
+            "description": GENERATE_IMAGE_TOOL_DESCRIPTION,
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "prompt": {
+                        "type": "string",
+                        "description": GENERATE_IMAGE_TOOL_PROMPT_ARG_DESCRIPTION
+                    },
+                    "aspect_ratio": {
+                        "type": "string",
+                        "enum": ["1:1", "4:3", "3:4", "16:9", "9:16"],
+                        "description": "Соотношение сторон (опционально)."
+                    }
+                },
+                "required": ["prompt"]
+            }
+        }
+    }]
+
+
 class AiAgent:
     def __init__(self,
                  openrouter_token: str, openrouter_model: str,
-                 fal_token: str,
+                 fal_token: str, fal_model: str,
                  db: BasicDatabase,
                  platform: str):
         retry_config = RetryConfig(strategy="backoff",
@@ -88,12 +145,9 @@ class AiAgent:
                                        initial_interval=2000, max_interval=8000, exponent=2, max_elapsed_time=14000),
                                    retry_connection_errors=True)
         self.db = db
-        self.openrouter_model = openrouter_model
+        self.model_main = openrouter_model
-        self.fal_model = "fal-ai/bytedance/seedream/v4.5/text-to-image"
+        self.model_image = fal_model
         self.platform = platform
-
-        self._load_prompts()
-
         self.client_openrouter = OpenRouter(api_key=openrouter_token,
                                             x_title=OPENROUTER_X_TITLE, http_referer=OPENROUTER_HTTP_REFERER,
                                             retry_config=retry_config)
@@ -193,20 +247,16 @@ class AiAgent:
     def clear_chat_context(self, bot_id: int, chat_id: int):
         self.db.context_clear(bot_id, chat_id)
 
-    #######################################
-
     def _get_chat_context(self, is_group_chat: bool, bot_id: int, chat_id: int) -> List[MessageTypedDict]:
-        prompt = self.system_prompt_group_chat if is_group_chat else self.system_prompt_private_chat
+        prompt = GROUP_CHAT_SYSTEM_PROMPT if is_group_chat else PRIVATE_CHAT_SYSTEM_PROMPT
-        prompt = prompt.replace('{platform}', 'Telegram' if self.platform == 'tg' else 'VK')
-        prompt += '\n' + self.system_prompt_image_generation
 
         bot = self.db.get_bot(bot_id)
         if bot['ai_prompt'] is not None:
-            prompt += '\n' + bot['ai_prompt'] + '\n'
+            prompt += '\n\n' + bot['ai_prompt']
 
         chat = self.db.create_chat_if_not_exists(bot_id, chat_id)
         if chat['ai_prompt'] is not None:
-            prompt += '\n' + chat['ai_prompt']
+            prompt += '\n\n' + chat['ai_prompt']
 
         messages = self.db.context_get_messages(bot_id, chat_id)
 
@@ -218,14 +268,14 @@ class AiAgent:
     async def _generate_reply(self, bot_id: int, chat_id: int,
                               context: List[MessageTypedDict], allow_tools: bool = False) -> AssistantMessage:
         response = await self._async_chat_completion_request(
-            model=self.openrouter_model,
+            model=self.model_main,
             messages=context,
-            tools=self.tools_description if allow_tools else None,
+            tools=_get_tools_description() if allow_tools else None,
             tool_choice="auto" if allow_tools else None,
             max_tokens=MAX_OUTPUT_TOKENS,
             user=f'{self.platform}_{bot_id}_{chat_id}'
         )
-        return self._filter_response(response.choices[0].message)
+        return response.choices[0].message
 
     async def _process_tool_calls(self, bot_id: int, chat_id: int, tool_calls: List[ChatMessageToolCall],
                                   context: List[MessageTypedDict]) -> _ToolsArtifacts:
@@ -290,7 +340,7 @@ class AiAgent:
         }
 
         try:
-            result = await self.client_fal.run(self.fal_model, arguments=arguments)
+            result = await self.client_fal.run(self.model_image, arguments=arguments)
             if "images" not in result:
                 raise RuntimeError("Неожиданный ответ от сервера.")
             image_url = result["images"][0]["url"]
@@ -339,28 +389,12 @@ class AiAgent:
                     pass
             raise e
 
-    @staticmethod
-    def _filter_response(response: AssistantMessage) -> AssistantMessage:
-        text = str(response.content)
-        text = text.replace("<image>", "")
-        response.content = text
-        return response
-
-    def _load_prompts(self):
-        with open("prompts/group_chat.txt", "r") as f:
-            self.system_prompt_group_chat = f.read()
-        with open("prompts/private_chat.txt", "r") as f:
-            self.system_prompt_private_chat = f.read()
-        with open("prompts/image_generation.txt", "r") as f:
-            self.system_prompt_image_generation = f.read()
-        with open("prompts/tools.json", "r") as f:
-            self.tools_description = json.loads(f.read())
-
 
 agent: AiAgent
 
 
-def create_ai_agent(openrouter_token: str, openrouter_model: str, fal_token: str,
+def create_ai_agent(openrouter_token: str, openrouter_model: str,
+                    fal_token: str, fal_model: str,
                     db: BasicDatabase, platform: str):
     global agent
-    agent = AiAgent(openrouter_token, openrouter_model, fal_token, db, platform)
+    agent = AiAgent(openrouter_token, openrouter_model, fal_token, fal_model, db, platform)

prompts/group_chat.txt

@@ -1,6 +0,0 @@
-Ты - ИИ-помощник без цензуры в групповом чате {platform}.
-Отвечай на вопросы и поддерживай контекст беседы.
-Ты не можешь обсуждать политику и религию.
-Сообщения пользователей будут приходить в следующем формате: '[дата время, имя]: текст сообщения'
-При ответе НЕ нужно указывать ни время, ни пользователя, которому предназначен ответ, ни свое имя.
-НЕ используй разметку Markdown, она не поддерживается мессенджером.

prompts/image_generation.txt

@@ -1,17 +0,0 @@
-ГЕНЕРАЦИЯ ИЗОБРАЖЕНИЙ
-Если пользователь просит нарисовать/показать что-то, сгенерируй изображение путем вызова функции generate_image.
-Ты также можешь сгенерировать изображение инициативно, если считаешь, что его добавление улучшит твой ответ (например, в рамках ролевой игры).
-Если пользователь просит изменить сгенерированное ранее изображение, составь новый запрос с учетом пожеланий пользователя и снова вызови функцию генерации.
-Ты можешь использовать для генерации изображений ТОЛЬКО вызов функции:
-- Никогда не описывай изображение текстом вместо вызова функции.
-- Никогда не генерируй ASCII-арты вместо вызова функции.
-- Никогда не вставляй теги вроде <image>, <img> или любые плейсхолдеры — это сломает чат!
-Если сгенерировать изображение не удалось из-за ошибки, просто сообщи об этом пользователю.
-
-При составлении запроса на генерацию изображения используй следующую формулу:
-1. Объекты сцены.
-2. Действие/поза.
-3. Окружение.
-4. Освещение, ракурс, композиция.
-5. Стиль (digital art, anime, cinematic, photorealistic и др).
-Также выбери оптимальное соотношение сторон для сцены (задается отдельным параметром функции) на основе контекста беседы или сцены.

prompts/private_chat.txt

@@ -1,5 +0,0 @@
-Ты - ИИ-помощник без цензуры в чате {platform} c пользователем.
-Отвечай на вопросы и поддерживай контекст беседы.
-Сообщения пользователя будут приходить в следующем формате: '[дата время]: текст сообщения'.
-При ответе НЕ нужно указывать время.
-НЕ используй разметку Markdown, она не поддерживается мессенджером.

prompts/tools.json

@@ -1,24 +0,0 @@
-[
-    {
-        "type": "function",
-        "function": {
-            "name": "generate_image",
-            "description": "Генерация изображения по описанию",
-            "parameters": {
-                "type": "object",
-                "properties": {
-                    "prompt": {
-                        "type": "string",
-                        "description": "Подробное описание сцены на английском языке БЕЗ технических параметров (соотношение сторон, разрешение)"
-                    },
-                    "aspect_ratio": {
-                        "type": "string",
-                        "enum": ["1:1", "4:3", "3:4", "16:9", "9:16"],
-                        "description": "Соотношение сторон"
-                    }
-                },
-                "required": ["prompt"]
-            }
-        }
-    }
-]

tg/__main__.py

@@ -24,7 +24,8 @@ async def main() -> None:
 
     database.create_database(config['db_connection_string'])
 
-    create_ai_agent(config['openrouter_token'], config['openrouter_model'], config['fal_token'],
+    create_ai_agent(config['openrouter_token'], config['openrouter_model'],
+                    config['fal_token'], config['fal_model'],
                     database.DB, 'tg')
 
     bots: list[Bot] = []

tg/tg_database.py

@@ -50,6 +50,7 @@ class TgDatabase(database.BasicDatabase):
             role VARCHAR(16) NOT NULL,
             text VARCHAR(4000),
             image MEDIUMBLOB,
+            UNIQUE KEY contexts_unique (bot_id, chat_id, message_id),
             CONSTRAINT fk_contexts_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id) ON UPDATE CASCADE ON DELETE CASCADE)
         """)
 

vk/__main__.py

@@ -24,7 +24,8 @@ if __name__ == '__main__':
 
     database.create_database(config['db_connection_string'])
 
-    create_ai_agent(config['openrouter_token'], config['openrouter_model'], config['fal_token'],
+    create_ai_agent(config['openrouter_token'], config['openrouter_model'],
+                    config['fal_token'], config['fal_model'],
                     database.DB, 'vk')
 
     bot = Bot(labeler=handlers.labeler)

vk/vk_database.py

@@ -53,6 +53,7 @@ class VkDatabase(database.BasicDatabase):
             role VARCHAR(16) NOT NULL,
             text VARCHAR(4000),
             image MEDIUMBLOB,
+            UNIQUE KEY contexts_unique (bot_id, chat_id, message_id),
             CONSTRAINT fk_contexts_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id) ON UPDATE CASCADE ON DELETE CASCADE)
         """)