diff --git a/assistant_bot/agent.py b/assistant_bot/agent.py index b8d215a..5b69fb0 100644 --- a/assistant_bot/agent.py +++ b/assistant_bot/agent.py @@ -70,30 +70,6 @@ MUTATION_SUCCESS_MESSAGES = { "delete_status": "Отслеживаемый объект{suffix} удалён.", } -MUTATION_ACTION_LABELS = { - "remember": "сохранить информацию в памяти", - "delete_memory": "удалить запись памяти", - "create_note": "сохранить заметку", - "update_note": "обновить заметку", - "delete_note": "удалить заметку", - "create_reminder": "создать напоминание", - "update_reminder": "обновить напоминание", - "cancel_reminder": "отменить напоминание", - "create_status": "создать отслеживаемый объект", - "update_status": "обновить статус объекта", - "delete_status": "удалить отслеживаемый объект", -} - -MUTATION_ERROR_MESSAGES = { - "text is required": "не указан текст", - "valid id is required": "не указан корректный идентификатор", - "when and text are required": "не указаны время или текст", - "when or text is required": "не указаны новое время или текст", - "could not parse reminder time": "не удалось распознать время", - "reminder time must be in the future": "время должно быть в будущем", -} - - def mutating_tool_call_key( name: str, arguments: dict[str, Any], @@ -130,6 +106,12 @@ def render_mutating_tool_results( for item in tool_results ): return None + if any( + not isinstance(item.get("result"), dict) + or item["result"].get("ok") is not True + for item in tool_results + ): + return None messages: list[str] = [] for item in tool_results: @@ -161,16 +143,6 @@ def render_mutating_tool_results( ) continue - raw_error = result.get("error") - reason = ( - MUTATION_ERROR_MESSAGES.get(str(raw_error), str(raw_error)) - if raw_error - else "объект не найден или уже отсутствует" - ) - messages.append( - f"Не удалось {MUTATION_ACTION_LABELS[name]}{suffix}: {reason}." - ) - return " ".join(messages) @@ -235,6 +207,31 @@ def cancel_reminder_is_authorized(current_request: str) -> bool: ) +def note_update_target_is_grounded( + current_request: str, + arguments: dict[str, Any], + listed_note_ids: set[int], +) -> bool: + """Reject note ids guessed without an explicit id or a fresh list result.""" + try: + note_id = int(str(arguments.get("id", "")).strip()) + except ValueError: + return False + if note_id <= 0: + return False + if note_id in listed_note_ids: + return True + + escaped_id = re.escape(str(note_id)) + explicit_id_pattern = re.compile( + rf"(?:[#№]\s*{escaped_id}\b|" + rf"\b(?:id|ид|номер)\s*[:#№]?\s*{escaped_id}\b|" + rf"\bзаметк\w*\s+{escaped_id}\b)", + re.IGNORECASE, + ) + return explicit_id_pattern.search(current_request) is not None + + def build_agent_tool_prompt(_tz: ZoneInfo | None = None) -> str: return ( "Ты работаешь как агент с внутренними tools. Пользователь пишет свободным текстом, " @@ -277,7 +274,13 @@ def build_agent_tool_prompt(_tz: ZoneInfo | None = None) -> str: "и отмены используй соответствующие list_*, delete_* и cancel_reminder.\n" "- Если пользователь просит дополнить заметку или список, используй update_note: " "в text передай полный новый текст, сохранив прежнее содержимое и добавив новые " - "данные. Если id или прежний текст неизвестны, сначала вызови list_notes. " + "данные. Если пользователь не указал id заметки явно в текущем запросе, " + "сначала обязательно вызови list_notes и выбери id из его результата. " + "Если прежний текст неизвестен, также сначала вызови list_notes. " + "Если после list_notes подходящей заметки нет, создай ее через create_note: " + "в text передай понятное название заметки или списка и данные, которые просил " + "добавить пользователь. Не сообщай об отсутствии и не спрашивай подтверждение, " + "если название и новое содержимое однозначны. " "Никогда не используй delete_note для изменения, замены или дополнения заметки.\n" "- Если пользователь просит перенести напоминание или изменить его текст, " "используй update_reminder. Передай только изменяемые поля when и/или text; " @@ -487,6 +490,7 @@ async def run_agent_prompt( try: last_tool_results: list[dict[str, Any]] = [] successful_mutating_calls: set[str] = set() + listed_note_ids: set[int] = set() reset_context: bool | None = None final_answer: str | None = None @@ -548,6 +552,49 @@ async def run_agent_prompt( ) continue + ungrounded_note_updates = [ + call + for call in decision.tool_calls + if ( + str(call.get("name", "")).strip().lower() + == "update_note" + and not note_update_target_is_grounded( + prompt, + call.get("arguments", {}), + listed_note_ids, + ) + ) + ] + if ungrounded_note_updates: + messages.append( + { + "role": "assistant", + "content": raw_answer, + } + ) + messages.append( + { + "role": "user", + "content": json_dumps( + { + "kind": "protocol_error", + "error": ( + "update_note запрещен: id заметки не " + "указан явно в current_request и не " + "получен из list_notes в этом запросе. " + "Сначала вызови только list_notes, затем " + "выбери подходящую заметку из result.items " + "и передай ее полный обновленный текст. " + "Если подходящей заметки нет, создай ее " + "через create_note с понятным названием " + "и данными из current_request." + ), + } + ), + } + ) + continue + tool_results = [] for call in decision.tool_calls: tool_name = str(call.get("name", "")).strip().lower() @@ -582,6 +629,16 @@ async def run_agent_prompt( and result.get("ok") is True ): successful_mutating_calls.add(call_key) + if ( + tool_name == "list_notes" + and result.get("ok") is True + ): + for item in result.get("items", []): + if not isinstance(item, dict): + continue + item_id = item.get("id") + if isinstance(item_id, int) and item_id > 0: + listed_note_ids.add(item_id) tool_results.append( { "name": tool_name, diff --git a/tests/test_core.py b/tests/test_core.py index 2cddaa9..d12fdcb 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -417,7 +417,13 @@ class AgentLoopTests(unittest.IsolatedAsyncioTestCase): '"text":"Список покупок: бумажные полотенца, чеснок"' '}}],"reset_context":false}' ) - ai_client = ScriptedAIClient([bad_delete, corrected_update]) + list_notes = ( + '{"tool_calls":[{"name":"list_notes","arguments":{}}],' + '"reset_context":false}' + ) + ai_client = ScriptedAIClient( + [bad_delete, list_notes, corrected_update] + ) update, context, message = self.make_update_and_context( storage, ai_client, @@ -432,7 +438,7 @@ class AgentLoopTests(unittest.IsolatedAsyncioTestCase): notes = storage.list_notes(42) conversation = storage.list_conversation_messages(42, 100) - self.assertEqual(len(ai_client.calls), 2) + self.assertEqual(len(ai_client.calls), 3) repair_messages = ai_client.calls[1]["messages"] assert isinstance(repair_messages, list) repair_envelope = json.loads(repair_messages[-1]["content"]) @@ -453,6 +459,168 @@ class AgentLoopTests(unittest.IsolatedAsyncioTestCase): f"Заметка #{note_id} обновлена.", ) + async def test_add_to_named_note_resolves_id_before_updating(self) -> None: + with tempfile.TemporaryDirectory() as directory: + storage = create_storage(Path(directory) / "assistant.sqlite3") + obsolete_id = storage.add_note(42, "Старая заметка") + storage.delete_note(42, obsolete_id) + note_id = storage.add_note( + 42, + "Список покупок: бумажные полотенца", + ) + guessed_update = ( + '{"tool_calls":[{"name":"update_note","arguments":{"id":' + f"{obsolete_id}," + '"text":"Список покупок: сахар"' + '}}],"reset_context":false}' + ) + list_notes = ( + '{"tool_calls":[{"name":"list_notes","arguments":{}}],' + '"reset_context":false}' + ) + resolved_update = ( + '{"tool_calls":[{"name":"update_note","arguments":{"id":' + f"{note_id}," + '"text":"Список покупок: бумажные полотенца, сахар"' + '}}],"reset_context":false}' + ) + ai_client = ScriptedAIClient( + [guessed_update, list_notes, resolved_update] + ) + update, context, _message = self.make_update_and_context( + storage, + ai_client, + ) + + await run_agent_prompt( + update, + context, + "Добавь сахар в список покупок", + ) + + notes = storage.list_notes(42) + conversation = storage.list_conversation_messages(42, 100) + + self.assertEqual(len(ai_client.calls), 3) + repair_messages = ai_client.calls[1]["messages"] + assert isinstance(repair_messages, list) + repair_envelope = json.loads(repair_messages[-1]["content"]) + self.assertEqual(repair_envelope["kind"], "protocol_error") + self.assertIn("Сначала вызови только list_notes", repair_envelope["error"]) + update_messages = ai_client.calls[2]["messages"] + assert isinstance(update_messages, list) + list_envelope = json.loads(update_messages[-1]["content"]) + self.assertEqual(list_envelope["kind"], "tool_results") + self.assertEqual( + list_envelope["tool_results"][0]["result"]["items"][0]["id"], + note_id, + ) + self.assertEqual( + notes, + [ + { + "id": note_id, + "text": "Список покупок: бумажные полотенца, сахар", + "created_at": notes[0]["created_at"], + } + ], + ) + self.assertEqual( + conversation[-1]["content"], + f"Заметка #{note_id} обновлена.", + ) + + async def test_add_to_missing_named_note_creates_it(self) -> None: + guessed_update = ( + '{"tool_calls":[{"name":"update_note","arguments":{"id":1,' + '"text":"Список покупок: сахар"}}],"reset_context":false}' + ) + list_notes = ( + '{"tool_calls":[{"name":"list_notes","arguments":{}}],' + '"reset_context":false}' + ) + create_note = ( + '{"tool_calls":[{"name":"create_note","arguments":' + '{"text":"Список покупок: сахар"}}],"reset_context":false}' + ) + ai_client = ScriptedAIClient( + [guessed_update, list_notes, create_note] + ) + + with tempfile.TemporaryDirectory() as directory: + storage = create_storage(Path(directory) / "assistant.sqlite3") + update, context, _message = self.make_update_and_context( + storage, + ai_client, + ) + + await run_agent_prompt( + update, + context, + "Добавь сахар в список покупок", + ) + + notes = storage.list_notes(42) + conversation = storage.list_conversation_messages(42, 100) + + self.assertEqual(len(ai_client.calls), 3) + list_messages = ai_client.calls[2]["messages"] + assert isinstance(list_messages, list) + list_envelope = json.loads(list_messages[-1]["content"]) + self.assertEqual(list_envelope["kind"], "tool_results") + self.assertEqual( + list_envelope["tool_results"][0]["result"]["items"], + [], + ) + self.assertEqual(len(notes), 1) + self.assertEqual(notes[0]["text"], "Список покупок: сахар") + self.assertEqual( + conversation[-1]["content"], + f"Заметка #{notes[0]['id']} сохранена.", + ) + + async def test_failed_mutation_is_returned_to_model_for_recovery(self) -> None: + missing_id = 99 + failed_update = ( + '{"tool_calls":[{"name":"update_note","arguments":{"id":' + f"{missing_id}," + '"text":"Новый текст"' + '}}],"reset_context":false}' + ) + ai_client = ScriptedAIClient( + [ + failed_update, + '{"final":"Заметка #99 не найдена.","reset_context":false}', + ] + ) + + with tempfile.TemporaryDirectory() as directory: + storage = create_storage(Path(directory) / "assistant.sqlite3") + update, context, message = self.make_update_and_context( + storage, + ai_client, + ) + + await run_agent_prompt( + update, + context, + "Замени текст заметки #99 на «Новый текст»", + ) + + self.assertEqual(len(ai_client.calls), 2) + recovery_messages = ai_client.calls[1]["messages"] + assert isinstance(recovery_messages, list) + tool_envelope = json.loads(recovery_messages[-1]["content"]) + self.assertEqual(tool_envelope["kind"], "tool_results") + self.assertFalse( + tool_envelope["tool_results"][0]["result"]["ok"] + ) + message.reply_text.assert_awaited_once() + self.assertEqual( + message.reply_text.await_args.args[0], + r"Заметка \#99 не найдена\.", + ) + async def test_reschedule_cannot_be_misread_as_cancellation(self) -> None: with tempfile.TemporaryDirectory() as directory: storage = create_storage(Path(directory) / "assistant.sqlite3") @@ -614,6 +782,8 @@ class AgentToolContractTests(unittest.TestCase): "result.items, а не по неполному reference_data", "Не добавляй благодарности", "одним предложением обычного текста только о результате", + "сначала обязательно вызови list_notes", + "Если после list_notes подходящей заметки нет, создай ее", ): with self.subTest(clause=clause): self.assertIn(clause, prompt)