""" Unit tests for Output Chunking and Streaming System. Tests the ChunkedResult, ProgressiveProcessor, and ContentSummarizer to ensure proper chunked output handling and streaming capabilities. """ import pytest import time import asyncio from unittest.mock import Mock, AsyncMock, MagicMock from typing import List, Dict, Any import sys from pathlib import Path sys.path.insert(0, str(Path(__file__).parent.parent.parent / "src")) from pipeline.chunking import ( ChunkedResult, ProgressiveProcessor, ContentSummarizer ) class TestChunkedResult: """Test ChunkedResult functionality.""" def test_init(self): """Test ChunkedResult initialization.""" result = ChunkedResult( chunk_id="chunk_1", task_id="task_123", agent_id="agent_456", content_chunk="This is a test chunk.", chunk_index=0, is_final=False, timestamp=time.time(), continuation_token="token_abc" ) assert result.chunk_id == "chunk_1" assert result.task_id == "task_123" assert result.agent_id == "agent_456" assert result.content_chunk == "This is a test chunk." assert result.chunk_index == 0 assert not result.is_final assert result.continuation_token == "token_abc" assert isinstance(result.timestamp, float) assert isinstance(result.metadata, dict) def test_init_with_defaults(self): """Test ChunkedResult with default values.""" result = ChunkedResult( chunk_id="chunk_1", task_id="task_123", agent_id="agent_456", content_chunk="Test content", chunk_index=0, is_final=True, timestamp=time.time() ) assert result.metadata == {} assert result.continuation_token is None class TestProgressiveProcessor: """Test ProgressiveProcessor functionality.""" def test_init(self): """Test ProgressiveProcessor initialization.""" content = "This is a test content for chunking. " * 10 # 400+ chars processor = ProgressiveProcessor( task_id="task_123", agent_id="agent_456", full_content=content, chunk_size=100 ) assert processor.task_id == "task_123" assert processor.agent_id == "agent_456" assert processor.full_content == content assert processor.chunk_size == 100 assert processor._current_chunk_index == 0 assert processor.total_chunks > 1 # Should need multiple chunks assert len(processor._continuation_tokens) == processor.total_chunks def test_post_init_calculations(self): """Test post-initialization calculations.""" content = "A" * 250 # 250 characters processor = ProgressiveProcessor( task_id="task_1", agent_id="agent_1", full_content=content, chunk_size=100 ) # Should need 3 chunks: 100 + 100 + 50 assert processor.total_chunks == 3 assert len(processor._continuation_tokens) == 3 # All tokens should be unique tokens = list(processor._continuation_tokens.values()) assert len(tokens) == len(set(tokens)) def test_get_first_chunk(self): """Test getting the first chunk.""" content = "This is chunk one. This is chunk two. This is chunk three." processor = ProgressiveProcessor( task_id="task_1", agent_id="agent_1", full_content=content, chunk_size=20 ) first_chunk = processor.get_next_chunk() assert first_chunk is not None assert first_chunk.chunk_index == 0 assert first_chunk.content_chunk == content[:20] assert not first_chunk.is_final # Should have more chunks assert first_chunk.continuation_token is not None assert first_chunk.task_id == "task_1" assert first_chunk.agent_id == "agent_1" def test_get_chunk_sequence(self): """Test getting a sequence of chunks.""" content = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" # 26 characters processor = ProgressiveProcessor( task_id="task_1", agent_id="agent_1", full_content=content, chunk_size=10 ) chunks = [] current_token = None # Get all chunks while True: chunk = processor.get_next_chunk(current_token) if chunk is None: break chunks.append(chunk) current_token = chunk.continuation_token if chunk.is_final: break assert len(chunks) == 3 # 10 + 10 + 6 characters assert chunks[0].content_chunk == "ABCDEFGHIJ" assert chunks[1].content_chunk == "KLMNOPQRST" assert chunks[2].content_chunk == "UVWXYZ" # Only last chunk should be final assert not chunks[0].is_final assert not chunks[1].is_final assert chunks[2].is_final # Final chunk should have no continuation token assert chunks[2].continuation_token is None def test_get_chunk_with_invalid_token(self): """Test getting chunk with invalid continuation token.""" content = "Test content for invalid token testing." processor = ProgressiveProcessor( task_id="task_1", agent_id="agent_1", full_content=content, chunk_size=15 ) # Try with invalid token chunk = processor.get_next_chunk("invalid_token_12345") assert chunk is None def test_get_chunk_by_index(self): """Test getting chunk by specific index.""" content = "Index test: " + "ABCDEFGHIJ" * 5 # 62 characters total processor = ProgressiveProcessor( task_id="task_1", agent_id="agent_1", full_content=content, chunk_size=20 ) # Get chunk at index 1 chunk = processor.get_chunk_by_index(1) assert chunk is not None assert chunk.chunk_index == 1 assert chunk.content_chunk == content[20:40] assert chunk.task_id == "task_1" # Get chunk at index 0 first_chunk = processor.get_chunk_by_index(0) assert first_chunk is not None assert first_chunk.chunk_index == 0 assert first_chunk.content_chunk == content[:20] # Get final chunk last_index = processor.total_chunks - 1 last_chunk = processor.get_chunk_by_index(last_index) assert last_chunk is not None assert last_chunk.is_final assert last_chunk.continuation_token is None def test_get_chunk_by_invalid_index(self): """Test getting chunk with invalid index.""" content = "Short content" processor = ProgressiveProcessor( task_id="task_1", agent_id="agent_1", full_content=content, chunk_size=50 ) # Index out of bounds assert processor.get_chunk_by_index(-1) is None assert processor.get_chunk_by_index(10) is None def test_single_chunk_content(self): """Test content that fits in a single chunk.""" content = "Short" processor = ProgressiveProcessor( task_id="task_1", agent_id="agent_1", full_content=content, chunk_size=100 ) assert processor.total_chunks == 1 chunk = processor.get_next_chunk() assert chunk is not None assert chunk.chunk_index == 0 assert chunk.content_chunk == content assert chunk.is_final assert chunk.continuation_token is None def test_empty_content(self): """Test processing empty content.""" processor = ProgressiveProcessor( task_id="task_1", agent_id="agent_1", full_content="", chunk_size=100 ) assert processor.total_chunks == 1 # Empty content still creates one chunk chunk = processor.get_next_chunk() assert chunk is not None assert chunk.content_chunk == "" assert chunk.is_final def test_chunk_metadata_consistency(self): """Test that chunk metadata is consistent across requests.""" content = "Metadata consistency test content for chunking." processor = ProgressiveProcessor( task_id="task_123", agent_id="agent_456", full_content=content, chunk_size=15 ) # Get same chunk multiple times chunk1 = processor.get_chunk_by_index(0) chunk2 = processor.get_chunk_by_index(0) assert chunk1.task_id == chunk2.task_id assert chunk1.agent_id == chunk2.agent_id assert chunk1.chunk_index == chunk2.chunk_index assert chunk1.content_chunk == chunk2.content_chunk assert chunk1.is_final == chunk2.is_final # Note: chunk_id and timestamp will be different as they're generated fresh def test_large_content_performance(self): """Test performance with large content.""" # Create large content (10KB) content = "Large content test. " * 500 # ~10,000 characters start_time = time.time() processor = ProgressiveProcessor( task_id="perf_test", agent_id="agent_1", full_content=content, chunk_size=1000 ) # Should handle large content quickly assert time.time() - start_time < 1.0 # Less than 1 second # Should create reasonable number of chunks assert processor.total_chunks <= 15 # 10KB / 1KB + buffer # First chunk should work quickly start_time = time.time() first_chunk = processor.get_next_chunk() assert time.time() - start_time < 0.1 # Very fast assert first_chunk is not None assert len(first_chunk.content_chunk) == 1000 class TestContentSummarizer: """Test ContentSummarizer functionality.""" def setup_method(self): """Set up test fixtures.""" self.mock_llm_client = Mock() self.summarizer = ContentSummarizer(self.mock_llm_client) def test_init(self): """Test ContentSummarizer initialization.""" assert self.summarizer.llm_client == self.mock_llm_client @pytest.mark.asyncio async def test_summarize_content_short_content(self): """Test summarization of short content that doesn't need summarization.""" short_content = "This is a short text." result = await self.summarizer.summarize_content( content=short_content, target_tokens=100, agent_id="agent_1", task_id="task_1" ) # Short content should be returned as-is assert result == short_content self.mock_llm_client.complete_async.assert_not_called() @pytest.mark.asyncio async def test_summarize_content_empty(self): """Test summarization of empty content.""" result = await self.summarizer.summarize_content( content="", target_tokens=100, agent_id="agent_1", task_id="task_1" ) assert result == "" self.mock_llm_client.complete_async.assert_not_called() @pytest.mark.asyncio async def test_summarize_content_needs_summarization(self): """Test summarization of long content.""" # Create long content that needs summarization (>100 words) long_content = "This is a very long piece of content that needs to be summarized. " * 50 # Mock LLM response mock_response = Mock() mock_response.content = "This is a concise summary of the long content." self.mock_llm_client.complete_async = AsyncMock(return_value=mock_response) result = await self.summarizer.summarize_content( content=long_content, target_tokens=50, agent_id="agent_1", task_id="task_1" ) assert result == "This is a concise summary of the long content." # Check LLM was called with correct parameters self.mock_llm_client.complete_async.assert_called_once() call_args = self.mock_llm_client.complete_async.call_args assert call_args.kwargs["agent_id"] == "agent_1" assert call_args.kwargs["max_tokens"] == 50 assert call_args.kwargs["temperature"] == 0.3 assert "Summarize the following content" in call_args.kwargs["user_prompt"] assert long_content in call_args.kwargs["user_prompt"] @pytest.mark.asyncio async def test_summarize_content_llm_error(self): """Test summarization fallback when LLM fails.""" long_content = "Content that needs summarization. " * 100 # Mock LLM to raise an exception self.mock_llm_client.complete_async = AsyncMock(side_effect=Exception("LLM Error")) result = await self.summarizer.summarize_content( content=long_content, target_tokens=50, agent_id="agent_1", task_id="task_1" ) # Should fallback to simple truncation assert result.endswith("...") assert len(result) <= 200 + 3 # 50 tokens * 4 chars + "..." def test_simple_truncate(self): """Test simple truncation fallback.""" content = "This is a test content that will be truncated." result = self.summarizer._simple_truncate(content, target_tokens=5) # Should be truncated to ~20 characters (5 tokens * 4 chars) assert len(result) <= 23 # 20 + "..." assert result.endswith("...") def test_simple_truncate_short_content(self): """Test simple truncation with content shorter than limit.""" short_content = "Short" result = self.summarizer._simple_truncate(short_content, target_tokens=10) # Should return content as-is assert result == short_content @pytest.mark.asyncio async def test_system_prompt_includes_target_tokens(self): """Test that system prompt includes target token limit.""" long_content = "Long content for testing system prompt. " * 20 target_tokens = 75 mock_response = Mock() mock_response.content = "Summary" self.mock_llm_client.complete_async = AsyncMock(return_value=mock_response) await self.summarizer.summarize_content( content=long_content, target_tokens=target_tokens, agent_id="agent_1", task_id="task_1" ) call_args = self.mock_llm_client.complete_async.call_args system_prompt = call_args.kwargs["system_prompt"] assert str(target_tokens) in system_prompt assert "Content Summarizer" in system_prompt assert "essential information" in system_prompt class TestIntegrationScenarios: """Test realistic integration scenarios.""" def test_chunked_blog_post_simulation(self): """Test chunking a realistic blog post.""" # Simulate a blog post blog_content = """ # The Future of Artificial Intelligence Artificial intelligence is rapidly transforming our world. In this comprehensive analysis, we explore the key trends and implications for the future. ## Current State of AI Today's AI systems demonstrate remarkable capabilities in various domains including natural language processing, computer vision, and decision-making. These systems are being deployed across industries from healthcare to finance. ## Emerging Trends Several key trends are shaping the future of AI: 1. Increased model sophistication 2. Better human-AI collaboration 3. Improved ethical frameworks 4. Enhanced accessibility ## Challenges Ahead Despite progress, significant challenges remain including bias in algorithms, privacy concerns, and the need for better interpretability. ## Conclusion The future of AI holds great promise, but requires careful consideration of ethical implications and societal impact. """ processor = ProgressiveProcessor( task_id="blog_post", agent_id="blog_writer", full_content=blog_content.strip(), chunk_size=300 ) chunks = [] current_token = None while True: chunk = processor.get_next_chunk(current_token) if chunk is None: break chunks.append(chunk) current_token = chunk.continuation_token if chunk.is_final: break # Should create multiple chunks assert len(chunks) >= 3 # Reconstruct content from chunks reconstructed = "".join(chunk.content_chunk for chunk in chunks) assert reconstructed == blog_content.strip() # Check chunk properties for i, chunk in enumerate(chunks): assert chunk.chunk_index == i assert chunk.task_id == "blog_post" assert chunk.agent_id == "blog_writer" # Only last chunk should be final for chunk in chunks[:-1]: assert not chunk.is_final assert chunk.continuation_token is not None assert chunks[-1].is_final assert chunks[-1].continuation_token is None @pytest.mark.asyncio async def test_summarization_fallback_scenario(self): """Test realistic summarization scenario with fallback.""" # Long research content that exceeds token limit research_content = """ Research Analysis: Climate Change Impact on Agricultural Systems Executive Summary: This comprehensive study examines the multifaceted impacts of climate change on global agricultural systems, analyzing temperature variations, precipitation patterns, and extreme weather events across multiple geographic regions. Methodology: We employed a mixed-methods approach combining quantitative climate data analysis with qualitative assessments from agricultural stakeholders across 15 countries. Key Findings: 1. Temperature increases of 2-3°C significantly reduce crop yields in tropical regions 2. Changing precipitation patterns affect irrigation-dependent systems most severely 3. Extreme weather events cause both immediate and long-term agricultural disruption 4. Adaptation strategies show varying effectiveness across different crop types Regional Analysis: Sub-Saharan Africa shows greatest vulnerability while Northern European regions may experience some agricultural benefits from moderate warming. Recommendations: Immediate implementation of climate-resilient farming practices, investment in drought-resistant crop varieties, and improved early warning systems. """ * 3 # Triple the content to ensure it needs summarization mock_llm_client = Mock() mock_response = Mock() mock_response.content = "Climate change significantly impacts global agriculture through temperature increases, changing precipitation, and extreme weather. Key recommendations include climate-resilient farming and drought-resistant crops." mock_llm_client.complete_async = AsyncMock(return_value=mock_response) summarizer = ContentSummarizer(mock_llm_client) summary = await summarizer.summarize_content( content=research_content, target_tokens=100, agent_id="research_agent", task_id="climate_research" ) # Should get summarized version assert summary != research_content assert "climate change" in summary.lower() assert "agriculture" in summary.lower() assert len(summary) < len(research_content) # Verify LLM was called with research content mock_llm_client.complete_async.assert_called_once() call_args = mock_llm_client.complete_async.call_args assert "climate change" in call_args.kwargs["user_prompt"].lower() def test_progressive_synthesis_workflow(self): """Test a progressive synthesis workflow with multiple agents.""" # Simulate multiple agent contributions contributions = [ "Research findings on quantum computing fundamentals and current capabilities.", "Analysis of market trends and commercial applications in the quantum computing sector.", "Technical review of hardware approaches: superconducting, trapped ion, and photonic systems.", "Risk assessment and timeline projections for quantum computing milestones." ] # Each contribution gets processed into chunks processors = [] all_chunks = [] for i, contribution in enumerate(contributions): processor = ProgressiveProcessor( task_id=f"synthesis_task", agent_id=f"agent_{i+1}", full_content=contribution, chunk_size=50 ) processors.append(processor) # Get first chunk from each processor first_chunk = processor.get_next_chunk() if first_chunk: all_chunks.append(first_chunk) # Should have chunks from multiple agents assert len(all_chunks) == 4 # Each chunk should have different agent_id but same task_id agent_ids = {chunk.agent_id for chunk in all_chunks} task_ids = {chunk.task_id for chunk in all_chunks} assert len(agent_ids) == 4 # Four different agents assert len(task_ids) == 1 # Same task assert "synthesis_task" in task_ids if __name__ == "__main__": pytest.main([__file__, "-v"])