def test_token_count_truncation(self, mock_search_results):
        """Test truncating context based on token count."""
        # Function to truncate context to max tokens
        def truncate_context_by_tokens(search_results, max_tokens=1000):
            # Simple token counting function (in real code, use a tokenizer)
            def estimate_tokens(text):
                # Rough approximation: 4 chars ~ 1 token
                return len(text) // 4

            context_items = []
            current_tokens = 0

            # Add chunks until we hit the token limit
            for result in search_results["chunk_search_results"]:
                chunk_text = result["text"]
                chunk_tokens = estimate_tokens(chunk_text)

                if current_tokens + chunk_tokens > max_tokens:
                    # If this chunk would exceed the limit, stop
                    break

                # Add this chunk and update token count
                context_items.append(chunk_text)
                current_tokens += chunk_tokens

            return "\n\n".join(context_items)

        # Truncate to a small token limit (should fit ~2-3 chunks)
        small_context = truncate_context_by_tokens(mock_search_results, max_tokens=50)

        # Check truncation
        chunk_count = small_context.count("search result")
        assert 1 <= chunk_count <= 3, "Should only include 1-3 chunks with small token limit"

        # Truncate with larger limit (should fit all chunks)
        large_context = truncate_context_by_tokens(mock_search_results, max_tokens=1000)
        large_chunk_count = large_context.count("search result")
        assert large_chunk_count == 5, "Should include all 5 chunks with large token limit"
Test truncating context based on token count.
py/tests/unit/retrieval/test_rag_processing.py