kolega-code 0.1.0__py3-none-any.whl

kolega_code/agent/tests/test_coder_attachments.py

@@ -0,0 +1,330 @@
+"""Tests for coder agent image attachment handling."""
+
+import base64
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+
+from kolega_code.agent.baseagent import BaseAgent
+from kolega_code.agent.coder import CoderAgent
+from kolega_code.config import AgentConfig, ModelConfig, ModelProvider, RateLimitConfig
+from kolega_code.llm.models import ImageBlock, Message, TextBlock
+from kolega_code.llm.providers.models import TokenCount
+from kolega_code.agent.prompt_provider import AgentMode
+
+
+def _deepseek_config() -> AgentConfig:
+    model_config = ModelConfig(
+        provider=ModelProvider.DEEPSEEK,
+        model="deepseek-v4-pro",
+        rate_limits=RateLimitConfig(),
+    )
+    return AgentConfig(
+        deepseek_api_key="test-key",
+        long_context_config=model_config,
+        fast_config=model_config,
+        edit_model_config=model_config,
+        thinking_config=model_config,
+    )
+
+
+def _image_attachment() -> dict:
+    return {
+        "type": "image",
+        "media_type": "image/png",
+        "data": base64.b64encode(b"fake-image-data").decode("utf-8"),
+        "filename": "test-image.png",
+    }
+
+
+class _EmptyStream:
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        return False
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        raise StopAsyncIteration
+
+    async def get_final_message(self):
+        return Message("assistant", [TextBlock("done")], stop_reason="end_turn")
+
+
+def test_deepseek_image_attachment_is_rejected_by_provider_check():
+    agent = object.__new__(BaseAgent)
+    agent.config = SimpleNamespace(
+        long_context_config=SimpleNamespace(provider=ModelProvider.DEEPSEEK.value),
+    )
+
+    assert (
+        agent._unsupported_attachment_message([_image_attachment()])
+        == BaseAgent.deepseek_image_unsupported_message
+    )
+
+
+def test_deepseek_attachment_check_allows_non_images_and_other_providers():
+    agent = object.__new__(BaseAgent)
+    agent.config = SimpleNamespace(
+        long_context_config=SimpleNamespace(provider=ModelProvider.DEEPSEEK),
+    )
+    assert agent._unsupported_attachment_message(None) is None
+    assert agent._unsupported_attachment_message([{"type": "document", "data": "abc"}]) is None
+
+    agent.config.long_context_config.provider = ModelProvider.ANTHROPIC
+    assert agent._unsupported_attachment_message([_image_attachment()]) is None
+
+
+@pytest.mark.asyncio
+async def test_coder_agent_rejects_deepseek_image_without_llm_call(tmp_path):
+    connection_manager = Mock()
+    connection_manager.broadcast_event = AsyncMock()
+    agent = CoderAgent(
+        project_path=tmp_path,
+        workspace_id="workspace-123",
+        thread_id="thread-123",
+        connection_manager=connection_manager,
+        config=_deepseek_config(),
+        agent_mode=AgentMode.CLI,
+    )
+    agent.llm = Mock()
+
+    chunks = [
+        chunk
+        async for chunk in agent.process_message_stream("What is in this image?", [_image_attachment()])
+    ]
+
+    assert len(chunks) == 1
+    assert chunks[0]["type"] == "response"
+    assert chunks[0]["content"] == BaseAgent.deepseek_image_unsupported_message
+    assert chunks[0]["complete"] is True
+    assert agent.history == []
+    agent.llm.stream.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_coder_agent_does_not_print_context_token_counts(tmp_path, capsys):
+    connection_manager = Mock()
+    connection_manager.broadcast_event = AsyncMock()
+    agent = CoderAgent(
+        project_path=tmp_path,
+        workspace_id="workspace-123",
+        thread_id="thread-123",
+        connection_manager=connection_manager,
+        config=_deepseek_config(),
+        agent_mode=AgentMode.CLI,
+    )
+    agent.count_current_context = AsyncMock(return_value=TokenCount(input_tokens=42))
+    agent.llm = Mock()
+    agent.llm.stream = AsyncMock(return_value=_EmptyStream())
+
+    chunks = [chunk async for chunk in agent.process_message_stream("hello")]
+
+    assert chunks[-1]["complete"] is True
+    assert capsys.readouterr().out == ""
+
+
+class TestCoderAgentAttachments:
+    """Test suite for verifying image attachment handling logic."""
+
+    def test_single_image_attachment_processing(self):
+        """Test that a single image attachment is correctly processed into an ImageBlock."""
+        test_image_data = base64.b64encode(b"fake-image-data").decode("utf-8")
+        attachment = {
+            "type": "image",
+            "media_type": "image/png",
+            "data": test_image_data,
+            "filename": "test-image.png",
+        }
+
+        # Process the attachment as the coder agent would
+        image_block = ImageBlock(
+            image_type="base64", media_type=attachment.get("media_type", "image/png"), data=attachment["data"]
+        )
+
+        assert image_block.image_type == "base64"
+        assert image_block.media_type == "image/png"
+        assert image_block.data == test_image_data
+
+    def test_multiple_image_attachments_processing(self):
+        """Test that multiple image attachments are correctly processed."""
+        attachments = [
+            {
+                "type": "image",
+                "media_type": "image/png",
+                "data": base64.b64encode(b"image1").decode("utf-8"),
+                "filename": "image1.png",
+            },
+            {
+                "type": "image",
+                "media_type": "image/jpeg",
+                "data": base64.b64encode(b"image2").decode("utf-8"),
+                "filename": "image2.jpg",
+            },
+        ]
+
+        # Process attachments as the coder agent would
+        image_blocks = []
+        for attachment in attachments:
+            if attachment.get("type") == "image":
+                image_block = ImageBlock(
+                    image_type="base64", media_type=attachment.get("media_type", "image/png"), data=attachment["data"]
+                )
+                image_blocks.append(image_block)
+
+        assert len(image_blocks) == 2
+        assert image_blocks[0].media_type == "image/png"
+        assert image_blocks[1].media_type == "image/jpeg"
+
+    def test_non_image_attachments_filtered(self):
+        """Test that non-image attachments are filtered out."""
+        attachments = [
+            {
+                "type": "document",
+                "media_type": "application/pdf",
+                "data": base64.b64encode(b"pdf-data").decode("utf-8"),
+                "filename": "document.pdf",
+            },
+            {
+                "type": "image",
+                "media_type": "image/png",
+                "data": base64.b64encode(b"image-data").decode("utf-8"),
+                "filename": "image.png",
+            },
+        ]
+
+        # Process attachments, filtering non-images
+        image_blocks = []
+        for attachment in attachments:
+            if attachment.get("type") == "image":
+                image_block = ImageBlock(
+                    image_type="base64", media_type=attachment.get("media_type", "image/png"), data=attachment["data"]
+                )
+                image_blocks.append(image_block)
+
+        assert len(image_blocks) == 1
+        assert image_blocks[0].media_type == "image/png"
+
+    def test_empty_attachments_handling(self):
+        """Test that empty or None attachments are handled gracefully."""
+        # Test with None
+        image_blocks = []
+        attachments = None
+        if attachments:
+            for attachment in attachments:
+                if attachment.get("type") == "image":
+                    image_block = ImageBlock(
+                        image_type="base64",
+                        media_type=attachment.get("media_type", "image/png"),
+                        data=attachment["data"],
+                    )
+                    image_blocks.append(image_block)
+
+        assert len(image_blocks) == 0
+
+        # Test with empty list
+        image_blocks = []
+        attachments = []
+        for attachment in attachments:
+            if attachment.get("type") == "image":
+                image_block = ImageBlock(
+                    image_type="base64", media_type=attachment.get("media_type", "image/png"), data=attachment["data"]
+                )
+                image_blocks.append(image_block)
+
+        assert len(image_blocks) == 0
+
+    def test_message_content_with_attachments(self):
+        """Test that message content is correctly structured with text and image blocks."""
+        test_message = "What is in this image?"
+        test_image_data = base64.b64encode(b"fake-image-data").decode("utf-8")
+        attachments = [
+            {
+                "type": "image",
+                "media_type": "image/png",
+                "data": test_image_data,
+                "filename": "test-image.png",
+            }
+        ]
+
+        # Build content blocks as the coder agent would
+        content_blocks = [TextBlock(text=test_message)]
+
+        if attachments:
+            for attachment in attachments:
+                if attachment.get("type") == "image":
+                    image_block = ImageBlock(
+                        image_type="base64",
+                        media_type=attachment.get("media_type", "image/png"),
+                        data=attachment["data"],
+                    )
+                    content_blocks.append(image_block)
+
+        assert len(content_blocks) == 2
+        assert isinstance(content_blocks[0], TextBlock)
+        assert content_blocks[0].text == test_message
+        assert isinstance(content_blocks[1], ImageBlock)
+        assert content_blocks[1].data == test_image_data
+
+
+@pytest.mark.asyncio
+async def test_coder_agent_process_message_imports():
+    """Test that the coder agent has the necessary imports for image handling."""
+    # This test verifies the imports are correct
+    from kolega_code.agent.coder import CoderAgent
+    from kolega_code.llm.models import ImageBlock, TextBlock
+
+    # Just verify the imports work
+    assert ImageBlock is not None
+    assert TextBlock is not None
+
+
+def test_attachment_blocks_mixes_images_and_files():
+    agent = object.__new__(BaseAgent)
+    blocks = agent._attachment_blocks(
+        [
+            _image_attachment(),
+            {"type": "file", "path": "src/app.py", "content": "print('hi')"},
+            {"type": "unknown", "data": "ignored"},
+        ]
+    )
+
+    assert len(blocks) == 2
+    assert isinstance(blocks[0], ImageBlock)
+    assert isinstance(blocks[1], TextBlock)
+    assert blocks[1].text == '<attached-file path="src/app.py">\nprint(\'hi\')\n</attached-file>'
+
+
+def test_attachment_blocks_handles_none():
+    agent = object.__new__(BaseAgent)
+    assert agent._attachment_blocks(None) == []
+
+
+@pytest.mark.asyncio
+async def test_coder_agent_file_attachment_added_to_history(tmp_path):
+    connection_manager = Mock()
+    connection_manager.broadcast_event = AsyncMock()
+    agent = CoderAgent(
+        project_path=tmp_path,
+        workspace_id="workspace-123",
+        thread_id="thread-123",
+        connection_manager=connection_manager,
+        config=_deepseek_config(),
+        agent_mode=AgentMode.CLI,
+    )
+    agent.count_current_context = AsyncMock(return_value=TokenCount(input_tokens=42))
+    agent.llm = Mock()
+    agent.llm.stream = AsyncMock(return_value=_EmptyStream())
+
+    attachments = [{"type": "file", "path": "notes.md", "content": "remember the milk"}]
+    chunks = [chunk async for chunk in agent.process_message_stream("see the notes", attachments)]
+
+    assert chunks[-1]["complete"] is True
+    user_message = agent.history[0]
+    texts = [block.text for block in user_message.content if isinstance(block, TextBlock)]
+    assert texts[0] == "see the notes"
+    assert texts[1] == '<attached-file path="notes.md">\nremember the milk\n</attached-file>'

kolega_code/agent/tests/test_coder_prompt_extensions.py

@@ -0,0 +1,61 @@
+from unittest.mock import AsyncMock, Mock
+
+from kolega_code.agent.coder import CoderAgent
+from kolega_code.config import AgentConfig, ModelConfig, ModelProvider, RateLimitConfig
+from kolega_code.agent.prompt_provider import AgentMode, AgentType, PromptExtension, PromptProvider
+
+
+def test_coder_agent_includes_matching_prompt_extensions(tmp_path):
+    config = AgentConfig(
+        anthropic_api_key="test-key",
+        openai_api_key="test-key",
+        long_context_config=ModelConfig(
+            provider=ModelProvider.ANTHROPIC,
+            model="claude-sonnet-4-20250514",
+            rate_limits=RateLimitConfig(),
+        ),
+        fast_config=ModelConfig(
+            provider=ModelProvider.ANTHROPIC,
+            model="claude-3-haiku-20240307",
+            rate_limits=RateLimitConfig(),
+        ),
+        thinking_config=ModelConfig(
+            provider=ModelProvider.ANTHROPIC,
+            model="claude-3-7-sonnet-20250219",
+            rate_limits=RateLimitConfig(),
+            thinking_tokens=1024,
+        ),
+    )
+
+    connection_manager = Mock()
+    connection_manager.broadcast_event = AsyncMock()
+    template_dir = tmp_path / "prompt_templates"
+    agents_dir = template_dir / "agents"
+    agents_dir.mkdir(parents=True)
+    (agents_dir / "coder_code_mode.j2").write_text(
+        "{% for extension in prompt_extensions %}{{ extension.title }}\n{{ extension.markdown }}{% endfor %}",
+        encoding="utf-8",
+    )
+
+    agent = CoderAgent(
+        project_path=tmp_path,
+        workspace_id="workspace-123",
+        thread_id="thread-123",
+        connection_manager=connection_manager,
+        config=config,
+        agent_mode=AgentMode.CODE,
+        prompt_provider=PromptProvider(template_dirs=[template_dir]),
+        prompt_extensions=[
+            PromptExtension(
+                id="host-context",
+                title="Host Context",
+                markdown="Injected host-specific context.",
+                agent_types=[AgentType.CODER],
+                modes=[AgentMode.CODE],
+            )
+        ],
+    )
+
+    prompt = agent.system_prompt.content[0].text
+    assert "Host Context" in prompt
+    assert "Injected host-specific context." in prompt

kolega_code/agent/tests/test_commands.py

@@ -0,0 +1,179 @@
+from unittest.mock import AsyncMock, MagicMock, create_autospec
+
+import pytest
+
+from kolega_code.llm.models import Message, MessageHistory
+from ..utils.commands import CommandProcessor
+
+
+class MockAgent:
+    """Mock agent class for testing CommandProcessor"""
+
+    def __init__(self):
+        self.history = MessageHistory()
+        self.command_processor = CommandProcessor(self)
+        self.compress_history = AsyncMock()
+        self.count_current_context = AsyncMock()
+
+
+@pytest.fixture
+def mock_agent():
+    """Fixture to create a mock agent for testing"""
+    return MockAgent()
+
+
+@pytest.fixture
+def command_processor(mock_agent):
+    """Fixture to create a CommandProcessor instance with a mock agent"""
+    return mock_agent.command_processor
+
+
+@pytest.fixture
+def mock_message():
+    """Fixture to create a mock Message object"""
+    message = create_autospec(Message)
+    message.role = "user"
+    message.content = "Test message"
+    return message
+
+
+@pytest.mark.asyncio
+async def test_handle_help(command_processor):
+    """Test the /help command handler"""
+    help_text = await command_processor._handle_help()
+
+    # Verify help text contains all available commands
+    assert "/help" in help_text
+    assert "/compress" in help_text
+    assert "/clear" in help_text
+    assert "/reset" in help_text
+    assert "/context" in help_text
+
+    # Verify help text formatting
+    assert help_text.startswith("# Available Commands")
+    assert all(line.startswith("- `/") for line in help_text.split("\n")[2:])
+
+
+@pytest.mark.asyncio
+async def test_handle_compress(command_processor, mock_agent):
+    """Test the /compress command handler"""
+    result = await command_processor._handle_compress()
+
+    # Verify compress was called and returned expected message
+    mock_agent.compress_history.assert_called_once()
+    assert result == "Message history compressed."
+
+
+@pytest.mark.asyncio
+async def test_handle_clear(command_processor, mock_agent, mock_message):
+    """Test the /clear command handler"""
+    # Add some messages to history first
+    mock_agent.history.append(mock_message)
+    mock_agent.history.append(mock_message)
+    assert len(mock_agent.history) == 2
+
+    result = await command_processor._handle_clear()
+
+    # Verify history was cleared and returned expected message
+    assert len(mock_agent.history) == 0
+    assert result == "Message history cleared."
+
+
+@pytest.mark.asyncio
+async def test_handle_reset_alias(command_processor, mock_agent, mock_message):
+    """Test the /reset command alias"""
+    mock_agent.history.append(mock_message)
+    mock_agent.history.append(mock_message)
+    assert len(mock_agent.history) == 2
+
+    result = await command_processor.commands["/reset"]()
+
+    assert len(mock_agent.history) == 0
+    assert result == "Message history cleared."
+
+
+@pytest.mark.asyncio
+async def test_handle_context_empty_history(command_processor):
+    """Test the /context command handler with empty history"""
+    result = await command_processor._handle_context()
+    assert result == "Current context token count: 0"
+
+
+@pytest.mark.asyncio
+async def test_handle_context_with_history(command_processor, mock_agent, mock_message):
+    """Test the /context command handler with non-empty history"""
+    # Add mock message to history
+    mock_agent.history.append(mock_message)
+
+    # Mock token count response
+    mock_token_count = MagicMock()
+    mock_token_count.input_tokens = 100
+    mock_agent.count_current_context.return_value = mock_token_count
+
+    result = await command_processor._handle_context()
+
+    # Verify token count was called and returned expected message
+    mock_agent.count_current_context.assert_called_once()
+    assert result == "Current context token count: 100"
+
+
+@pytest.mark.asyncio
+async def test_process_commands_decorator():
+    """Test the process_commands decorator functionality"""
+
+    @CommandProcessor.process_commands
+    class TestAgent:
+        def __init__(self):
+            self.history = MessageHistory()
+            self.command_processor = CommandProcessor(self)
+
+        async def process_message_stream(self, message, attachments=None):
+            yield {"type": "response", "content": "Normal processing", "complete": True}
+
+    agent = TestAgent()
+
+    # Test command processing
+    responses = []
+    async for response in agent.process_message_stream("/help"):
+        responses.append(response)
+
+    assert len(responses) == 1
+    assert responses[0]["type"] == "response"
+    assert "Available Commands" in responses[0]["content"]
+    assert responses[0]["complete"] is True
+
+    # Test normal message processing
+    responses = []
+    async for response in agent.process_message_stream("normal message"):
+        responses.append(response)
+
+    assert len(responses) == 1
+    assert responses[0]["type"] == "response"
+    assert responses[0]["content"] == "Normal processing"
+    assert responses[0]["complete"] is True
+
+
+@pytest.mark.asyncio
+async def test_process_commands_invalid_command():
+    """Test processing an invalid command"""
+
+    @CommandProcessor.process_commands
+    class TestAgent:
+        def __init__(self):
+            self.history = MessageHistory()
+            self.command_processor = CommandProcessor(self)
+
+        async def process_message_stream(self, message, attachments=None):
+            yield {"type": "response", "content": "Normal processing", "complete": True}
+
+    agent = TestAgent()
+
+    # Test invalid command
+    responses = []
+    async for response in agent.process_message_stream("/invalid"):
+        responses.append(response)
+
+    assert len(responses) == 1
+    assert responses[0]["type"] == "response"
+    assert responses[0]["content"] == "Normal processing"
+    assert responses[0]["complete"] is True