learningcircuit--local-deep-research/tests/text_optimization/test_integration.py

"""Integration tests for text optimization with the research service."""

import pytest
from unittest.mock import patch

from local_deep_research.web.services.research_service import (
    get_citation_formatter,
    export_report_to_memory,
)
from local_deep_research.text_optimization import (
    CitationFormatter,
    CitationMode,
)


class TestResearchServiceIntegration:
    """Test integration of text optimization with research service."""

    @patch("local_deep_research.config.search_config.get_setting_from_snapshot")
    def test_get_citation_formatter_number_mode(self, mock_get_setting):
        """Test getting formatter with number hyperlinks mode."""
        mock_get_setting.return_value = "number_hyperlinks"

        formatter = get_citation_formatter()
        assert isinstance(formatter, CitationFormatter)
        assert formatter.mode == CitationMode.NUMBER_HYPERLINKS

    @patch("local_deep_research.config.search_config.get_setting_from_snapshot")
    def test_get_citation_formatter_domain_mode(self, mock_get_setting):
        """Test getting formatter with domain hyperlinks mode."""
        mock_get_setting.return_value = "domain_hyperlinks"

        formatter = get_citation_formatter()
        assert isinstance(formatter, CitationFormatter)
        assert formatter.mode == CitationMode.DOMAIN_HYPERLINKS

    @patch("local_deep_research.config.search_config.get_setting_from_snapshot")
    def test_get_citation_formatter_no_hyperlinks_mode(self, mock_get_setting):
        """Test getting formatter with no hyperlinks mode."""
        mock_get_setting.return_value = "no_hyperlinks"

        formatter = get_citation_formatter()
        assert isinstance(formatter, CitationFormatter)
        assert formatter.mode == CitationMode.NO_HYPERLINKS

    @patch("local_deep_research.config.search_config.get_setting_from_snapshot")
    def test_get_citation_formatter_invalid_mode(self, mock_get_setting):
        """Test getting formatter with invalid mode falls back to default."""
        mock_get_setting.return_value = "invalid_mode"

        formatter = get_citation_formatter()
        assert isinstance(formatter, CitationFormatter)
        assert formatter.mode == CitationMode.NUMBER_HYPERLINKS  # Default

    def test_export_report_to_latex(self):
        """Test LaTeX export functionality."""
        markdown_content = """# Test Report

This is a test with citation [1].

## Sources

[1] Test Source
    URL: https://example.com
"""
        # Export to LaTeX
        latex_bytes, filename, mimetype = export_report_to_memory(
            markdown_content, "latex"
        )

        # Check export results
        assert filename.endswith(".tex")
        assert mimetype == "text/plain"

        # Check LaTeX content
        latex_content = latex_bytes.decode("utf-8")

        assert r"\documentclass[12pt]{article}" in latex_content
        assert r"\section{Test Report}" in latex_content
        assert r"\cite{1}" in latex_content
        assert r"\bibitem{1}" in latex_content

    @patch("local_deep_research.config.search_config.get_setting_from_snapshot")
    def test_real_world_citation_formatting(self, mock_get_setting):
        """Test citation formatting with real-world example."""
        mock_get_setting.return_value = "number_hyperlinks"

        # Simulate research report content
        content = """# Deep Learning Research Summary

Query: What are the latest advances in transformer architectures?

## Executive Summary

Recent advances in transformer architectures have focused on efficiency improvements [1],
novel attention mechanisms [2], and scaling laws [3]. The field has seen rapid progress
with models like GPT-4 [4] and Claude [5] demonstrating impressive capabilities.

## Key Findings

### Efficiency Improvements
Multiple research groups have proposed methods to reduce computational complexity [1, 2, 3]:
- Flash Attention reduces memory usage significantly [1]
- Sparse transformers achieve O(n√n) complexity [2]
- Linear attention approximations show promise [3]

### Novel Architectures
Recent architectural innovations include:
- Mixture of Experts (MoE) models [4][5]
- Retrieval-augmented generation [6]
- Tool-use capabilities [7][8][9]

## Sources

[1] FlashAttention: Fast and Memory-Efficient Exact Attention
    URL: https://arxiv.org/abs/2205.14135

[2] Efficient Transformers: A Survey
    URL: https://arxiv.org/abs/2009.06732

[3] Linformer: Self-Attention with Linear Complexity
    URL: https://arxiv.org/abs/2006.04768

[4] Mixtral of Experts
    URL: https://arxiv.org/abs/2401.04088

[5] Switch Transformers
    URL: https://arxiv.org/abs/2101.03961

[6] Retrieval-Augmented Generation
    URL: https://arxiv.org/abs/2005.11401

[7] Toolformer: Language Models Can Teach Themselves
    URL: https://arxiv.org/abs/2302.04761

[8] WebGPT: Browser-assisted question-answering
    URL: https://arxiv.org/abs/2112.09332

[9] Constitutional AI: Harmlessness from AI Feedback
    URL: https://arxiv.org/abs/2212.08073
"""

        formatter = get_citation_formatter()
        result = formatter.format_document(content)

        # Verify all citation formats are handled correctly
        assert "[[1]](https://arxiv.org/abs/2205.14135)" in result
        assert (
            "[[1]](https://arxiv.org/abs/2205.14135)[[2]](https://arxiv.org/abs/2009.06732)[[3]](https://arxiv.org/abs/2006.04768)"
            in result
        )
        assert (
            "[[4]](https://arxiv.org/abs/2401.04088)[[5]](https://arxiv.org/abs/2101.03961)"
            in result
        )
        assert (
            "[[7]](https://arxiv.org/abs/2302.04761)[[8]](https://arxiv.org/abs/2112.09332)[[9]](https://arxiv.org/abs/2212.08073)"
            in result
        )

        # Ensure sources section is preserved
        assert "## Sources" in result
        assert "[1] FlashAttention" in result

    def test_export_report_to_quarto(self):
        """Test Quarto export functionality."""
        markdown_content = """# AI Research Report

This report discusses recent advances [1] and challenges [2].

## Sources

[1] Advances in AI
    URL: https://arxiv.org/abs/2024.1234
[2] AI Challenges
    URL: https://example.com/challenges
"""
        # Export to Quarto
        zip_bytes, filename, mimetype = export_report_to_memory(
            markdown_content, "quarto", "AI Research"
        )

        # Check export results
        assert filename.endswith(".zip")
        assert mimetype == "application/zip"

        # Extract and check content from zip
        import zipfile
        import io

        zip_buffer = io.BytesIO(zip_bytes)
        with zipfile.ZipFile(zip_buffer, "r") as zipf:
            # Check files in zip
            file_list = zipf.namelist()
            qmd_files = [f for f in file_list if f.endswith(".qmd")]
            bib_files = [f for f in file_list if f.endswith(".bib")]

            assert len(qmd_files) == 1, "Should have one .qmd file"
            assert len(bib_files) == 1, "Should have one .bib file"

            # Check Quarto content
            quarto_content = zipf.read(qmd_files[0]).decode("utf-8")

            # Verify YAML header
            assert 'title: "AI Research"' in quarto_content
            assert "bibliography: references.bib" in quarto_content

            # Verify citation conversion
            assert "[@ref1]" in quarto_content
            assert "[@ref2]" in quarto_content

            # Check bibliography content
            bib_content = zipf.read(bib_files[0]).decode("utf-8")

            assert "@misc{ref1," in bib_content
            assert 'title = "{Advances in AI}"' in bib_content

    def test_export_with_invalid_format(self):
        """Test export with invalid format raises error."""
        markdown_content = "# Test"

        with pytest.raises(ValueError, match="Unsupported export format"):
            export_report_to_memory(markdown_content, "invalid_format")

    @patch("local_deep_research.config.search_config.get_setting_from_snapshot")
    def test_automatic_export_formats(self, mock_get_setting):
        """Test automatic export to multiple formats based on settings."""
        # This would be called in the actual research service when saving reports
        mock_get_setting.return_value = ["markdown", "latex", "quarto"]

        # Simulate the export logic
        export_formats = mock_get_setting("report.export_formats", ["markdown"])

        assert "markdown" in export_formats
        assert "latex" in export_formats
        assert "quarto" in export_formats