learningcircuit--local-deep-research/tests/programmatic_access/test_ollama_integration.py

"""Integration tests for Ollama LLM with real text generation."""

import pytest
import os
from typing import List
from datetime import datetime
from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_core.retrievers import Document

from local_deep_research.api import quick_summary


# Skip these tests if SKIP_OLLAMA_TESTS is set
pytestmark = pytest.mark.skipif(
    os.environ.get("SKIP_OLLAMA_TESTS", "true").lower() == "true",
    reason="Ollama integration tests skipped (set SKIP_OLLAMA_TESTS=false to run)",
)


def create_test_documents() -> List[Document]:
    """Create a small set of test documents."""
    return [
        Document(
            page_content="Python is a high-level, interpreted programming language known for its readability and versatility. It supports multiple programming paradigms including procedural, object-oriented, and functional programming.",
            metadata={"source": "python_overview.txt", "topic": "programming"},
        ),
        Document(
            page_content="Machine learning is a subset of artificial intelligence that enables systems to learn and improve from experience without being explicitly programmed. It uses algorithms to parse data, learn from it, and make decisions.",
            metadata={"source": "ml_intro.txt", "topic": "machine_learning"},
        ),
        Document(
            page_content="Deep learning is a specialized subset of machine learning that uses artificial neural networks with multiple layers. It excels at tasks like image recognition, natural language processing, and speech recognition.",
            metadata={"source": "deep_learning.txt", "topic": "deep_learning"},
        ),
    ]


@pytest.fixture
def ollama_llm_factory():
    """Create a factory function for Ollama LLM."""

    def create_llm(model_name="gemma3:12b", temperature=0.7, **kwargs):
        """Factory that creates ChatOllama instances."""
        # Use the provided model_name or default
        actual_model = model_name
        return ChatOllama(
            model=actual_model,
            temperature=temperature,
            num_predict=kwargs.get("max_tokens", 256),
        )

    return create_llm


@pytest.fixture
def memory_retriever():
    """Create an in-memory retriever with test documents."""
    documents = create_test_documents()

    # Create embeddings
    embeddings = OllamaEmbeddings(
        model="jeffh/intfloat-multilingual-e5-large-instruct:f16"
    )

    # Create vector store
    vectorstore = FAISS.from_documents(
        documents=documents, embedding=embeddings
    )

    # Return retriever
    return vectorstore.as_retriever(
        search_kwargs={"k": 2}  # Return top 2 documents
    )


def write_test_summary(
    test_name: str, result: dict, output_dir: str = "test_outputs"
):
    """Write test results to a summary file."""
    os.makedirs(output_dir, exist_ok=True)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"{output_dir}/ollama_test_{test_name}_{timestamp}.md"

    with open(filename, "w") as f:
        f.write(f"# Ollama Integration Test: {test_name}\n\n")
        f.write(f"**Timestamp**: {datetime.now().isoformat()}\n\n")
        f.write(f"**Query**: {result.get('query', 'N/A')}\n\n")
        f.write("## Generated Summary\n\n")
        f.write(f"{result.get('summary', 'No summary generated')}\n\n")

        if result.get("findings"):
            f.write("## Findings\n\n")
            for i, finding in enumerate(result["findings"], 1):
                f.write(f"{i}. {finding}\n")
            f.write("\n")

        if result.get("sources"):
            f.write("## Sources\n\n")
            for source in result["sources"]:
                f.write(f"- {source}\n")

    return filename


def test_ollama_quick_summary_real_generation(memory_retriever):
    """Test quick_summary with real Ollama text generation."""

    # Create ChatOllama LLM instance directly
    llm = ChatOllama(
        model="gemma3:12b",
        temperature=0.3,
        num_predict=256,
    )

    # Perform quick summary with real generation
    result = quick_summary(
        query="What is Python and how is it used in machine learning?",
        llms={"ollama": llm},  # Pass LLM instance directly
        retrievers={"test_docs": memory_retriever},
        provider="ollama",
        search_tool="test_docs",
    )

    # Verify we got a real response
    assert "summary" in result
    assert isinstance(result["summary"], str)
    assert len(result["summary"]) > 50  # Should be a meaningful summary

    # The summary should mention Python and ML based on our documents
    summary_lower = result["summary"].lower()
    assert any(
        term in summary_lower for term in ["python", "programming", "language"]
    )
    assert any(
        term in summary_lower
        for term in ["machine learning", "ml", "learning", "ai"]
    )

    # Check other fields
    assert "findings" in result
    assert isinstance(result["findings"], list)

    # Write summary to file
    result["query"] = "What is Python and how is it used in machine learning?"
    output_file = write_test_summary("quick_summary", result)

    # Print the actual generated summary for verification
    print("\n=== GENERATED SUMMARY ===")
    print(result["summary"])
    print("\n=== FINDINGS ===")
    for i, finding in enumerate(result.get("findings", [])[:3]):
        print(f"{i + 1}. {finding}")
    print(f"\n=== Summary written to: {output_file} ===")


def test_ollama_with_multiple_queries(ollama_llm_factory, memory_retriever):
    """Test multiple queries to verify consistent operation."""

    queries = [
        "What is deep learning?",
        "How does Python relate to AI development?",
        "Explain the difference between machine learning and deep learning",
    ]

    all_results = []
    summaries = []

    for query in queries:
        result = quick_summary(
            query=query,
            llms={"ollama": ollama_llm_factory},
            retrievers={"docs": memory_retriever},
            provider="ollama",
            search_tool="docs",
            temperature=0.5,
        )

        # Verify each query produces a summary
        assert "summary" in result
        assert len(result["summary"]) > 30

        result["query"] = query
        all_results.append(result)
        summaries.append(result["summary"])

    # All summaries should be different (not cached or static)
    assert len(set(summaries)) == len(summaries), (
        "All summaries should be unique"
    )

    # Write combined summary
    combined_result = {
        "summary": "\n\n---\n\n".join(
            f"**Query**: {r['query']}\n\n{r['summary']}" for r in all_results
        ),
        "findings": [],
        "query": "Multiple queries test",
    }
    output_file = write_test_summary("multiple_queries", combined_result)

    # Print summaries for manual verification
    print("\n=== MULTIPLE QUERY RESULTS ===")
    for query, summary in zip(queries, summaries):
        print(f"\nQuery: {query}")
        print(f"Summary: {summary[:200]}...")
    print(f"\n=== Combined summary written to: {output_file} ===")


def test_ollama_factory_with_different_parameters(memory_retriever):
    """Test that factory parameters are properly passed through."""

    def custom_factory(model_name="gemma3:12b", temperature=0.7, **kwargs):
        """Factory with custom defaults."""
        # Track what parameters were received
        print(
            f"\nFactory called with: model_name={model_name}, temp={temperature}, kwargs={kwargs}"
        )

        return ChatOllama(
            model=model_name,
            temperature=temperature,
            num_predict=kwargs.get("max_tokens", 100),
        )

    # Test with custom parameters
    result = quick_summary(
        query="Brief explanation of Python",
        llms={"custom": custom_factory},
        retrievers={"docs": memory_retriever},
        provider="custom",
        search_tool="docs",
        temperature=0.1,  # Should override factory default
        max_tokens=150,  # Should be passed to factory
    )

    assert "summary" in result
    assert len(result["summary"]) > 20

    # Write summary
    result["query"] = "Brief explanation of Python"
    output_file = write_test_summary("custom_parameters", result)
    print(f"\nCustom parameters test summary written to: {output_file}")


def test_retriever_actually_retrieves_documents(memory_retriever):
    """Verify the retriever is working correctly."""

    # Test retriever directly
    docs = memory_retriever.get_relevant_documents("Python programming")

    assert len(docs) > 0
    assert all(isinstance(doc.page_content, str) for doc in docs)

    # Should retrieve Python-related content
    combined_content = " ".join(doc.page_content for doc in docs).lower()
    assert "python" in combined_content


@pytest.mark.parametrize("temperature", [0.1, 0.5, 0.9])
def test_temperature_affects_generation(
    ollama_llm_factory, memory_retriever, temperature
):
    """Test that different temperatures produce different outputs."""

    result = quick_summary(
        query="Describe machine learning",
        llms={"ollama": ollama_llm_factory},
        retrievers={"docs": memory_retriever},
        provider="ollama",
        search_tool="docs",
        temperature=temperature,
    )

    assert "summary" in result
    print(f"\nTemp {temperature} summary: {result['summary'][:100]}...")


if __name__ == "__main__":
    # Allow running directly for debugging
    print("Running Ollama integration tests...")
    print("Make sure Ollama is running and models are available:")
    print("  ollama pull gemma3:12b")
    print("  ollama pull jeffh/intfloat-multilingual-e5-large-instruct:f16")

    # Run a simple test
    try:

        def factory(**kwargs):
            return ChatOllama(model="gemma3:12b", **kwargs)

        # Create simple retriever
        docs = [Document(page_content="Test content about Python programming.")]
        embeddings = OllamaEmbeddings(
            model="jeffh/intfloat-multilingual-e5-large-instruct:f16"
        )
        vectorstore = FAISS.from_documents(docs, embeddings)
        retriever = vectorstore.as_retriever()

        result = quick_summary(
            query="What is Python?",
            llms={"test": factory},
            retrievers={"test": retriever},
            provider="test",
            search_tool="test",
        )

        print(
            f"\nSuccess! Generated summary: {result.get('summary', 'No summary')}"
        )

    except Exception as e:
        print(f"\nError: {e}")
        import traceback

        traceback.print_exc()