aixtools 0.1.9__py3-none-any.whl → 0.1.11__py3-none-any.whl

tests/unit/a2a/google_sdk/test_utils.py

@@ -1,208 +0,0 @@
-"""Tests for the A2A utils module."""
-
-import unittest
-from unittest.mock import AsyncMock, MagicMock, patch
-
-import httpx
-from a2a.client import ClientConfig, ClientFactory, A2ACardResolver
-from a2a.server.agent_execution import RequestContext
-from a2a.types import AgentCard
-
-from aixtools.a2a.google_sdk.utils import (
-    _AgentCardResolver,
-    get_a2a_clients,
-    get_session_id_tuple,
-)
-from aixtools.a2a.google_sdk.remote_agent_connection import RemoteAgentConnection
-from aixtools.context import DEFAULT_USER_ID, DEFAULT_SESSION_ID
-
-
-class TestAgentCardResolver(unittest.IsolatedAsyncioTestCase):
-    """Tests for the _AgentCardResolver class."""
-
-    def setUp(self):
-        self.mock_client = AsyncMock(spec=httpx.AsyncClient)
-        self.resolver = _AgentCardResolver(self.mock_client)
-
-    @patch("aixtools.a2a.google_sdk.utils.ClientFactory")
-    def test_init(self, mock_client_factory_class):
-        """Test _AgentCardResolver initialization."""
-        mock_factory = MagicMock(spec=ClientFactory)
-        mock_client_factory_class.return_value = mock_factory
-        
-        resolver = _AgentCardResolver(self.mock_client)
-        
-        # Verify ClientFactory was created with correct config
-        mock_client_factory_class.assert_called_once()
-        call_args = mock_client_factory_class.call_args[0][0]
-        self.assertIsInstance(call_args, ClientConfig)
-        self.assertEqual(call_args.httpx_client, self.mock_client)
-        
-        # Verify attributes are set
-        self.assertEqual(resolver._httpx_client, self.mock_client)
-        self.assertEqual(resolver._a2a_client_factory, mock_factory)
-        self.assertEqual(resolver.clients, {})
-
-    @patch("aixtools.a2a.google_sdk.utils.RemoteAgentConnection")
-    def test_register_agent_card(self, mock_connection_class):
-        """Test registering an agent card."""
-        mock_card = MagicMock(spec=AgentCard)
-        mock_card.name = "test_agent"
-        
-        mock_client = MagicMock()
-        mock_factory = MagicMock(spec=ClientFactory)
-        mock_factory.create.return_value = mock_client
-        self.resolver._a2a_client_factory = mock_factory
-        
-        mock_connection = MagicMock(spec=RemoteAgentConnection)
-        mock_connection_class.return_value = mock_connection
-        
-        self.resolver.register_agent_card(mock_card)
-        
-        # Verify client was created
-        mock_factory.create.assert_called_once_with(mock_card)
-        
-        # Verify RemoteAgentConnection was created
-        mock_connection_class.assert_called_once_with(mock_card, mock_client)
-        
-        # Verify connection was stored
-        self.assertEqual(self.resolver.clients["test_agent"], mock_connection)
-
-    @patch("aixtools.a2a.google_sdk.utils.A2ACardResolver")
-    async def test_retrieve_card(self, mock_resolver_class):
-        """Test retrieving a card from an address."""
-        mock_resolver = AsyncMock(spec=A2ACardResolver)
-        mock_resolver_class.return_value = mock_resolver
-        
-        mock_card = MagicMock(spec=AgentCard)
-        mock_card.name = "test_agent"
-        mock_resolver.get_agent_card.return_value = mock_card
-        
-        with patch.object(self.resolver, 'register_agent_card') as mock_register:
-            await self.resolver.retrieve_card("http://test.com")
-            
-            # Verify resolver was created correctly (it tries the first card path)
-            mock_resolver_class.assert_called_with(self.mock_client, "http://test.com", "/.well-known/agent-card.json")
-            
-            # Verify card was retrieved
-            mock_resolver.get_agent_card.assert_called_once()
-            
-            # Verify card was registered
-            mock_register.assert_called_once_with(mock_card)
-
-    async def test_get_a2a_clients(self):
-        """Test getting A2A clients for multiple hosts."""
-        agent_hosts = ["http://agent1.com", "http://agent2.com"]
-        
-        with patch.object(self.resolver, 'retrieve_card') as mock_retrieve:
-            mock_retrieve.return_value = None  # Mock async function
-            
-            # Set up some mock clients
-            mock_connection1 = MagicMock(spec=RemoteAgentConnection)
-            mock_connection2 = MagicMock(spec=RemoteAgentConnection)
-            self.resolver.clients = {
-                "agent1": mock_connection1,
-                "agent2": mock_connection2
-            }
-            
-            result = await self.resolver.get_a2a_clients(agent_hosts)
-            
-            # Verify retrieve_card was called for each host
-            self.assertEqual(mock_retrieve.call_count, 2)
-            mock_retrieve.assert_any_call("http://agent1.com")
-            mock_retrieve.assert_any_call("http://agent2.com")
-            
-            # Verify result contains the clients
-            self.assertEqual(result, self.resolver.clients)
-
-
-class TestGetA2AClients(unittest.IsolatedAsyncioTestCase):
-    """Tests for the get_a2a_clients function."""
-
-    @patch("aixtools.a2a.google_sdk.utils._AgentCardResolver")
-    @patch("aixtools.a2a.google_sdk.utils.httpx.AsyncClient")
-    async def test_get_a2a_clients(self, mock_client_class, mock_resolver_class):
-        """Test the get_a2a_clients function."""
-        mock_client = AsyncMock()
-        mock_client_class.return_value = mock_client
-        
-        mock_resolver = AsyncMock(spec=_AgentCardResolver)
-        mock_resolver_class.return_value = mock_resolver
-        
-        mock_clients = {"agent1": MagicMock(), "agent2": MagicMock()}
-        mock_resolver.get_a2a_clients.return_value = mock_clients
-        
-        ctx = ("user123", "session456")
-        agent_hosts = ["http://agent1.com", "http://agent2.com"]
-        
-        result = await get_a2a_clients(ctx, agent_hosts)
-        
-        # Verify httpx client was created with correct headers
-        expected_headers = {
-            "user-id": "user123",
-            "session-id": "session456",
-        }
-        mock_client_class.assert_called_once_with(headers=expected_headers, timeout=60.0)
-        
-        # Verify resolver was created with the client
-        mock_resolver_class.assert_called_once_with(mock_client)
-        
-        # Verify get_a2a_clients was called with the hosts
-        mock_resolver.get_a2a_clients.assert_called_once_with(agent_hosts)
-        
-        # Verify result
-        self.assertEqual(result, mock_clients)
-
-
-class TestGetSessionIdTuple(unittest.TestCase):
-    """Tests for the get_session_id_tuple function."""
-
-    def test_get_session_id_tuple_with_headers(self):
-        """Test getting session ID tuple when headers are present."""
-        mock_context = MagicMock(spec=RequestContext)
-        mock_context.call_context.state = {
-            "headers": {
-                "user-id": "test_user",
-                "session-id": "test_session"
-            }
-        }
-        
-        result = get_session_id_tuple(mock_context)
-        
-        self.assertEqual(result, ("test_user", "test_session"))
-
-    def test_get_session_id_tuple_with_partial_headers(self):
-        """Test getting session ID tuple when only some headers are present."""
-        mock_context = MagicMock(spec=RequestContext)
-        mock_context.call_context.state = {
-            "headers": {
-                "user-id": "test_user"
-                # session-id is missing
-            }
-        }
-        
-        result = get_session_id_tuple(mock_context)
-        
-        self.assertEqual(result, ("test_user", DEFAULT_SESSION_ID))
-
-    def test_get_session_id_tuple_no_headers(self):
-        """Test getting session ID tuple when no headers are present."""
-        mock_context = MagicMock(spec=RequestContext)
-        mock_context.call_context.state = {}
-        
-        result = get_session_id_tuple(mock_context)
-        
-        self.assertEqual(result, (DEFAULT_USER_ID, DEFAULT_SESSION_ID))
-
-    def test_get_session_id_tuple_empty_headers(self):
-        """Test getting session ID tuple when headers dict is empty."""
-        mock_context = MagicMock(spec=RequestContext)
-        mock_context.call_context.state = {"headers": {}}
-        
-        result = get_session_id_tuple(mock_context)
-        
-        self.assertEqual(result, (DEFAULT_USER_ID, DEFAULT_SESSION_ID))
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/unit/agents/test_prompt.py

@@ -1,363 +0,0 @@
-"""Unit tests for aixtools.agents.prompt module."""
-
-import tempfile
-import unittest
-from pathlib import Path
-from unittest.mock import Mock, patch, mock_open
-
-from pydantic_ai import BinaryContent
-
-from aixtools.agents.prompt import (
-    CLAUDE_IMAGE_MAX_FILE_SIZE_IN_CONTEXT,
-    CLAUDE_MAX_FILE_SIZE_IN_CONTEXT,
-    build_user_input,
-    file_to_binary_content,
-    should_be_included_into_context,
-)
-
-
-class TestShouldBeIncludedIntoContext(unittest.TestCase):
-    """Test cases for should_be_included_into_context function."""
-
-    def test_non_binary_content_returns_false(self):
-        """Test that non-BinaryContent returns False."""
-        result = should_be_included_into_context("text content", 100)
-        self.assertFalse(result)
-        
-        result = should_be_included_into_context(None, 100)
-        self.assertFalse(result)
-
-    def test_text_media_type_returns_false(self):
-        """Test that text media types return False."""
-        binary_content = BinaryContent(data=b"test", media_type="text/plain")
-        result = should_be_included_into_context(binary_content, 100)
-        self.assertFalse(result)
-        
-        binary_content = BinaryContent(data=b"test", media_type="text/html")
-        result = should_be_included_into_context(binary_content, 100)
-        self.assertFalse(result)
-
-    def test_archive_types_return_false(self):
-        """Test that archive media types return False."""
-        archive_types = [
-            "application/zip",
-            "application/x-tar",
-            "application/gzip",
-            "application/x-gzip",
-            "application/x-rar-compressed",
-            "application/x-7z-compressed",
-        ]
-        
-        for media_type in archive_types:
-            binary_content = BinaryContent(data=b"test", media_type=media_type)
-            result = should_be_included_into_context(binary_content, 100)
-            self.assertFalse(result, f"Archive type {media_type} should return False")
-
-    def test_image_within_size_limit_returns_true(self):
-        """Test that images within size limit return True."""
-        with patch.object(BinaryContent, 'is_image', new_callable=lambda: True):
-            binary_content = BinaryContent(data=b"image_data", media_type="image/png")
-            
-            # Test with size under limit
-            result = should_be_included_into_context(binary_content, 1024)
-            self.assertTrue(result)
-
-    def test_image_over_size_limit_returns_false(self):
-        """Test that images over size limit return False."""
-        with patch.object(BinaryContent, 'is_image', new_callable=lambda: True):
-            binary_content = BinaryContent(data=b"image_data", media_type="image/png")
-            
-            # Test with size over limit
-            result = should_be_included_into_context(binary_content, CLAUDE_IMAGE_MAX_FILE_SIZE_IN_CONTEXT + 1)
-            self.assertFalse(result)
-
-    def test_non_image_within_size_limit_returns_true(self):
-        """Test that non-images within size limit return True."""
-        with patch.object(BinaryContent, 'is_image', new_callable=lambda: False):
-            binary_content = BinaryContent(data=b"pdf_data", media_type="application/pdf")
-            
-            # Test with size under limit
-            result = should_be_included_into_context(binary_content, 1024)
-            self.assertTrue(result)
-
-    def test_non_image_over_size_limit_returns_false(self):
-        """Test that non-images over size limit return False."""
-        with patch.object(BinaryContent, 'is_image', new_callable=lambda: False):
-            binary_content = BinaryContent(data=b"pdf_data", media_type="application/pdf")
-            
-            # Test with size over limit
-            result = should_be_included_into_context(binary_content, CLAUDE_MAX_FILE_SIZE_IN_CONTEXT + 1)
-            self.assertFalse(result)
-
-    def test_custom_size_limits(self):
-        """Test with custom size limits."""
-        # Test non-image content with custom limits (since image detection is complex)
-        pdf_content = BinaryContent(data=b"pdf_data", media_type="application/pdf")
-        
-        # Test non-image over custom file limit
-        result = should_be_included_into_context(
-            pdf_content, 5000, max_img_size_bytes=1024, max_file_size_bytes=4096
-        )
-        self.assertFalse(result)
-        
-        # Test non-image under custom file limit
-        result = should_be_included_into_context(
-            pdf_content, 2000, max_img_size_bytes=1024, max_file_size_bytes=4096
-        )
-        self.assertTrue(result)
-        
-        # Test with very small custom limits
-        result = should_be_included_into_context(
-            pdf_content, 100, max_img_size_bytes=50, max_file_size_bytes=80
-        )
-        self.assertFalse(result)
-        
-        # Test with very large custom limits
-        result = should_be_included_into_context(
-            pdf_content, 100, max_img_size_bytes=200, max_file_size_bytes=200
-        )
-        self.assertTrue(result)
-
-
-class TestFileToBinaryContent(unittest.TestCase):
-    """Test cases for file_to_binary_content function."""
-
-    def setUp(self):
-        """Set up test fixtures."""
-        self.temp_dir = tempfile.mkdtemp()
-        self.temp_path = Path(self.temp_dir)
-
-    def tearDown(self):
-        """Clean up test fixtures."""
-        import shutil
-        shutil.rmtree(self.temp_dir)
-
-    @patch('aixtools.agents.prompt.is_text_content')
-    @patch('mimetypes.guess_type')
-    def test_text_file_returns_string(self, mock_guess_type, mock_is_text):
-        """Test that text files return decoded strings."""
-        mock_guess_type.return_value = ("text/plain", None)
-        mock_is_text.return_value = True
-        
-        test_file = self.temp_path / "test.txt"
-        test_content = "Hello, world!"
-        test_file.write_text(test_content, encoding="utf-8")
-        
-        result = file_to_binary_content(test_file)
-        
-        self.assertEqual(result, test_content)
-        mock_is_text.assert_called_once()
-
-    @patch('aixtools.agents.prompt.is_text_content')
-    @patch('mimetypes.guess_type')
-    def test_binary_file_returns_binary_content(self, mock_guess_type, mock_is_text):
-        """Test that binary files return BinaryContent."""
-        mock_guess_type.return_value = ("image/png", None)
-        mock_is_text.return_value = False
-        
-        test_file = self.temp_path / "test.png"
-        test_data = b'\x89PNG\r\n\x1a\n'
-        test_file.write_bytes(test_data)
-        
-        result = file_to_binary_content(test_file)
-        
-        self.assertIsInstance(result, BinaryContent)
-        if isinstance(result, BinaryContent):
-            self.assertEqual(result.data, test_data)
-            self.assertEqual(result.media_type, "image/png")
-
-    @patch('aixtools.agents.prompt.is_text_content')
-    @patch('mimetypes.guess_type')
-    def test_unknown_mime_type_defaults_to_octet_stream(self, mock_guess_type, mock_is_text):
-        """Test that unknown mime types default to application/octet-stream."""
-        mock_guess_type.return_value = (None, None)
-        mock_is_text.return_value = False
-        
-        test_file = self.temp_path / "test.unknown"
-        test_data = b'unknown data'
-        test_file.write_bytes(test_data)
-        
-        result = file_to_binary_content(test_file)
-        
-        self.assertIsInstance(result, BinaryContent)
-        if isinstance(result, BinaryContent):
-            self.assertEqual(result.media_type, "application/octet-stream")
-
-    @patch('aixtools.agents.prompt.is_text_content')
-    def test_explicit_mime_type(self, mock_is_text):
-        """Test with explicitly provided mime type."""
-        mock_is_text.return_value = False
-        
-        test_file = self.temp_path / "test.data"
-        test_data = b'test data'
-        test_file.write_bytes(test_data)
-        
-        result = file_to_binary_content(test_file, "application/custom")
-        
-        self.assertIsInstance(result, BinaryContent)
-        if isinstance(result, BinaryContent):
-            self.assertEqual(result.media_type, "application/custom")
-
-
-class TestBuildUserInput(unittest.TestCase):
-    """Test cases for build_user_input function."""
-
-    def setUp(self):
-        """Set up test fixtures."""
-        self.session_tuple = ("user123", "session456")
-        self.user_text = "Please analyze these files"
-
-    def test_no_file_paths_returns_text_only(self):
-        """Test that no file paths returns just the user text."""
-        result = build_user_input(self.session_tuple, self.user_text, [])
-        
-        self.assertEqual(result, self.user_text)
-
-    @patch('aixtools.agents.prompt.should_be_included_into_context')
-    @patch('aixtools.agents.prompt.file_to_binary_content')
-    @patch('aixtools.agents.prompt.container_to_host_path')
-    @patch('mimetypes.guess_type')
-    def test_single_file_not_included_in_context(
-        self, mock_guess_type, mock_container_to_host,
-        mock_file_to_binary, mock_should_include
-    ):
-        """Test with single file that should not be included in context."""
-        # Setup mocks
-        mock_guess_type.return_value = ("text/plain", None)
-        mock_should_include.return_value = False
-        
-        mock_host_path = Mock()
-        mock_host_path.stat.return_value.st_size = 1024
-        mock_container_to_host.return_value = mock_host_path
-        
-        mock_file_to_binary.return_value = "file content"
-        
-        file_paths = [Path("/workspace/test.txt")]
-        
-        result = build_user_input(self.session_tuple, self.user_text, file_paths)
-        
-        self.assertIsInstance(result, list)
-        self.assertEqual(len(result), 1)  # Only the prompt, no binary attachments
-        self.assertIsInstance(result[0], str)
-        prompt_text = str(result[0])
-        self.assertIn("Please analyze these files", prompt_text)
-        self.assertIn("Attachments:", prompt_text)
-        self.assertIn("test.txt", prompt_text)
-        self.assertIn("file_size=1024 bytes", prompt_text)
-
-    @patch('aixtools.agents.prompt.should_be_included_into_context')
-    @patch('aixtools.agents.prompt.file_to_binary_content')
-    @patch('aixtools.agents.prompt.container_to_host_path')
-    @patch('mimetypes.guess_type')
-    def test_single_file_included_in_context(
-        self, mock_guess_type, mock_container_to_host,
-        mock_file_to_binary, mock_should_include
-    ):
-        """Test with single file that should be included in context."""
-        # Setup mocks
-        mock_guess_type.return_value = ("image/png", None)
-        mock_should_include.return_value = True
-        
-        mock_host_path = Mock()
-        mock_host_path.stat.return_value.st_size = 2048
-        mock_container_to_host.return_value = mock_host_path
-        
-        mock_binary_content = BinaryContent(data=b"image data", media_type="image/png")
-        mock_file_to_binary.return_value = mock_binary_content
-        
-        file_paths = [Path("/workspace/image.png")]
-        
-        result = build_user_input(self.session_tuple, self.user_text, file_paths)
-        
-        self.assertIsInstance(result, list)
-        self.assertEqual(len(result), 2)  # Prompt + 1 binary attachment
-        self.assertIsInstance(result[0], str)
-        prompt_text = str(result[0])
-        self.assertIn("Please analyze these files", prompt_text)
-        self.assertIn("Attachments:", prompt_text)
-        self.assertIn("image.png", prompt_text)
-        self.assertIn("file_size=2048 bytes", prompt_text)
-        self.assertIn("provided to model context at index 0", prompt_text)
-        self.assertEqual(result[1], mock_binary_content)
-
-    @patch('aixtools.agents.prompt.should_be_included_into_context')
-    @patch('aixtools.agents.prompt.file_to_binary_content')
-    @patch('aixtools.agents.prompt.container_to_host_path')
-    @patch('mimetypes.guess_type')
-    def test_multiple_files_mixed_inclusion(
-        self, mock_guess_type, mock_container_to_host,
-        mock_file_to_binary, mock_should_include
-    ):
-        """Test with multiple files, some included in context, some not."""
-        # Setup mocks
-        mock_guess_type.side_effect = [("text/plain", None), ("image/png", None)]
-        mock_should_include.side_effect = [False, True]  # First file not included, second included
-        
-        mock_host_path1 = Mock()
-        mock_host_path1.stat.return_value.st_size = 1024
-        mock_host_path2 = Mock()
-        mock_host_path2.stat.return_value.st_size = 2048
-        mock_container_to_host.side_effect = [mock_host_path1, mock_host_path2]
-        
-        mock_text_content = "text content"
-        mock_binary_content = BinaryContent(data=b"image data", media_type="image/png")
-        mock_file_to_binary.side_effect = [mock_text_content, mock_binary_content]
-        
-        file_paths = [Path("/workspace/text.txt"), Path("/workspace/image.png")]
-        
-        result = build_user_input(self.session_tuple, self.user_text, file_paths)
-        
-        self.assertIsInstance(result, list)
-        self.assertEqual(len(result), 2)  # Prompt + 1 binary attachment
-        self.assertIsInstance(result[0], str)
-        prompt_text = str(result[0])
-        self.assertIn("Please analyze these files", prompt_text)
-        self.assertIn("Attachments:", prompt_text)
-        self.assertIn("text.txt", prompt_text)
-        self.assertIn("image.png", prompt_text)
-        self.assertIn("file_size=1024 bytes", prompt_text)
-        self.assertIn("file_size=2048 bytes", prompt_text)
-        self.assertIn("provided to model context at index 0", prompt_text)
-        self.assertEqual(result[1], mock_binary_content)
-
-    def test_non_workspace_path_raises_error(self):
-        """Test that non-workspace paths raise ValueError."""
-        file_paths = [Path("/invalid/path.txt")]
-        
-        with self.assertRaises(ValueError) as context:
-            build_user_input(self.session_tuple, self.user_text, file_paths)
-        
-        self.assertIn(
-            "Container path must be a subdir of '/workspace', got '/invalid/path.txt' instead",
-            str(context.exception),
-        )
-
-    @patch('aixtools.agents.prompt.should_be_included_into_context')
-    @patch('aixtools.agents.prompt.file_to_binary_content')
-    @patch('aixtools.agents.prompt.container_to_host_path')
-    @patch('mimetypes.guess_type')
-    def test_unknown_mime_type_defaults(
-        self, mock_guess_type, mock_container_to_host,
-        mock_file_to_binary, mock_should_include
-    ):
-        """Test that unknown mime types default to application/octet-stream."""
-        # Setup mocks
-        mock_guess_type.return_value = (None, None)  # Unknown mime type
-        mock_should_include.return_value = False
-        
-        mock_host_path = Mock()
-        mock_host_path.stat.return_value.st_size = 1024
-        mock_container_to_host.return_value = mock_host_path
-        
-        mock_file_to_binary.return_value = "file content"
-        
-        file_paths = [Path("/workspace/unknown.dat")]
-        
-        build_user_input(self.session_tuple, self.user_text, file_paths)
-        
-        # Verify that file_to_binary_content was called with the default mime type
-        mock_file_to_binary.assert_called_once_with(mock_host_path, "application/octet-stream")
-
-
-if __name__ == '__main__':
-    unittest.main()