This commit is contained in:
kandrusyak
2026-07-25 15:39:03 +03:00
parent 0163208484
commit 3f63d0305c
6 changed files with 148 additions and 37 deletions

View File

@@ -29,7 +29,9 @@ from .time_utils import format_local_dt, to_utc_iso, utc_now
logger = logging.getLogger(__name__) 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) memories = storage.list_memories(user_id, limit=20)
notes = storage.list_notes(user_id, limit=10) notes = storage.list_notes(user_id, limit=10)
reminders = storage.list_reminders(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_calls. Если действие уже выполнено "
"или tool не нужен, верни final.\n\n" "или tool не нужен, верни final.\n\n"
"Форматы ответа:\n" "Форматы ответа:\n"
"{\"tool_calls\":[{\"name\":\"create_note\",\"arguments\":{\"text\":\"...\"}}],\"reset_context\":false}\n" '{"tool_calls":[{"name":"create_note","arguments":{"text":"..."}}],"reset_context":false}\n'
"{\"final\":\"Короткий ответ пользователю\",\"reset_context\":false}\n\n" '{"final":"Короткий ответ пользователю","reset_context":false}\n\n'
"Доступные tools:\n" "Доступные tools:\n"
"- get_current_datetime {} - узнать текущее локальное и UTC-время.\n" "- get_current_datetime {} - узнать текущее локальное и UTC-время.\n"
"- remember {\"text\": string} - сохранить важный долгосрочный факт о пользователе.\n" '- remember {"text": string} - сохранить важный долгосрочный факт о пользователе.\n'
"- list_memory {\"limit\": number} - показать сохраненную память.\n" '- list_memory {"limit": number} - показать сохраненную память.\n'
"- delete_memory {\"id\": number} - удалить запись памяти.\n" '- delete_memory {"id": number} - удалить запись памяти.\n'
"- create_note {\"text\": string} - сохранить заметку.\n" '- create_note {"text": string} - сохранить заметку.\n'
"- list_notes {\"limit\": number} - показать заметки.\n" '- list_notes {"limit": number} - показать заметки.\n'
"- delete_note {\"id\": number} - удалить заметку.\n" '- delete_note {"id": number} - удалить заметку.\n'
"- create_reminder {\"when\": string, \"text\": string} - поставить напоминание. " '- create_reminder {"when": string, "text": string} - поставить напоминание. '
"when можно указывать как '30m', 'через 2 часа', '18:30', '2026-07-16 18:30'. " "when можно указывать как '30m', 'через 2 часа', '18:30', '2026-07-16 18:30'. "
"Если пользователь говорит 'завтра/послезавтра/через неделю', сам рассчитай " "Если пользователь говорит 'завтра/послезавтра/через неделю', сам рассчитай "
"дату от текущего локального времени и передай 'YYYY-MM-DD HH:MM'.\n" "дату от текущего локального времени и передай 'YYYY-MM-DD HH:MM'.\n"
"- list_reminders {\"limit\": number} - показать активные напоминания.\n" '- list_reminders {"limit": number} - показать активные напоминания.\n'
"- cancel_reminder {\"id\": number} - отменить напоминание.\n" '- cancel_reminder {"id": number} - отменить напоминание.\n'
"- create_status {\"title\": string, \"status\": string} - начать отслеживать статус.\n" '- create_status {"title": string, "status": string} - начать отслеживать статус.\n'
"- list_statuses {\"limit\": number} - показать отслеживаемые статусы.\n" '- list_statuses {"limit": number} - показать отслеживаемые статусы.\n'
"- update_status {\"id\": number, \"status\": string} - обновить статус.\n" '- update_status {"id": number, "status": string} - обновить статус.\n'
"- delete_status {\"id\": number} - удалить отслеживаемый объект.\n\n" '- delete_status {"id": number} - удалить отслеживаемый объект.\n\n'
"- search_conversation {\"query\": string, \"limit\": number} - найти старое обсуждение " '- search_conversation {"query": string, "limit": number} - найти старое обсуждение '
"во всей сохраненной переписке по содержательным ключевым словам.\n\n" "во всей сохраненной переписке по содержательным ключевым словам.\n\n"
"Правила:\n" "Правила:\n"
"- Для просьб 'запомни', 'сохрани как факт', 'учти на будущее' используй remember.\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: def extract_json_object(raw_text: str) -> dict[str, Any] | None:
candidates = [raw_text.strip()] 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: if code_block:
candidates.insert(0, code_block.group(1).strip()) 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() 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) parsed = parse_reminder(f"{when} {text}".strip(), tz)
if parsed: if parsed:
return parsed return parsed
@@ -254,7 +260,9 @@ def execute_agent_tool(
return {"ok": True, "id": memory_id, "text": text} return {"ok": True, "id": memory_id, "text": text}
if tool_name == "list_memory": 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} return {"ok": True, "items": rows}
if tool_name == "delete_memory": if tool_name == "delete_memory":
@@ -271,7 +279,9 @@ def execute_agent_tool(
return {"ok": True, "id": note_id, "text": text} return {"ok": True, "id": note_id, "text": text}
if tool_name == "list_notes": 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} return {"ok": True, "items": rows}
if tool_name == "delete_note": if tool_name == "delete_note":
@@ -291,7 +301,9 @@ def execute_agent_tool(
if parsed.remind_at_utc <= utc_now(): if parsed.remind_at_utc <= utc_now():
return {"ok": False, "error": "reminder time must be in the future"} 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 { return {
"ok": True, "ok": True,
"id": reminder_id, "id": reminder_id,
@@ -300,7 +312,9 @@ def execute_agent_tool(
} }
if tool_name == "list_reminders": 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: for row in rows:
row["remind_at_local"] = format_local_dt(row["remind_at"], tz) row["remind_at_local"] = format_local_dt(row["remind_at"], tz)
return {"ok": True, "items": rows} return {"ok": True, "items": rows}
@@ -320,7 +334,9 @@ def execute_agent_tool(
return {"ok": True, "id": item_id, "title": title, "status": status} return {"ok": True, "id": item_id, "title": title, "status": status}
if tool_name == "list_statuses": 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} return {"ok": True, "items": rows}
if tool_name == "update_status": if tool_name == "update_status":
@@ -328,7 +344,11 @@ def execute_agent_tool(
status = coerce_required_text(arguments, "status") status = coerce_required_text(arguments, "status")
if item_id is None or not status: if item_id is None or not status:
return {"ok": False, "error": "valid id and status are required"} 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": if tool_name == "delete_status":
item_id = parse_positive_int(str(arguments.get("id", ""))) item_id = parse_positive_int(str(arguments.get("id", "")))
@@ -440,7 +460,12 @@ async def run_agent_prompt(
) )
last_tool_results = tool_results 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( messages.append(
{ {
"role": "user", "role": "user",
@@ -496,8 +521,7 @@ async def run_agent_prompt(
hint = f"Проверь, что Ollama запущена и модель установлена: ollama pull {model}" hint = f"Проверь, что Ollama запущена и модель установлена: ollama pull {model}"
else: else:
hint = ( hint = (
"Проверь YANDEX_CLOUD_FOLDER, YC_API_KEY " "Проверь YANDEX_CLOUD_FOLDER, YC_API_KEY и доступ к выбранной модели."
"и доступ к выбранной модели."
) )
await placeholder.edit_text( await placeholder.edit_text(
f"Не получилось вызвать {ai_client.display_name}.\n{exc}\n\n{hint}" f"Не получилось вызвать {ai_client.display_name}.\n{exc}\n\n{hint}"

View File

@@ -37,7 +37,7 @@ from .handlers import (
private_voice, private_voice,
start, start,
) )
from .jobs import post_init from .jobs import post_init, post_shutdown
from .ollama import OllamaClient from .ollama import OllamaClient
from .storage import AssistantStorage from .storage import AssistantStorage
from .speech import SpeechRecognizer, SpeechTranscriber from .speech import SpeechRecognizer, SpeechTranscriber
@@ -59,6 +59,7 @@ def create_application() -> Application:
Application.builder() Application.builder()
.token(get_bot_token()) .token(get_bot_token())
.post_init(post_init) .post_init(post_init)
.post_shutdown(post_shutdown)
.build() .build()
) )
application.bot_data["storage"] = storage application.bot_data["storage"] = storage

View File

@@ -1,5 +1,6 @@
import asyncio import asyncio
import logging import logging
from contextlib import suppress
from zoneinfo import ZoneInfo from zoneinfo import ZoneInfo
from telegram.ext import Application from telegram.ext import Application
@@ -38,5 +39,17 @@ async def reminder_loop(application: Application) -> None:
async def post_init(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

View File

@@ -18,6 +18,30 @@ def describe_yandex_error(exc: Exception) -> str:
return str(exc) 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: def create_async_sdk(folder_id: str) -> Any:
try: try:
from yandex_ai_studio_sdk import AsyncAIStudio from yandex_ai_studio_sdk import AsyncAIStudio
@@ -83,13 +107,7 @@ class YandexAIClient:
messages: list[dict[str, str]], messages: list[dict[str, str]],
json_mode: bool = False, json_mode: bool = False,
) -> str: ) -> str:
sdk_messages = [ sdk_messages = prepare_chat_messages(messages)
{
"role": message.get("role", "user"),
"content": message.get("content", ""),
}
for message in messages
]
try: try:
completion = self._sdk.chat.completions(self.normalize_model(model)) completion = self._sdk.chat.completions(self.normalize_model(model))

32
tests/test_jobs.py Normal file
View File

@@ -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()

View File

@@ -7,6 +7,7 @@ from assistant_bot.yandex_ai import (
YandexAIClient, YandexAIClient,
YandexSpeechRecognizer, YandexSpeechRecognizer,
describe_yandex_error, describe_yandex_error,
prepare_chat_messages,
) )
@@ -38,6 +39,28 @@ class FakeChatCompletions:
class YandexAIClientTests(unittest.IsolatedAsyncioTestCase): 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: def test_http_error_includes_safe_response_body(self) -> None:
error = RuntimeError("forbidden") error = RuntimeError("forbidden")
error.response = SimpleNamespace( error.response = SimpleNamespace(