diff --git a/ai_agent.py b/ai_agent.py index 561f212..7c89e6c 100644 --- a/ai_agent.py +++ b/ai_agent.py @@ -16,6 +16,7 @@ from openrouter.utils import BackoffStrategy from fal_client import AsyncClient as FalClient from replicate import Client as ReplicateClient +from tavily import TavilyClient from utils import download_file from database import BasicDatabase @@ -43,7 +44,7 @@ class Message: class AiAgent: def __init__(self, openrouter_token: str, openrouter_model: str, - fal_token: str, replicate_token: str, + fal_token: str, replicate_token: str, tavily_token: str, db: BasicDatabase, platform: str): retry_config = RetryConfig(strategy="backoff", @@ -61,6 +62,7 @@ class AiAgent: retry_config=retry_config) self.client_fal = FalClient(key=fal_token) self.replicate_client = ReplicateClient(api_token=replicate_token) + self.tavily_client = TavilyClient(api_key=tavily_token) @dataclass() class _ToolsArtifacts: @@ -168,7 +170,7 @@ class AiAgent: def _construct_system_prompt(self, is_group_chat: bool, bot_id: int, chat_id: int) -> SystemMessageTypedDict: prompt = self.system_prompt_group_chat if is_group_chat else self.system_prompt_private_chat prompt = prompt.replace('{platform}', 'Telegram' if self.platform == 'tg' else 'VK') - prompt += '\n' + self.system_prompt_image_generation + prompt += '\n' + self.system_prompt_tools bot = self.db.get_bot(bot_id) if bot['ai_prompt'] is not None: @@ -202,7 +204,8 @@ class AiAgent: Callable[[int, int, Dict, AiAgent._ToolsArtifacts], Awaitable[List[ChatMessageContentItemTypedDict]]]] = { "generate_image": self._process_tool_generate_image, - "generate_image_anime": self._process_tool_generate_image_anime + "generate_image_anime": self._process_tool_generate_image_anime, + "tavily_search": self._process_tool_tavily_search } for tool_call in tool_calls: @@ -293,6 +296,30 @@ class AiAgent: print(f"Ошибка генерации изображения: {e}") return _serialize_message_content(text=f"Не удалось сгенерировать изображение: {e}") + async def _process_tool_tavily_search(self, _bot_id: int, _chat_id: int, args: dict, + _artifacts: _ToolsArtifacts) -> List[ChatMessageContentItemTypedDict]: + query = args.get("query", "") + print(f"Веб-поиск: {query}") + + try: + results = self.tavily_client.search(query=query, max_results=5) + + if not results or "results" not in results: + return _serialize_message_content(text="Не удалось получить результаты поиска.") + + answer_parts = [] + for i, result in enumerate(results["results"], 1): + title = result.get("title", "Без названия") + url = result.get("url", "") + content = result.get("content", "") + answer_parts.append(f"{i}. {title}\n {url}\n {content}\n") + + answer = "\n".join(answer_parts) + return _serialize_message_content(text=f"По запросу \"{query}\" найдено:\n\n{answer}") + except Exception as e: + print(f"Ошибка веб-поиска: {e}") + return _serialize_message_content(text=f"Не удалось выполнить веб-поиск: {e}") + async def _async_chat_completion_request(self, **kwargs): try: return await self.client_openrouter.chat.send_async(**kwargs) @@ -331,8 +358,8 @@ class AiAgent: self.system_prompt_group_chat = f.read() with open("prompts/private_chat.md", "r") as f: self.system_prompt_private_chat = f.read() - with open("prompts/image_generation.md", "r") as f: - self.system_prompt_image_generation = f.read() + with open("prompts/tools.md", "r") as f: + self.system_prompt_tools = f.read() with open("prompts/tools.json", "r") as f: self.tools_description = json.loads(f.read()) @@ -341,10 +368,10 @@ agent: AiAgent def create_ai_agent(openrouter_token: str, openrouter_model: str, - fal_token: str, replicate_token: str, + fal_token: str, replicate_token: str, tavily_token: str, db: BasicDatabase, platform: str): global agent - agent = AiAgent(openrouter_token, openrouter_model, fal_token, replicate_token, db, platform) + agent = AiAgent(openrouter_token, openrouter_model, fal_token, replicate_token, tavily_token, db, platform) def _add_message_prefix(text: Optional[str], username: Optional[str] = None) -> str: diff --git a/prompts/tools.json b/prompts/tools.json index 05469f5..d5523fa 100644 --- a/prompts/tools.json +++ b/prompts/tools.json @@ -46,5 +46,22 @@ "required": ["prompt", "negative_prompt"] } } + }, + { + "type": "function", + "function": { + "name": "tavily_search", + "description": "Веб-поиск по теме запроса. Используй для поиска информации о событиях, фактах, новостях, определениях, статистике и других актуальных данных.", + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Запрос для поиска (на русском или английском языке)" + } + }, + "required": ["query"] + } + } } ] diff --git a/prompts/image_generation.md b/prompts/tools.md similarity index 74% rename from prompts/image_generation.md rename to prompts/tools.md index 3ed4a17..6fadcb2 100644 --- a/prompts/image_generation.md +++ b/prompts/tools.md @@ -1,4 +1,6 @@ -# Генерация изображений +# Доступные инструменты + +## Генерация изображений Если пользователь просит "нарисовать" или "показать" что-то, сгенерируй изображение путем вызова одной из функций. При вызове функции не нужно добавлять сообщение - оно будет отброшено. Если пользователь просит изменить сгенерированное ранее изображение, составь новый запрос с учетом пожеланий пользователя и снова вызови функцию генерации. @@ -7,10 +9,12 @@ - Никогда не генерируй ASCII-арты вместо вызова функции. - Никогда не вставляй теги вроде ``, `` или любые плейсхолдеры — это сломает чат! +При вызове функции выбери оптимальное соотношение сторон для сцены (задается отдельным параметром функции) на основе контекста беседы или сцены. + НИКОГДА НЕ добавляй в ответ параметры или код генерации - пользователю это не нужно! Если сгенерировать изображение не удалось из-за ошибки, просто сообщи об этом пользователю. -## Генерация обычных (не аниме) изображений +### Генерация обычных (не аниме) изображений Для генерации используй функцию `generate_image` и составляй запрос на естественном языке по следующей формуле: 1. Объекты сцены. 2. Действие/поза. @@ -18,11 +22,14 @@ 4. Освещение, ракурс, композиция. 5. Стиль (digital art, cinematic, photorealistic и др). -## Генерация изображений в стиле аниме +### Генерация изображений в стиле аниме Для генерации используй функцию `generate_image_anime` и составляй запрос, следуя правилам: 1. Описывай сцену набором тегов Danbooru для SDXL, обязательно разделяй теги запятыми. 2. Положительный запрос должен начинаться с `masterpiece, best quality, amazing quality, 4k, very aesthetic, high resolution, ultra-detailed, absurdres, newest, scenery`, а заканчиваться `depth of field, volumetric lighting`. 3. Отрицательный запрос должен заканчиваться `modern, recent, old, oldest, cartoon, graphic, text, painting, crayon, graphite, abstract, glitch, deformed, mutated, ugly, disfigured, long body, lowres, bad anatomy, bad hands, missing fingers, extra digits, fewer digits, cropped, very displeasing, (worst quality, bad quality:1.2), bad anatomy, sketch, jpeg artifacts, signature, watermark, username, signature, simple background, conjoined, bad ai-generated`. 4. Ты можешь добавлять тегам веса, например: `1girl, (long hair:1.2), pink hair`. -Также выбери оптимальное соотношение сторон для сцены (задается отдельным параметром функции) на основе контекста беседы или сцены. +## Веб-поиск +Для поиска информации о событиях, фактах, новостях, определениях, статистике и других актуальных данных используй функцию `tavily_search`. +- Вызывай функцию поиска, когда нужна актуальная информация из интернета. +- После получения результатов дай пользователю краткую сводку найденной информации. diff --git a/tg/__main__.py b/tg/__main__.py index 2b0e77b..ef682bf 100644 --- a/tg/__main__.py +++ b/tg/__main__.py @@ -26,7 +26,7 @@ async def main() -> None: create_ai_agent(config['openrouter_token'], config['openrouter_model'], config['fal_token'], config['replicate_token'], - database.DB, 'tg') + config['tavily_token'], database.DB, 'tg') bots: list[Bot] = [] for item in database.DB.get_bots(): diff --git a/vk/__main__.py b/vk/__main__.py index 1d6bd75..889ecf1 100644 --- a/vk/__main__.py +++ b/vk/__main__.py @@ -26,7 +26,7 @@ if __name__ == '__main__': create_ai_agent(config['openrouter_token'], config['openrouter_model'], config['fal_token'], config['replicate_token'], - database.DB, 'vk') + config['tavily_token'], database.DB, 'vk') bot = Bot(labeler=handlers.labeler)