"""Test demonstrating programmatic access with Langchain Ollama LLM and in-memory vector retriever.""" import pytest from unittest.mock import patch, MagicMock import requests from langchain_ollama import OllamaEmbeddings, OllamaLLM from langchain_community.vectorstores import FAISS from langchain_core.retrievers import Document from local_deep_research.api import ( quick_summary, detailed_research, generate_report, ) from local_deep_research.llm import clear_llm_registry def _is_ollama_running(): """Check if Ollama service is running.""" try: response = requests.get("http://localhost:11434/api/tags", timeout=1) return response.status_code == 200 except Exception: return False # Skip entire module if Ollama is not running to avoid fixture initialization hangs pytestmark = pytest.mark.skipif( not _is_ollama_running(), reason="Ollama is not running - skipping all Ollama integration tests", ) @pytest.fixture(autouse=True) def clear_registries(): """Clear registries before and after each test.""" clear_llm_registry() yield clear_llm_registry() @pytest.fixture def sample_documents(): """Create sample documents for the vector store.""" docs = [ Document( page_content="Machine learning is a subset of artificial intelligence that enables systems to learn from data.", metadata={"source": "ml_basics.txt", "topic": "machine_learning"}, ), Document( page_content="Deep learning uses neural networks with multiple layers to extract features from raw data.", metadata={"source": "dl_intro.txt", "topic": "deep_learning"}, ), Document( page_content="Natural language processing allows computers to understand and generate human language.", metadata={"source": "nlp_guide.txt", "topic": "nlp"}, ), Document( page_content="Computer vision enables machines to interpret and analyze visual information from images and videos.", metadata={"source": "cv_overview.txt", "topic": "computer_vision"}, ), Document( page_content="Reinforcement learning trains agents to make decisions by rewarding desired behaviors.", metadata={ "source": "rl_basics.txt", "topic": "reinforcement_learning", }, ), ] return docs @pytest.fixture def ollama_llm(): """Create an Ollama LLM instance.""" # Using gemma3n:e4b as requested return OllamaLLM( model="gemma3n:e4b", temperature=0.7, ) @pytest.fixture def memory_retriever(sample_documents): """Create an in-memory vector store retriever.""" # Create embeddings using the specified multilingual model embeddings = OllamaEmbeddings( model="jeffh/intfloat-multilingual-e5-large-instruct:f16" ) # Create FAISS vector store from documents vectorstore = FAISS.from_documents( documents=sample_documents, embedding=embeddings ) # Create retriever from vector store retriever = vectorstore.as_retriever( search_kwargs={"k": 3} # Return top 3 most relevant documents ) return retriever @pytest.fixture def mock_search_system(): """Create a mock search system for testing.""" system = MagicMock() system.analyze_topic.return_value = { "current_knowledge": "Analysis using Ollama LLM and in-memory retriever", "findings": [ "Successfully retrieved relevant documents from vector store", "Ollama LLM provided coherent responses", ], "iterations": 2, "questions": { "iteration_1": ["What is the main concept?"], "iteration_2": ["How is it applied in practice?"], }, "formatted_findings": "## Research Summary\n- Vector retrieval worked effectively", "all_links_of_system": [], } system.model = MagicMock() return system @pytest.mark.skipif( not _is_ollama_running(), reason="Ollama is not running - skipping integration test", ) def test_quick_summary_with_ollama_and_memory_retriever( ollama_llm, memory_retriever, mock_search_system ): """Test quick_summary using Ollama LLM and in-memory vector retriever.""" with patch( "local_deep_research.api.research_functions._init_search_system" ) as mock_init: mock_init.return_value = mock_search_system # Use programmatic API with Ollama and memory retriever result = quick_summary( query="What is deep learning and how does it relate to machine learning?", llms={"ollama_llm": ollama_llm}, retrievers={"memory_docs": memory_retriever}, provider="ollama_llm", search_tool="memory_docs", temperature=0.5, ) # Verify result structure assert "summary" in result assert ( result["summary"] == "Analysis using Ollama LLM and in-memory retriever" ) assert len(result["findings"]) == 2 assert "vector store" in result["findings"][0].lower() # Verify components were configured correctly init_kwargs = mock_init.call_args[1] assert init_kwargs["provider"] == "ollama_llm" assert init_kwargs["search_tool"] == "memory_docs" assert init_kwargs["temperature"] == 0.5 @pytest.mark.skipif( not _is_ollama_running(), reason="Ollama is not running - skipping integration test", ) def test_detailed_research_with_ollama_and_memory_retriever( ollama_llm, memory_retriever, mock_search_system ): """Test detailed_research with Ollama and memory retriever.""" with patch( "local_deep_research.api.research_functions._init_search_system" ) as mock_init: mock_init.return_value = mock_search_system result = detailed_research( query="Explain the differences between various machine learning approaches", llms={"ollama": ollama_llm}, retrievers={"local_docs": memory_retriever}, provider="ollama", search_tool="local_docs", iterations=3, questions_per_iteration=2, ) # Verify detailed research results assert ( result["query"] == "Explain the differences between various machine learning approaches" ) assert ( result["summary"] == "Analysis using Ollama LLM and in-memory retriever" ) assert result["metadata"]["iterations_requested"] == 3 assert result["metadata"]["search_tool"] == "local_docs" @pytest.mark.skipif( not _is_ollama_running(), reason="Ollama is not running - skipping integration test", ) def test_generate_report_with_ollama_and_memory_retriever( ollama_llm, memory_retriever ): """Test report generation using Ollama and memory retriever.""" with patch( "local_deep_research.api.research_functions._init_search_system" ) as mock_init: with patch( "local_deep_research.api.research_functions.IntegratedReportGenerator" ) as mock_report_gen: # Setup mocks mock_system = MagicMock() mock_system.analyze_topic.return_value = { "findings": "Initial ML findings" } mock_init.return_value = mock_system mock_generator = MagicMock() mock_generator.generate_report.return_value = { "content": "# Machine Learning Overview\n\n## Introduction\nThis report covers key ML concepts from local documents.", "metadata": { "query": "machine learning overview", "sources_used": 5, }, } mock_report_gen.return_value = mock_generator # Generate report result = generate_report( query="Create a comprehensive overview of machine learning concepts", llms={"ollama": ollama_llm}, retrievers={"vector_store": memory_retriever}, provider="ollama", search_tool="vector_store", searches_per_section=2, ) # Verify report generation assert "content" in result assert "# Machine Learning Overview" in result["content"] assert "local documents" in result["content"] # Verify configuration init_kwargs = mock_init.call_args[1] assert init_kwargs["provider"] == "ollama" assert init_kwargs["search_tool"] == "vector_store" @pytest.mark.skipif( not _is_ollama_running(), reason="Ollama is not running - skipping integration test", ) def test_custom_vector_store_with_more_documents(): """Test creating a larger in-memory vector store.""" # Create more documents extended_docs = [ Document( page_content="Transfer learning allows models trained on one task to be adapted for related tasks.", metadata={"source": "transfer_learning.txt"}, ), Document( page_content="Attention mechanisms help models focus on relevant parts of the input data.", metadata={"source": "attention.txt"}, ), Document( page_content="Gradient descent is an optimization algorithm used to train neural networks.", metadata={"source": "optimization.txt"}, ), Document( page_content="Convolutional neural networks are particularly effective for image processing tasks.", metadata={"source": "cnn.txt"}, ), Document( page_content="Recurrent neural networks can process sequential data like text or time series.", metadata={"source": "rnn.txt"}, ), ] # Create vector store with extended documents embeddings = OllamaEmbeddings( model="jeffh/intfloat-multilingual-e5-large-instruct:f16" ) vectorstore = FAISS.from_documents( documents=extended_docs, embedding=embeddings ) # Create retriever with different search parameters retriever = vectorstore.as_retriever( search_type="similarity", search_kwargs={"k": 5}, # Return top 5 documents ) # Create Ollama LLM llm = OllamaLLM(model="gemma3n:e4b", temperature=0.3) with patch( "local_deep_research.api.research_functions._init_search_system" ) as mock_init: mock_system = MagicMock() mock_system.analyze_topic.return_value = { "current_knowledge": "Extended document analysis complete", "findings": ["Found relevant information about neural networks"], } mock_init.return_value = mock_system result = quick_summary( query="How do different types of neural networks work?", llms={"ollama": llm}, retrievers={"extended_docs": retriever}, provider="ollama", search_tool="extended_docs", ) assert "summary" in result assert result["summary"] == "Extended document analysis complete" @pytest.mark.skipif( not _is_ollama_running(), reason="Ollama is not running - skipping integration test", ) def test_multiple_retrievers_with_ollama(): """Test using multiple in-memory retrievers with Ollama.""" # Create first retriever for ML topics ml_docs = [ Document( page_content="Supervised learning uses labeled data for training." ), Document( page_content="Unsupervised learning finds patterns in unlabeled data." ), ] # Create second retriever for application topics app_docs = [ Document( page_content="ML is used in recommendation systems for personalized content." ), Document( page_content="ML powers autonomous vehicles through computer vision and sensor fusion." ), ] embeddings = OllamaEmbeddings( model="jeffh/intfloat-multilingual-e5-large-instruct:f16" ) ml_vectorstore = FAISS.from_documents(ml_docs, embeddings) app_vectorstore = FAISS.from_documents(app_docs, embeddings) ml_retriever = ml_vectorstore.as_retriever() app_retriever = app_vectorstore.as_retriever() ollama_llm = OllamaLLM(model="gemma3n:e4b") with patch( "local_deep_research.api.research_functions._init_search_system" ) as mock_init: mock_system = MagicMock() mock_system.analyze_topic.return_value = { "current_knowledge": "Analysis from multiple vector stores", "findings": ["ML concepts retrieved", "Applications identified"], } mock_init.return_value = mock_system result = quick_summary( query="What are ML techniques and their applications?", llms={"ollama": ollama_llm}, retrievers={ "ml_concepts": ml_retriever, "ml_applications": app_retriever, }, provider="ollama", search_tool="ml_concepts", # Use a registered retriever ) assert "summary" in result assert "multiple vector stores" in result["summary"] @pytest.mark.skipif( not _is_ollama_running(), reason="Ollama is not running - skipping integration test", ) def test_simple_ollama_factory_pattern(): """Test using a factory function to create Ollama instances.""" def create_ollama_llm(model_name="gemma3n:e4b", temperature=0.7, **kwargs): """Factory function for creating configured Ollama instances.""" return OllamaLLM( model=model_name, temperature=temperature, num_predict=kwargs.get("max_tokens", 256), ) # Create simple in-memory retriever docs = [Document(page_content="Test content for factory pattern demo.")] embeddings = OllamaEmbeddings( model="jeffh/intfloat-multilingual-e5-large-instruct:f16" ) vectorstore = FAISS.from_documents(docs, embeddings) retriever = vectorstore.as_retriever() with patch( "local_deep_research.api.research_functions._init_search_system" ) as mock_init: mock_system = MagicMock() mock_system.analyze_topic.return_value = { "current_knowledge": "Factory pattern test successful", "findings": [], } mock_init.return_value = mock_system result = quick_summary( query="Test factory pattern", llms={"ollama_factory": create_ollama_llm}, retrievers={"test_docs": retriever}, provider="ollama_factory", search_tool="test_docs", model_name="gemma3n:e4b", temperature=0.2, max_tokens=512, ) assert "summary" in result assert "Factory pattern test successful" in result["summary"]