fix
This commit is contained in:
@@ -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}"
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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))
|
||||
|
||||
32
tests/test_jobs.py
Normal file
32
tests/test_jobs.py
Normal 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()
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user