multi-bot-tg/tests/test_ai_command.py

from __future__ import annotations

import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, call, patch

from app.agentic import (
    AgentSessionHandle,
    AgentSessionStore,
    AgentSyntaxError,
    _build_context_messages,
    _create_llm,
    _download_message_media,
    _provider,
    render_prompt_template,
    parse_agent_invocation,
    run_agent_turn,
    strip_prompt_comments,
)
from app.commands.builtins import ai, command_registry
from app.config import AiConfig, AiProviderConfig


class AgentInvocationTests(unittest.TestCase):
    def test_glm_uses_chat_completions_and_deepseek_uses_responses(self) -> None:
        provider = AiProviderConfig("https://example.test", "key", "model")
        with (
            patch("app.agentic.ChatCompletionsLLM") as chat_llm,
            patch("app.agentic.OpenAIBasedLLM") as responses_llm,
            patch("app.agentic.AnthropicMessagesLLM") as claude_llm,
            patch("app.agentic.OpenRouterChatLLM") as openrouter_llm,
        ):
            _create_llm(provider, parse_agent_invocation(".glm prompt", AiConfig()))
            _create_llm(provider, parse_agent_invocation(".ds prompt", AiConfig()))
            _create_llm(provider, parse_agent_invocation(".cl prompt", AiConfig()))
            _create_llm(provider, parse_agent_invocation("prompt", AiConfig()))

        chat_llm.assert_called_once_with(
            api_url="https://example.test",
            api_key="key",
            model="model",
        )
        responses_llm.assert_called_once_with(
            api_url="https://example.test",
            api_key="key",
            model="model",
        )
        claude_llm.assert_called_once_with(
            api_url="https://example.test",
            api_key="key",
            model="model",
        )
        openrouter_llm.assert_called_once_with(
            api_url="https://example.test",
            api_key="key",
            model="model",
        )

    def test_prompt_comments_are_removed_without_breaking_urls(self) -> None:
        value = strip_prompt_comments(
            "// first\nkeep https://example.com // trailing\n/* block\ncomment */\nlast"
        )

        self.assertNotIn("first", value)
        self.assertNotIn("trailing", value)
        self.assertNotIn("block", value)
        self.assertIn("https://example.com", value)
        self.assertIn("last", value)

    def test_prompt_template_replaces_nested_values_and_text(self) -> None:
        rendered = render_prompt_template(
            "caller={{caller_user.id}} owner={owner.name} text={{text}}",
            {
                "caller_user": {"id": 42},
                "owner": {"name": "Owner"},
                "text": "old",
                "source_text": "new",
            },
        )

        self.assertEqual(rendered, "caller=42 owner=Owner text=new")

    def test_parses_context_provider_and_thinking_flags(self) -> None:
        invocation = parse_agent_invocation("12 .ds .th explain this", AiConfig())

        self.assertEqual(invocation.context_messages, 12)
        self.assertTrue(invocation.use_deepseek)
        self.assertFalse(invocation.use_claude)
        self.assertFalse(invocation.use_openrouter)
        self.assertTrue(invocation.thinking)
        self.assertEqual(invocation.prompt, "explain this")

        claude = parse_agent_invocation("3 .ds .cl .th summarize", AiConfig())
        self.assertFalse(claude.use_deepseek)
        self.assertTrue(claude.use_claude)
        self.assertFalse(claude.use_openrouter)
        self.assertTrue(claude.thinking)
        self.assertEqual(claude.prompt, "summarize")
        self.assertEqual(claude.context_messages, 3)
        selected = _provider(
            AiConfig(claude=AiProviderConfig("https://api.anthropic.com", "key", "claude-haiku-5-5")),
            claude,
        )
        self.assertEqual(selected.model, "claude-haiku-5-5")

        plain = parse_agent_invocation("just ask", AiConfig())
        self.assertTrue(plain.use_openrouter)
        self.assertFalse(plain.use_deepseek)
        self.assertFalse(plain.use_claude)
        self.assertEqual(
            _provider(AiConfig(), plain).api_url,
            "https://openrouter.ai/api/v1",
        )

        routed = parse_agent_invocation(".cl .or .th route this", AiConfig())
        self.assertFalse(routed.use_claude)
        self.assertTrue(routed.use_openrouter)
        self.assertEqual(routed.prompt, "route this")
        selected_router = _provider(
            AiConfig(
                openrouter=AiProviderConfig(
                    "https://openrouter.ai/api/v1",
                    "key",
                    "anthropic/claude-haiku-5.5",
                )
            ),
            routed,
        )
        self.assertEqual(selected_router.api_url, "https://openrouter.ai/api/v1")
        self.assertEqual(selected_router.model, "anthropic/claude-haiku-5.5")

    def test_ignores_unknown_leading_dot_arguments(self) -> None:
        invocation = parse_agent_invocation(".unknown .ds prompt", AiConfig())

        self.assertTrue(invocation.use_deepseek)
        self.assertEqual(invocation.prompt, "prompt")

    def test_rejects_missing_prompt(self) -> None:
        with self.assertRaises(AgentSyntaxError):
            parse_agent_invocation("10 .ds .th", AiConfig())

    def test_rejects_context_count_above_configured_limit(self) -> None:
        with self.assertRaises(AgentSyntaxError):
            parse_agent_invocation("101 prompt", AiConfig(max_context_messages=100))


class AgentContextTests(unittest.IsolatedAsyncioTestCase):
    async def test_reply_is_injected_once_and_history_is_chronological(self) -> None:
        older = SimpleNamespace(id=1, chat=SimpleNamespace(id=99), text="older")
        newer = SimpleNamespace(id=2, chat=SimpleNamespace(id=99), text="newer")
        command = SimpleNamespace(
            id=3,
            chat=SimpleNamespace(id=99),
            reply_to_message=older,
        )

        async def history():
            yield newer
            yield older

        client = SimpleNamespace(get_chat_history=lambda *args, **kwargs: history())
        invocation = parse_agent_invocation("2 prompt", AiConfig())
        with tempfile.TemporaryDirectory() as temp_dir:
            messages = await _build_context_messages(
                client,
                command,
                invocation,
                AiConfig(),
                Path(temp_dir),
                user_template="wrapped={{message.is_reply}} {{text}}",
                template_context={},
            )

        self.assertEqual(len(messages), 2)
        self.assertIn("wrapped=true", messages[0].message())
        self.assertIn("older", messages[0].message())
        self.assertIn("newer", messages[1].message())

    async def test_media_over_limit_is_not_downloaded(self) -> None:
        message = SimpleNamespace(
            id=1,
            document=SimpleNamespace(file_size=6, mime_type="text/plain"),
        )
        client = SimpleNamespace(download_media=AsyncMock())

        with tempfile.TemporaryDirectory() as temp_dir:
            media = await _download_message_media(client, message, Path(temp_dir), 5)

        self.assertEqual(media, [])
        client.download_media.assert_not_awaited()

    async def test_media_within_limit_is_attached(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            path = Path(temp_dir) / "note.txt"
            path.write_bytes(b"hello")
            message = SimpleNamespace(
                id=1,
                document=SimpleNamespace(
                    file_size=5,
                    mime_type="text/plain",
                    file_name="note.txt",
                ),
            )
            client = SimpleNamespace(download_media=AsyncMock(return_value=str(path)))

            media = await _download_message_media(client, message, Path(temp_dir), 5)

        self.assertEqual(len(media), 1)
        self.assertEqual(media[0].data, b"hello")
        self.assertEqual(media[0].kind, "file")

    async def test_deepseek_and_reasoning_are_passed_to_agent(self) -> None:
        config = AiConfig(
            glm=AiProviderConfig("https://glm.example", "glm-key", "glm-model"),
            deepseek=AiProviderConfig("https://ds.example", "ds-key", "ds-model"),
            brave_api_key="brave-key",
        )
        invocation = parse_agent_invocation(".ds .th prompt", config)
        command = SimpleNamespace(
            id=3,
            chat=SimpleNamespace(id=99),
            from_user=SimpleNamespace(id=42),
            reply_to_message=None,
            photo=SimpleNamespace(file_size=100),
        )
        account = SimpleNamespace(
            me=SimpleNamespace(id=7, full_name="Owner", username="owner"),
            config=SimpleNamespace(name="personal", session="personal"),
        )
        final_message = SimpleNamespace(message=lambda: "answer")
        session = SimpleNamespace(
            run=AsyncMock(
                return_value=SimpleNamespace(
                    succeeded=True,
                    final_message=final_message,
                    error=None,
                )
            )
        )
        agent = MagicMock()
        agent.create_session.return_value = session
        attached_media = MagicMock()

        with (
            patch("app.agentic.OpenAIBasedLLM") as llm_class,
            patch("app.agentic.Aggregator", return_value=agent) as aggregator_class,
            patch("app.agentic.BraveSearchTools") as brave_tools,
            patch("app.agentic.WebPageTools") as web_page_tools,
            patch("app.agentic._build_context_messages", AsyncMock(return_value=[])),
            patch(
                "app.agentic._download_message_media",
                AsyncMock(return_value=[attached_media]),
            ),
        ):
            brave_tools.return_value.tools.return_value = []
            web_page_tools.return_value.tools.return_value = []
            result = await run_agent_turn(
                SimpleNamespace(),
                command,
                account,
                AgentSessionStore(),
                invocation,
                config,
            )

        self.assertEqual(result.answer, "answer")
        llm_class.assert_called_once_with(
            api_url="https://ds.example",
            api_key="ds-key",
            model="ds-model",
        )
        self.assertTrue(aggregator_class.call_args.kwargs["thinking"])
        session.run.assert_awaited_once_with("prompt", media=[attached_media])
        brave_tools.assert_called_once_with(api_key="brave-key", timeout_seconds=10.0)
        web_page_tools.assert_called_once_with(timeout_seconds=45.0, max_chars=60_000)
        result.handle.close()

    async def test_reply_to_ai_response_reuses_existing_session(self) -> None:
        workspace = MagicMock()
        final_message = SimpleNamespace(message=lambda: "continued")
        session = SimpleNamespace(
            context={"old": True},
            run=AsyncMock(
                return_value=SimpleNamespace(
                    succeeded=True,
                    final_message=final_message,
                    error=None,
                )
            ),
        )
        handle = AgentSessionHandle(session, workspace, set())
        store = AgentSessionStore()
        ai_response = SimpleNamespace(id=50, chat=SimpleNamespace(id=99))
        store.register_response(handle, ai_response)
        command = SimpleNamespace(
            id=51,
            chat=SimpleNamespace(id=99),
            from_user=SimpleNamespace(id=7, full_name="Owner"),
            reply_to_message=ai_response,
        )
        account = SimpleNamespace(
            me=SimpleNamespace(id=7, full_name="Owner"),
            config=SimpleNamespace(name="personal", session="personal"),
        )

        result = await run_agent_turn(
            SimpleNamespace(),
            command,
            account,
            store,
            parse_agent_invocation("continue", AiConfig()),
            AiConfig(),
        )

        self.assertIs(result.handle, handle)
        self.assertEqual(result.answer, "continued")
        session.run.assert_awaited_once_with("continue")
        self.assertTrue(session.context["is_owner"])
        store.clear()
        workspace.cleanup.assert_called_once()


class AiCommandTests(unittest.IsolatedAsyncioTestCase):
    def _ctx(self, args: str = "prompt") -> SimpleNamespace:
        sent = SimpleNamespace(id=124, chat=SimpleNamespace(id=99))
        message = SimpleNamespace(
            id=123,
            react=AsyncMock(),
            reply_text=AsyncMock(return_value=sent),
        )
        session_store = MagicMock()
        return SimpleNamespace(
            args=args,
            message=message,
            client=SimpleNamespace(),
            account=SimpleNamespace(
                me=SimpleNamespace(id=7, full_name="Owner"),
                config=SimpleNamespace(name="personal", session="personal"),
                ai_sessions=session_store,
            ),
            manager=SimpleNamespace(config=SimpleNamespace(ai=AiConfig())),
            reply=message.reply_text,
        )

    def test_ai_is_registered_as_shared(self) -> None:
        definition = command_registry.resolve("ai")

        self.assertIsNotNone(definition)
        self.assertTrue(definition.is_shared)
        self.assertEqual(definition.guest_prefix, "..")

    async def test_success_replies_with_agent_answer(self) -> None:
        ctx = self._ctx()
        handle = MagicMock()
        with patch(
            "app.commands.builtins.run_agent_turn",
            AsyncMock(return_value=SimpleNamespace(answer="answer", handle=handle)),
        ):
            await ai(ctx)

        ctx.reply.assert_awaited_once_with("answer")
        ctx.account.ai_sessions.register_response.assert_called_once_with(
            handle,
            ctx.reply.return_value,
        )
        self.assertEqual(ctx.message.react.await_args_list, [call("🤔"), call(None)])

    async def test_invalid_syntax_only_dislikes(self) -> None:
        ctx = self._ctx(".ds .th")

        with self.assertLogs("app.commands.builtins", level="ERROR"):
            await ai(ctx)

        self.assertEqual(ctx.message.react.await_args_list, [call("🤔"), call("👎")])
        ctx.reply.assert_not_awaited()

    async def test_agent_error_only_dislikes(self) -> None:
        ctx = self._ctx()
        with patch(
            "app.commands.builtins.run_agent_turn",
            AsyncMock(side_effect=RuntimeError("API failed")),
        ):
            with self.assertLogs("app.commands.builtins", level="ERROR"):
                await ai(ctx)

        self.assertEqual(ctx.message.react.await_args_list, [call("🤔"), call("👎")])
        ctx.reply.assert_not_awaited()

    async def test_thinking_reaction_failure_does_not_stop_ai(self) -> None:
        ctx = self._ctx()
        ctx.message.react.side_effect = [RuntimeError("reactions disabled"), None]

        with (
            patch(
                "app.commands.builtins.run_agent_turn",
                AsyncMock(
                    return_value=SimpleNamespace(answer="answer", handle=MagicMock())
                ),
            ),
            self.assertLogs("app.commands.builtins", level="ERROR"),
        ):
            await ai(ctx)

        ctx.reply.assert_awaited_once_with("answer")
        self.assertEqual(ctx.message.react.await_args_list, [call("🤔"), call(None)])


if __name__ == "__main__":
    unittest.main()