diff --git a/assistant_bot/agent.py b/assistant_bot/agent.py index 4b70d52..324d5cc 100644 --- a/assistant_bot/agent.py +++ b/assistant_bot/agent.py @@ -29,7 +29,9 @@ from .time_utils import format_local_dt, to_utc_iso, utc_now logger = logging.getLogger(__name__) -def build_assistant_context(storage: AssistantStorage, user_id: int, tz: ZoneInfo) -> str: +def build_assistant_context( + storage: AssistantStorage, user_id: int, tz: ZoneInfo +) -> str: memories = storage.list_memories(user_id, limit=20) notes = storage.list_notes(user_id, limit=10) reminders = storage.list_reminders(user_id, limit=10) @@ -82,27 +84,27 @@ def build_agent_tool_prompt(tz: ZoneInfo) -> str: "Если нужно выполнить действие, верни tool_calls. Если действие уже выполнено " "или tool не нужен, верни final.\n\n" "Форматы ответа:\n" - "{\"tool_calls\":[{\"name\":\"create_note\",\"arguments\":{\"text\":\"...\"}}],\"reset_context\":false}\n" - "{\"final\":\"Короткий ответ пользователю\",\"reset_context\":false}\n\n" + '{"tool_calls":[{"name":"create_note","arguments":{"text":"..."}}],"reset_context":false}\n' + '{"final":"Короткий ответ пользователю","reset_context":false}\n\n' "Доступные tools:\n" "- get_current_datetime {} - узнать текущее локальное и UTC-время.\n" - "- remember {\"text\": string} - сохранить важный долгосрочный факт о пользователе.\n" - "- list_memory {\"limit\": number} - показать сохраненную память.\n" - "- delete_memory {\"id\": number} - удалить запись памяти.\n" - "- create_note {\"text\": string} - сохранить заметку.\n" - "- list_notes {\"limit\": number} - показать заметки.\n" - "- delete_note {\"id\": number} - удалить заметку.\n" - "- create_reminder {\"when\": string, \"text\": string} - поставить напоминание. " + '- remember {"text": string} - сохранить важный долгосрочный факт о пользователе.\n' + '- list_memory {"limit": number} - показать сохраненную память.\n' + '- delete_memory {"id": number} - удалить запись памяти.\n' + '- create_note {"text": string} - сохранить заметку.\n' + '- list_notes {"limit": number} - показать заметки.\n' + '- delete_note {"id": number} - удалить заметку.\n' + '- create_reminder {"when": string, "text": string} - поставить напоминание. ' "when можно указывать как '30m', 'через 2 часа', '18:30', '2026-07-16 18:30'. " "Если пользователь говорит 'завтра/послезавтра/через неделю', сам рассчитай " "дату от текущего локального времени и передай 'YYYY-MM-DD HH:MM'.\n" - "- list_reminders {\"limit\": number} - показать активные напоминания.\n" - "- cancel_reminder {\"id\": number} - отменить напоминание.\n" - "- create_status {\"title\": string, \"status\": string} - начать отслеживать статус.\n" - "- list_statuses {\"limit\": number} - показать отслеживаемые статусы.\n" - "- update_status {\"id\": number, \"status\": string} - обновить статус.\n" - "- delete_status {\"id\": number} - удалить отслеживаемый объект.\n\n" - "- search_conversation {\"query\": string, \"limit\": number} - найти старое обсуждение " + '- list_reminders {"limit": number} - показать активные напоминания.\n' + '- cancel_reminder {"id": number} - отменить напоминание.\n' + '- create_status {"title": string, "status": string} - начать отслеживать статус.\n' + '- list_statuses {"limit": number} - показать отслеживаемые статусы.\n' + '- update_status {"id": number, "status": string} - обновить статус.\n' + '- delete_status {"id": number} - удалить отслеживаемый объект.\n\n' + '- search_conversation {"query": string, "limit": number} - найти старое обсуждение ' "во всей сохраненной переписке по содержательным ключевым словам.\n\n" "Правила:\n" "- Для просьб 'запомни', 'сохрани как факт', 'учти на будущее' используй remember.\n" @@ -121,7 +123,9 @@ def build_agent_tool_prompt(tz: ZoneInfo) -> str: def extract_json_object(raw_text: str) -> dict[str, Any] | None: candidates = [raw_text.strip()] - code_block = re.search(r"```(?:json)?\s*(.*?)```", raw_text, re.IGNORECASE | re.DOTALL) + code_block = re.search( + r"```(?:json)?\s*(.*?)```", raw_text, re.IGNORECASE | re.DOTALL + ) if code_block: candidates.insert(0, code_block.group(1).strip()) @@ -211,7 +215,9 @@ def coerce_required_text(arguments: dict[str, Any], key: str) -> str: return str(value).strip() -def parse_agent_reminder(when: str, text: str, tz: ZoneInfo) -> ReminderParseResult | None: +def parse_agent_reminder( + when: str, text: str, tz: ZoneInfo +) -> ReminderParseResult | None: parsed = parse_reminder(f"{when} {text}".strip(), tz) if parsed: return parsed @@ -254,7 +260,9 @@ def execute_agent_tool( return {"ok": True, "id": memory_id, "text": text} if tool_name == "list_memory": - rows = storage.list_memories(user_id, coerce_positive_int(arguments.get("limit"), 20)) + rows = storage.list_memories( + user_id, coerce_positive_int(arguments.get("limit"), 20) + ) return {"ok": True, "items": rows} if tool_name == "delete_memory": @@ -271,7 +279,9 @@ def execute_agent_tool( return {"ok": True, "id": note_id, "text": text} if tool_name == "list_notes": - rows = storage.list_notes(user_id, coerce_positive_int(arguments.get("limit"), 20)) + rows = storage.list_notes( + user_id, coerce_positive_int(arguments.get("limit"), 20) + ) return {"ok": True, "items": rows} if tool_name == "delete_note": @@ -291,7 +301,9 @@ def execute_agent_tool( if parsed.remind_at_utc <= utc_now(): return {"ok": False, "error": "reminder time must be in the future"} - reminder_id = storage.add_reminder(user_id, chat_id, parsed.text, parsed.remind_at_utc) + reminder_id = storage.add_reminder( + user_id, chat_id, parsed.text, parsed.remind_at_utc + ) return { "ok": True, "id": reminder_id, @@ -300,7 +312,9 @@ def execute_agent_tool( } if tool_name == "list_reminders": - rows = storage.list_reminders(user_id, coerce_positive_int(arguments.get("limit"), 20)) + rows = storage.list_reminders( + user_id, coerce_positive_int(arguments.get("limit"), 20) + ) for row in rows: row["remind_at_local"] = format_local_dt(row["remind_at"], tz) return {"ok": True, "items": rows} @@ -320,7 +334,9 @@ def execute_agent_tool( return {"ok": True, "id": item_id, "title": title, "status": status} if tool_name == "list_statuses": - rows = storage.list_tracked_items(user_id, coerce_positive_int(arguments.get("limit"), 30)) + rows = storage.list_tracked_items( + user_id, coerce_positive_int(arguments.get("limit"), 30) + ) return {"ok": True, "items": rows} if tool_name == "update_status": @@ -328,7 +344,11 @@ def execute_agent_tool( status = coerce_required_text(arguments, "status") if item_id is None or not status: return {"ok": False, "error": "valid id and status are required"} - return {"ok": storage.set_tracked_status(user_id, item_id, status), "id": item_id, "status": status} + return { + "ok": storage.set_tracked_status(user_id, item_id, status), + "id": item_id, + "status": status, + } if tool_name == "delete_status": item_id = parse_positive_int(str(arguments.get("id", ""))) @@ -440,7 +460,12 @@ async def run_agent_prompt( ) last_tool_results = tool_results - messages.append({"role": "assistant", "content": json_dumps({"tool_calls": decision.tool_calls})}) + messages.append( + { + "role": "assistant", + "content": json_dumps({"tool_calls": decision.tool_calls}), + } + ) messages.append( { "role": "user", @@ -496,8 +521,7 @@ async def run_agent_prompt( hint = f"Проверь, что Ollama запущена и модель установлена: ollama pull {model}" else: hint = ( - "Проверь YANDEX_CLOUD_FOLDER, YC_API_KEY " - "и доступ к выбранной модели." + "Проверь YANDEX_CLOUD_FOLDER, YC_API_KEY и доступ к выбранной модели." ) await placeholder.edit_text( f"Не получилось вызвать {ai_client.display_name}.\n{exc}\n\n{hint}" diff --git a/assistant_bot/application.py b/assistant_bot/application.py index f5cfc68..fdf705d 100644 --- a/assistant_bot/application.py +++ b/assistant_bot/application.py @@ -37,7 +37,7 @@ from .handlers import ( private_voice, start, ) -from .jobs import post_init +from .jobs import post_init, post_shutdown from .ollama import OllamaClient from .storage import AssistantStorage from .speech import SpeechRecognizer, SpeechTranscriber @@ -59,6 +59,7 @@ def create_application() -> Application: Application.builder() .token(get_bot_token()) .post_init(post_init) + .post_shutdown(post_shutdown) .build() ) application.bot_data["storage"] = storage diff --git a/assistant_bot/jobs.py b/assistant_bot/jobs.py index 93c0a3e..ae794cf 100644 --- a/assistant_bot/jobs.py +++ b/assistant_bot/jobs.py @@ -1,5 +1,6 @@ import asyncio import logging +from contextlib import suppress from zoneinfo import ZoneInfo from telegram.ext import Application @@ -38,5 +39,17 @@ async def reminder_loop(application: Application) -> None: async def post_init(application: Application) -> None: - application.create_task(reminder_loop(application)) + application.bot_data["reminder_task"] = asyncio.create_task( + reminder_loop(application), + name="reminder-loop", + ) + +async def post_shutdown(application: Application) -> None: + task = application.bot_data.pop("reminder_task", None) + if task is None: + return + + task.cancel() + with suppress(asyncio.CancelledError): + await task diff --git a/assistant_bot/yandex_ai.py b/assistant_bot/yandex_ai.py index 0c8db21..f7e36c5 100644 --- a/assistant_bot/yandex_ai.py +++ b/assistant_bot/yandex_ai.py @@ -18,6 +18,30 @@ def describe_yandex_error(exc: Exception) -> str: return str(exc) +def prepare_chat_messages( + messages: list[dict[str, str]], +) -> list[dict[str, str]]: + """Convert messages to the format accepted by Yandex prompt templates.""" + system_parts: list[str] = [] + chat_messages: list[dict[str, str]] = [] + + for message in messages: + role = message.get("role", "user") + content = message.get("content", "") + if role == "system": + system_parts.append(content) + else: + chat_messages.append({"role": role, "content": content}) + + if system_parts: + chat_messages.insert( + 0, + {"role": "system", "content": "\n\n".join(system_parts)}, + ) + + return chat_messages + + def create_async_sdk(folder_id: str) -> Any: try: from yandex_ai_studio_sdk import AsyncAIStudio @@ -83,13 +107,7 @@ class YandexAIClient: messages: list[dict[str, str]], json_mode: bool = False, ) -> str: - sdk_messages = [ - { - "role": message.get("role", "user"), - "content": message.get("content", ""), - } - for message in messages - ] + sdk_messages = prepare_chat_messages(messages) try: completion = self._sdk.chat.completions(self.normalize_model(model)) diff --git a/tests/test_jobs.py b/tests/test_jobs.py new file mode 100644 index 0000000..1fb766d --- /dev/null +++ b/tests/test_jobs.py @@ -0,0 +1,32 @@ +import asyncio +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from assistant_bot.jobs import post_init, post_shutdown + + +class ReminderJobLifecycleTests(unittest.IsolatedAsyncioTestCase): + async def test_reminder_task_is_cancelled_on_shutdown(self) -> None: + started = asyncio.Event() + + async def fake_reminder_loop(application) -> None: + started.set() + await asyncio.Event().wait() + + application = SimpleNamespace(bot_data={}) + with patch("assistant_bot.jobs.reminder_loop", fake_reminder_loop): + await post_init(application) + await started.wait() + + task = application.bot_data["reminder_task"] + self.assertFalse(task.done()) + + await post_shutdown(application) + + self.assertTrue(task.cancelled()) + self.assertNotIn("reminder_task", application.bot_data) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_yandex_ai.py b/tests/test_yandex_ai.py index c827ed4..e8d98b5 100644 --- a/tests/test_yandex_ai.py +++ b/tests/test_yandex_ai.py @@ -7,6 +7,7 @@ from assistant_bot.yandex_ai import ( YandexAIClient, YandexSpeechRecognizer, describe_yandex_error, + prepare_chat_messages, ) @@ -38,6 +39,28 @@ class FakeChatCompletions: class YandexAIClientTests(unittest.IsolatedAsyncioTestCase): + def test_multiple_system_messages_are_merged_at_the_beginning(self) -> None: + self.assertEqual( + prepare_chat_messages( + [ + {"role": "system", "content": "Primary instructions"}, + {"role": "user", "content": "Earlier question"}, + {"role": "system", "content": "Runtime context"}, + {"role": "assistant", "content": "Earlier answer"}, + {"role": "user", "content": "Current question"}, + ] + ), + [ + { + "role": "system", + "content": "Primary instructions\n\nRuntime context", + }, + {"role": "user", "content": "Earlier question"}, + {"role": "assistant", "content": "Earlier answer"}, + {"role": "user", "content": "Current question"}, + ], + ) + def test_http_error_includes_safe_response_body(self) -> None: error = RuntimeError("forbidden") error.response = SimpleNamespace(