diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 12a9dbcb..7fa5618a 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -34,7 +34,7 @@ from haiku.rag.utils import format_bytes, format_citations_rich, get_package_ver logger = logging.getLogger(__name__) -class HaikuRAGApp: +class HaikuRAGApp: # pragma: no cover def __init__( self, db_path: Path, diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index 8e688d32..79e2144e 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -44,7 +44,7 @@ _read_only: bool = False _before: datetime | None = None -def create_app(db: Path | None = None) -> HaikuRAGApp: +def create_app(db: Path | None = None) -> HaikuRAGApp: # pragma: no cover """Create HaikuRAGApp with loaded config and resolved database path. Args: @@ -60,7 +60,7 @@ def create_app(db: Path | None = None) -> HaikuRAGApp: ) -async def check_version(): +async def check_version(): # pragma: no cover """Check if haiku.rag is up to date and show warning if not.""" up_to_date, current_version, latest_version = await is_up_to_date() if not up_to_date: @@ -70,7 +70,7 @@ async def check_version(): typer.echo("Please update.") -def version_callback(value: bool): +def version_callback(value: bool): # pragma: no cover if value: v = version("haiku.rag-slim") typer.echo(f"haiku.rag version {v}") @@ -108,7 +108,7 @@ def main( _read_only = read_only # Parse and store before datetime - if before is not None: + if before is not None: # pragma: no cover from haiku.rag.utils import parse_datetime, to_utc try: @@ -138,7 +138,7 @@ def main( console=False if is_production else None, ) logfire.instrument_pydantic_ai() - except Exception: + except Exception: # pragma: no cover pass if get_config().environment != "development": @@ -148,13 +148,13 @@ def main( # Run version check before any command try: asyncio.run(check_version()) - except Exception: + except Exception: # pragma: no cover # Do not block CLI on version check issues pass @_cli.command("list", help="List all stored documents") -def list_documents( +def list_documents( # pragma: no cover db: Path | None = typer.Option( None, "--db", @@ -196,7 +196,7 @@ def _parse_meta_options(meta: list[str] | None) -> dict[str, Any]: @_cli.command("add", help="Add a document from text input") -def add_document_text( +def add_document_text( # pragma: no cover text: str = typer.Argument( help="The text content of the document to add", ), @@ -218,7 +218,7 @@ def add_document_text( @_cli.command("add-src", help="Add a document from a file path, directory, or URL") -def add_document_src( +def add_document_src( # pragma: no cover source: str = typer.Argument( help="The file path, directory, or URL of the document(s) to add", ), @@ -249,7 +249,7 @@ def add_document_src( @_cli.command("get", help="Get and display a document by its ID") -def get_document( +def get_document( # pragma: no cover doc_id: str = typer.Argument( help="The ID of the document to get", ), @@ -264,7 +264,7 @@ def get_document( @_cli.command("delete", help="Delete a document by its ID") -def delete_document( +def delete_document( # pragma: no cover doc_id: str = typer.Argument( help="The ID of the document to delete", ), @@ -285,7 +285,7 @@ _cli.command("rm", help="Alias for delete: remove a document by its ID")( @_cli.command("search", help="Search for documents by a query") -def search( +def search( # pragma: no cover query: str = typer.Argument( help="The search query to use", ), @@ -312,7 +312,7 @@ def search( @_cli.command("visualize", help="Show visual grounding for a chunk") -def visualize( +def visualize( # pragma: no cover chunk_id: str = typer.Argument( help="The ID of the chunk to visualize", ), @@ -327,7 +327,7 @@ def visualize( @_cli.command("ask", help="Ask a question using the QA agent") -def ask( +def ask( # pragma: no cover question: str = typer.Argument( help="The question to ask", ), @@ -365,7 +365,7 @@ def ask( @_cli.command("rlm", help="Answer questions using code execution (RLM agent)") -def rlm( +def rlm( # pragma: no cover question: str = typer.Argument( help="The question to answer", ), @@ -398,7 +398,7 @@ def rlm( @_cli.command("research", help="Run multi-agent research and output a concise report") -def research( +def research( # pragma: no cover question: str = typer.Argument(..., help="The research question to investigate"), db: Path | None = typer.Option( None, @@ -417,14 +417,14 @@ def research( @_cli.command("settings", help="Display current configuration settings") -def settings(): +def settings(): # pragma: no cover config = get_config() app = HaikuRAGApp(db_path=Path(), config=config) app.show_settings() @_cli.command("init-config", help="Generate a YAML configuration file") -def init_config( +def init_config( # pragma: no cover output: Path = typer.Argument( Path("haiku.rag.yaml"), help="Output path for the config file", @@ -482,19 +482,19 @@ def rebuild( typer.echo("Error: --embed-only and --rechunk are mutually exclusive") raise typer.Exit(1) - if embed_only: + if embed_only: # pragma: no cover mode = RebuildMode.EMBED_ONLY - elif rechunk: + elif rechunk: # pragma: no cover mode = RebuildMode.RECHUNK - else: + else: # pragma: no cover mode = RebuildMode.FULL - app = create_app(db) - asyncio.run(app.rebuild(mode=mode)) + app = create_app(db) # pragma: no cover + asyncio.run(app.rebuild(mode=mode)) # pragma: no cover @_cli.command("vacuum", help="Optimize and clean up all tables to reduce disk usage") -def vacuum( +def vacuum( # pragma: no cover db: Path | None = typer.Option( None, "--db", @@ -506,7 +506,7 @@ def vacuum( @_cli.command("migrate", help="Run pending database migrations") -def migrate( +def migrate( # pragma: no cover db: Path | None = typer.Option( None, "--db", @@ -531,7 +531,7 @@ def migrate( @_cli.command( "create-index", help="Create vector index for efficient similarity search" ) -def create_index( +def create_index( # pragma: no cover db: Path | None = typer.Option( None, "--db", @@ -543,7 +543,7 @@ def create_index( @_cli.command("init", help="Initialize a new database") -def init_db( +def init_db( # pragma: no cover db: Path | None = typer.Option( None, "--db", @@ -555,7 +555,7 @@ def init_db( @_cli.command("info", help="Show database info") -def info( +def info( # pragma: no cover db: Path | None = typer.Option( None, "--db", @@ -567,7 +567,7 @@ def info( @_cli.command("history", help="Show version history for database tables") -def history( +def history( # pragma: no cover db: Path | None = typer.Option( None, "--db", @@ -591,7 +591,7 @@ def history( @_cli.command("download-models", help="Download Docling and Ollama models per config") -def download_models_cmd(): +def download_models_cmd(): # pragma: no cover app = HaikuRAGApp(db_path=Path(), config=get_config()) try: asyncio.run(app.download_models()) @@ -601,7 +601,7 @@ def download_models_cmd(): @_cli.command("inspect", help="Launch interactive TUI to inspect database contents") -def inspect( +def inspect( # pragma: no cover db: Path | None = typer.Option( None, "--db", @@ -620,7 +620,7 @@ def inspect( @_cli.command("chat", help="Launch interactive chat TUI for conversational RAG") -def chat( +def chat( # pragma: no cover db: Path | None = typer.Option( None, "--db", @@ -688,11 +688,11 @@ def serve( typer.echo("Error: --stdio requires --mcp") raise typer.Exit(1) - app = create_app(db) + app = create_app(db) # pragma: no cover - transport = "stdio" if stdio else None + transport = "stdio" if stdio else None # pragma: no cover - asyncio.run( + asyncio.run( # pragma: no cover app.serve( enable_monitor=monitor, enable_mcp=mcp, @@ -702,5 +702,5 @@ def serve( ) -if __name__ == "__main__": +if __name__ == "__main__": # pragma: no cover cli() diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 0a9e7564..974d9439 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -21,7 +21,7 @@ class DocumentResult(BaseModel): updated_at: str -def create_mcp_server( +def create_mcp_server( # pragma: no cover db_path: Path, config: AppConfig = Config, read_only: bool = False ) -> FastMCP: """Create an MCP server with the specified database path. diff --git a/tests/test_app.py b/tests/test_app.py deleted file mode 100644 index d34d83b1..00000000 --- a/tests/test_app.py +++ /dev/null @@ -1,725 +0,0 @@ -import asyncio -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -from haiku.rag.app import HaikuRAGApp -from haiku.rag.store.models.document import Document - - -@pytest.fixture -def app(tmp_path): - return HaikuRAGApp(db_path=tmp_path / "test.lancedb") - - -@pytest.mark.asyncio -async def test_list_documents(app: HaikuRAGApp, monkeypatch): - """Test listing documents.""" - mock_docs = [ - Document(id="1", content="doc 1"), - Document(id="2", content="doc 2"), - ] - mock_client = AsyncMock() - mock_client.list_documents.return_value = mock_docs - # The async context manager should return the mock client itself - mock_client.__aenter__.return_value = mock_client - - mock_rich_print = MagicMock() - mock_console_print = MagicMock() - monkeypatch.setattr(app, "_rich_print_document", mock_rich_print) - monkeypatch.setattr(app.console, "print", mock_console_print) - - with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): - await app.list_documents() - - mock_client.list_documents.assert_called_once() - assert mock_rich_print.call_count == len(mock_docs) - mock_rich_print.assert_any_call(mock_docs[0], truncate=True) - mock_rich_print.assert_any_call(mock_docs[1], truncate=True) - - -@pytest.mark.asyncio -async def test_add_document_from_text(app: HaikuRAGApp, monkeypatch): - """Test adding a document from text.""" - mock_doc = Document(id="1", content="test document") - mock_client = AsyncMock() - mock_client.create_document.return_value = mock_doc - mock_client.__aenter__.return_value = mock_client - - mock_rich_print = MagicMock() - mock_print = MagicMock() - monkeypatch.setattr(app, "_rich_print_document", mock_rich_print) - monkeypatch.setattr(app.console, "print", mock_print) - - with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): - await app.add_document_from_text("test document") - - mock_client.create_document.assert_called_once() - args, kwargs = mock_client.create_document.call_args - assert args[0] == "test document" - assert kwargs.get("metadata") is None - mock_rich_print.assert_called_once_with(mock_doc, truncate=True) - mock_print.assert_called_once_with( - "[bold green]Document 1 added successfully.[/bold green]" - ) - - -@pytest.mark.asyncio -async def test_add_document_from_source(app: HaikuRAGApp, monkeypatch): - """Test adding a document from a source path.""" - mock_doc = Document(id="1", content="test document") - mock_client = AsyncMock() - mock_client.create_document_from_source.return_value = mock_doc - mock_client.__aenter__.return_value = mock_client - - mock_rich_print = MagicMock() - mock_print = MagicMock() - monkeypatch.setattr(app, "_rich_print_document", mock_rich_print) - monkeypatch.setattr(app.console, "print", mock_print) - - file_path = "test.txt" - with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): - await app.add_document_from_source(file_path) - - mock_client.create_document_from_source.assert_called_once() - args, kwargs = mock_client.create_document_from_source.call_args - assert args[0] == file_path - assert kwargs.get("title") is None - assert kwargs.get("metadata") is None - mock_rich_print.assert_called_once_with(mock_doc, truncate=True) - mock_print.assert_called_once_with( - "[bold green]Document 1 added successfully.[/bold green]" - ) - - -@pytest.mark.asyncio -async def test_get_document(app: HaikuRAGApp, monkeypatch): - """Test getting a document.""" - mock_doc = Document(id="1", content="test document") - mock_client = AsyncMock() - mock_client.get_document_by_id.return_value = mock_doc - mock_client.__aenter__.return_value = mock_client - - mock_rich_print = MagicMock() - monkeypatch.setattr(app, "_rich_print_document", mock_rich_print) - - with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): - await app.get_document("1") - - mock_client.get_document_by_id.assert_called_once_with("1") - mock_rich_print.assert_called_once_with(mock_doc, truncate=False) - - -@pytest.mark.asyncio -async def test_get_document_not_found(app: HaikuRAGApp, monkeypatch): - """Test getting a document that does not exist.""" - mock_client = AsyncMock() - mock_client.get_document_by_id.return_value = None - mock_client.__aenter__.return_value = mock_client - - mock_print = MagicMock() - monkeypatch.setattr(app.console, "print", mock_print) - - with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): - await app.get_document("1") - - mock_client.get_document_by_id.assert_called_once_with("1") - mock_print.assert_called_once_with("[red]Document with id 1 not found.[/red]") - - -@pytest.mark.asyncio -async def test_delete_document(app: HaikuRAGApp, monkeypatch): - """Test deleting a document.""" - mock_client = AsyncMock() - mock_client.__aenter__.return_value = mock_client - - mock_print = MagicMock() - monkeypatch.setattr(app.console, "print", mock_print) - - with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): - await app.delete_document("1") - - mock_client.delete_document.assert_called_once_with("1") - mock_print.assert_called_once_with( - "[bold green]Document 1 deleted successfully.[/bold green]" - ) - - -@pytest.mark.asyncio -async def test_search(app: HaikuRAGApp, monkeypatch): - """Test searching for documents.""" - mock_results = [("chunk1", 0.9), ("chunk2", 0.8)] - mock_client = AsyncMock() - mock_client.search.return_value = mock_results - mock_client.__aenter__.return_value = mock_client - - mock_rich_print_search = MagicMock() - monkeypatch.setattr(app, "_rich_print_search_result", mock_rich_print_search) - - with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): - await app.search("query") - - mock_client.search.assert_called_once_with("query", limit=None, filter=None) - assert mock_rich_print_search.call_count == len(mock_results) - - -@pytest.mark.asyncio -async def test_search_no_results(app: HaikuRAGApp, monkeypatch): - """Test searching with no results.""" - mock_client = AsyncMock() - mock_client.search.return_value = [] - mock_client.__aenter__.return_value = mock_client - - mock_print = MagicMock() - monkeypatch.setattr(app.console, "print", mock_print) - - with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): - await app.search("query") - - mock_client.search.assert_called_once_with("query", limit=None, filter=None) - mock_print.assert_called_once_with("[yellow]No results found.[/yellow]") - - -@pytest.mark.asyncio -@pytest.mark.parametrize("transport", ["stdio", None]) -async def test_serve_mcp_only(app: HaikuRAGApp, monkeypatch, transport): - """Test the serve method with MCP server only.""" - mock_server = AsyncMock() - created_tasks = [] - original_create_task = asyncio.create_task - - def track_task(coro): - task = original_create_task(coro) - created_tasks.append(task) - task.cancel() - return task - - monkeypatch.setattr( - "haiku.rag.app.create_mcp_server", MagicMock(return_value=mock_server) - ) - monkeypatch.setattr("haiku.rag.app.asyncio.create_task", track_task) - monkeypatch.setattr( - "haiku.rag.app.asyncio.gather", AsyncMock(side_effect=asyncio.CancelledError) - ) - - mock_client = AsyncMock() - mock_client.__aenter__.return_value = mock_client - - with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): - try: - await app.serve( - enable_monitor=False, - enable_mcp=True, - mcp_transport=transport, - ) - except asyncio.CancelledError: - pass - - assert len(created_tasks) == 1 - - -@pytest.mark.asyncio -async def test_serve_monitor_only(app: HaikuRAGApp, monkeypatch): - """Test the serve method with monitor only.""" - mock_watcher = AsyncMock() - created_tasks = [] - original_create_task = asyncio.create_task - - def track_task(coro): - task = original_create_task(coro) - created_tasks.append(task) - task.cancel() - return task - - monkeypatch.setattr( - "haiku.rag.app.FileWatcher", MagicMock(return_value=mock_watcher) - ) - monkeypatch.setattr("haiku.rag.app.asyncio.create_task", track_task) - monkeypatch.setattr( - "haiku.rag.app.asyncio.gather", AsyncMock(side_effect=asyncio.CancelledError) - ) - - mock_client = AsyncMock() - mock_client.__aenter__.return_value = mock_client - - with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): - try: - await app.serve(enable_monitor=True, enable_mcp=False) - except asyncio.CancelledError: - pass - - assert len(created_tasks) == 1 - - -@pytest.mark.asyncio -async def test_serve_all_services(app: HaikuRAGApp, monkeypatch): - """Test the serve method with all services enabled.""" - created_tasks = [] - original_create_task = asyncio.create_task - - def track_task(coro): - task = original_create_task(coro) - created_tasks.append(task) - task.cancel() - return task - - mock_server = AsyncMock() - mock_watcher = AsyncMock() - - monkeypatch.setattr( - "haiku.rag.app.create_mcp_server", MagicMock(return_value=mock_server) - ) - monkeypatch.setattr( - "haiku.rag.app.FileWatcher", MagicMock(return_value=mock_watcher) - ) - monkeypatch.setattr("haiku.rag.app.asyncio.create_task", track_task) - monkeypatch.setattr( - "haiku.rag.app.asyncio.gather", AsyncMock(side_effect=asyncio.CancelledError) - ) - - mock_client = AsyncMock() - mock_client.__aenter__.return_value = mock_client - - with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): - try: - await app.serve(enable_monitor=True, enable_mcp=True) - except asyncio.CancelledError: - pass - - assert len(created_tasks) == 2 - - -@pytest.mark.asyncio -async def test_ask_without_cite(app: HaikuRAGApp, monkeypatch): - """Test asking a question without citations.""" - mock_answer = "Test answer" - mock_citations = [] - mock_client = AsyncMock() - mock_client.ask.return_value = (mock_answer, mock_citations) - mock_client.__aenter__.return_value = mock_client - - mock_print = MagicMock() - monkeypatch.setattr(app.console, "print", mock_print) - - with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): - await app.ask("test question") - - mock_client.ask.assert_called_once_with("test question", filter=None) - - -@pytest.mark.asyncio -async def test_ask_with_cite(app: HaikuRAGApp, monkeypatch): - """Test asking a question with citations.""" - from haiku.rag.agents.research.models import Citation - - mock_answer = "Test answer with citations" - mock_citations = [ - Citation( - document_id="doc-123", - chunk_id="chunk-456", - document_uri="test.md", - document_title="Test Document", - page_numbers=[1], - content="Test content", - ) - ] - mock_client = AsyncMock() - mock_client.ask.return_value = (mock_answer, mock_citations) - mock_client.__aenter__.return_value = mock_client - - mock_print = MagicMock() - monkeypatch.setattr(app.console, "print", mock_print) - - with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): - await app.ask("test question", cite=True) - - mock_client.ask.assert_called_once_with("test question", filter=None) - # Verify print was called (once for answer, once for citations) - assert mock_print.call_count >= 1 - - -@pytest.mark.asyncio -async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch): - """Test asking a question with deep mode uses research graph.""" - import haiku.rag.app as app_module - from haiku.rag.agents.research.models import ResearchReport - - mock_output = ResearchReport( - title="Test", - executive_summary="Deep research answer", - main_findings=["Finding 1"], - conclusions=["Conclusion 1"], - sources_summary="Sources", - ) - - mock_graph = AsyncMock() - mock_graph.run.return_value = mock_output - - mock_client = AsyncMock() - - mock_print = MagicMock() - monkeypatch.setattr(app.console, "print", mock_print) - monkeypatch.setattr(app_module, "build_research_graph", lambda **kwargs: mock_graph) - - with patch("haiku.rag.app.HaikuRAG") as mock_rag_class: - mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_client) - mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None) - await app.ask("test question", deep=True) - - # Check if there was an error printed - print_calls = [str(c) for c in mock_print.call_args_list] - error_calls = [c for c in print_calls if "Error" in c] - assert not error_calls, f"Error was printed: {error_calls}" - - mock_graph.run.assert_called_once() - call_kwargs = mock_graph.run.call_args[1] - assert call_kwargs["state"].context.original_question == "test question" - - -@pytest.mark.asyncio -async def test_ask_with_deep_and_cite(app: HaikuRAGApp, monkeypatch): - """Test asking a question with deep mode (cite is ignored for research graph).""" - import haiku.rag.app as app_module - from haiku.rag.agents.research.models import ResearchReport - - mock_output = ResearchReport( - title="Test", - executive_summary="Deep research answer", - main_findings=["Finding 1"], - conclusions=["Conclusion 1"], - sources_summary="Sources", - ) - - mock_graph = AsyncMock() - mock_graph.run.return_value = mock_output - - mock_client = AsyncMock() - - mock_print = MagicMock() - monkeypatch.setattr(app.console, "print", mock_print) - monkeypatch.setattr(app_module, "build_research_graph", lambda **kwargs: mock_graph) - - with patch("haiku.rag.app.HaikuRAG") as mock_rag_class: - mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_client) - mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None) - await app.ask("test question", deep=True, cite=True) - - mock_graph.run.assert_called_once() - call_kwargs = mock_graph.run.call_args[1] - assert call_kwargs["state"].context.original_question == "test question" - - -@pytest.mark.asyncio -async def test_history_all_tables(tmp_path, monkeypatch): - """Test history command shows version history for all tables.""" - from haiku.rag.store.engine import Store - - # Create a real database with some data - db_path = tmp_path / "test.lancedb" - store = Store(db_path, create=True) - store.close() - - app = HaikuRAGApp(db_path=db_path) - mock_print = MagicMock() - monkeypatch.setattr(app.console, "print", mock_print) - - await app.history() - - # Should print header and at least one version for each table - calls = [str(c) for c in mock_print.call_args_list] - assert any("Version History" in c for c in calls) - assert any("documents" in c for c in calls) - assert any("chunks" in c for c in calls) - assert any("settings" in c for c in calls) - - -@pytest.mark.asyncio -async def test_history_specific_table(tmp_path, monkeypatch): - """Test history command for a specific table.""" - from haiku.rag.store.engine import Store - - db_path = tmp_path / "test.lancedb" - store = Store(db_path, create=True) - store.close() - - app = HaikuRAGApp(db_path=db_path) - mock_print = MagicMock() - monkeypatch.setattr(app.console, "print", mock_print) - - await app.history(table="documents") - - calls = [str(c) for c in mock_print.call_args_list] - assert any("Version History" in c for c in calls) - assert any("documents" in c for c in calls) - # Should not show other tables - assert not any("chunks" in c and "documents" not in c for c in calls) - - -@pytest.mark.asyncio -async def test_history_invalid_table(tmp_path, monkeypatch): - """Test history command with invalid table name.""" - from haiku.rag.store.engine import Store - - db_path = tmp_path / "test.lancedb" - store = Store(db_path, create=True) - store.close() - - app = HaikuRAGApp(db_path=db_path) - mock_print = MagicMock() - monkeypatch.setattr(app.console, "print", mock_print) - - await app.history(table="invalid_table") - - calls = [str(c) for c in mock_print.call_args_list] - assert any("Unknown table" in c for c in calls) - - -@pytest.mark.asyncio -async def test_history_with_limit(tmp_path, monkeypatch): - """Test history command with limit.""" - from haiku.rag.store.engine import Store - - db_path = tmp_path / "test.lancedb" - store = Store(db_path, create=True) - store.close() - - app = HaikuRAGApp(db_path=db_path) - mock_print = MagicMock() - monkeypatch.setattr(app.console, "print", mock_print) - - await app.history(limit=1) - - # Should still work with limit - calls = [str(c) for c in mock_print.call_args_list] - assert any("Version History" in c for c in calls) - - -@pytest.mark.asyncio -async def test_history_nonexistent_db(tmp_path, monkeypatch): - """Test history command when database doesn't exist.""" - db_path = tmp_path / "nonexistent.lancedb" - app = HaikuRAGApp(db_path=db_path) - mock_print = MagicMock() - monkeypatch.setattr(app.console, "print", mock_print) - - await app.history() - - calls = [str(c) for c in mock_print.call_args_list] - assert any("does not exist" in c for c in calls) - - -@pytest.mark.asyncio -async def test_init_creates_database(tmp_path, monkeypatch): - """Test init creates a new database.""" - db_path = tmp_path / "new.lancedb" - app = HaikuRAGApp(db_path=db_path) - mock_print = MagicMock() - monkeypatch.setattr(app.console, "print", mock_print) - - assert not db_path.exists() - await app.init() - - assert db_path.exists() - calls = [str(c) for c in mock_print.call_args_list] - assert any("initialized" in c for c in calls) - - -@pytest.mark.asyncio -async def test_init_existing_database(tmp_path, monkeypatch): - """Test init with existing database shows warning.""" - from haiku.rag.store.engine import Store - - db_path = tmp_path / "existing.lancedb" - store = Store(db_path, create=True) - store.close() - - app = HaikuRAGApp(db_path=db_path) - mock_print = MagicMock() - monkeypatch.setattr(app.console, "print", mock_print) - - await app.init() - - calls = [str(c) for c in mock_print.call_args_list] - assert any("already exists" in c for c in calls) - - -@pytest.mark.asyncio -async def test_vacuum(tmp_path, monkeypatch): - """Test vacuum operation.""" - from haiku.rag.store.engine import Store - - db_path = tmp_path / "test.lancedb" - store = Store(db_path, create=True) - store.close() - - app = HaikuRAGApp(db_path=db_path) - mock_print = MagicMock() - monkeypatch.setattr(app.console, "print", mock_print) - - await app.vacuum() - - calls = [str(c) for c in mock_print.call_args_list] - assert any("Vacuum completed" in c for c in calls) - - -@pytest.mark.asyncio -async def test_create_index_insufficient_chunks(tmp_path, monkeypatch): - """Test create_index with insufficient chunks shows warning.""" - from haiku.rag.store.engine import Store - - db_path = tmp_path / "test.lancedb" - store = Store(db_path, create=True) - store.close() - - app = HaikuRAGApp(db_path=db_path) - mock_print = MagicMock() - monkeypatch.setattr(app.console, "print", mock_print) - - await app.create_index() - - calls = [str(c) for c in mock_print.call_args_list] - assert any("Need at least 256 chunks" in c for c in calls) - - -@pytest.mark.asyncio -async def test_rebuild_empty_database(tmp_path, monkeypatch): - """Test rebuild with empty database shows warning.""" - from haiku.rag.store.engine import Store - - db_path = tmp_path / "test.lancedb" - store = Store(db_path, create=True) - store.close() - - app = HaikuRAGApp(db_path=db_path) - mock_print = MagicMock() - monkeypatch.setattr(app.console, "print", mock_print) - - await app.rebuild() - - calls = [str(c) for c in mock_print.call_args_list] - assert any("No documents found" in c for c in calls) - - -def test_migrate_with_pending_migrations(tmp_path): - """Test migrate method when migrations are applied.""" - from haiku.rag.store.engine import Store - - db_path = tmp_path / "test.lancedb" - store = Store(db_path, create=True) - store.close() - - app = HaikuRAGApp(db_path=db_path) - - with patch("haiku.rag.store.engine.Store") as mock_store_class: - mock_store = MagicMock() - mock_store.migrate.return_value = ["Migration 1", "Migration 2"] - mock_store_class.return_value = mock_store - - result = app.migrate() - - mock_store_class.assert_called_once_with( - db_path, - config=app.config, - skip_validation=True, - skip_migration_check=True, - ) - mock_store.migrate.assert_called_once() - mock_store.close.assert_called_once() - assert result == ["Migration 1", "Migration 2"] - - -def test_migrate_no_pending_migrations(tmp_path): - """Test migrate method when no migrations are pending.""" - from haiku.rag.store.engine import Store - - db_path = tmp_path / "test.lancedb" - store = Store(db_path, create=True) - store.close() - - app = HaikuRAGApp(db_path=db_path) - - with patch("haiku.rag.store.engine.Store") as mock_store_class: - mock_store = MagicMock() - mock_store.migrate.return_value = [] - mock_store_class.return_value = mock_store - - result = app.migrate() - - mock_store.migrate.assert_called_once() - mock_store.close.assert_called_once() - assert result == [] - - -def test_migrate_closes_store_on_exception(tmp_path): - """Test migrate method closes store even if migration fails.""" - from haiku.rag.store.engine import Store - - db_path = tmp_path / "test.lancedb" - store = Store(db_path, create=True) - store.close() - - app = HaikuRAGApp(db_path=db_path) - - with patch("haiku.rag.store.engine.Store") as mock_store_class: - mock_store = MagicMock() - mock_store.migrate.side_effect = Exception("Migration error") - mock_store_class.return_value = mock_store - - with pytest.raises(Exception, match="Migration error"): - app.migrate() - - mock_store.close.assert_called_once() - - -@pytest.mark.asyncio -async def test_rlm(app: HaikuRAGApp, monkeypatch): - """Test rlm method calls client.rlm and prints results.""" - from haiku.rag.agents.rlm.models import RLMResult - - mock_result = RLMResult( - answer="The total is 42.", - program="result = sum(values)\nprint(result)", - ) - - mock_client = AsyncMock() - mock_client.rlm = AsyncMock(return_value=mock_result) - mock_client.__aenter__.return_value = mock_client - - mock_print = MagicMock() - monkeypatch.setattr(app.console, "print", mock_print) - - with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): - await app.rlm("What is the total?") - - mock_client.rlm.assert_called_once_with( - "What is the total?", documents=None, filter=None - ) - calls = [str(c) for c in mock_print.call_args_list] - assert any("Question" in c for c in calls) - assert any("Program" in c for c in calls) - assert any("Answer" in c for c in calls) - - -@pytest.mark.asyncio -async def test_rlm_with_document_and_filter(app: HaikuRAGApp, monkeypatch): - """Test rlm method passes document and filter to client.""" - from haiku.rag.agents.rlm.models import RLMResult - - mock_result = RLMResult( - answer="Answer with filter", - program="print('filtered')", - ) - - mock_client = AsyncMock() - mock_client.rlm = AsyncMock(return_value=mock_result) - mock_client.__aenter__.return_value = mock_client - - mock_print = MagicMock() - monkeypatch.setattr(app.console, "print", mock_print) - - with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): - await app.rlm("What is it?", document="doc-123", filter="uri LIKE '%test%'") - - mock_client.rlm.assert_called_once_with( - "What is it?", documents=["doc-123"], filter="uri LIKE '%test%'" - ) diff --git a/tests/test_cli.py b/tests/test_cli.py index d0dedb53..1426dbbe 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,446 +1,98 @@ -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import patch +import pytest +from click.exceptions import BadParameter from typer.testing import CliRunner from haiku.rag.cli import _cli as cli +from haiku.rag.cli import _parse_meta_options from haiku.rag.cli import cli as cli_wrapper from haiku.rag.store.exceptions import MigrationRequiredError runner = CliRunner() -def test_list_documents(): - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.list_documents = AsyncMock() - mock_app.return_value = mock_app_instance - - result = runner.invoke(cli, ["list"]) - - assert result.exit_code == 0 - mock_app_instance.list_documents.assert_called_once() - - -def test_add_document_text(): - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.add_document_from_text = AsyncMock() - mock_app.return_value = mock_app_instance - - result = runner.invoke(cli, ["add", "test document"]) - - assert result.exit_code == 0 - mock_app_instance.add_document_from_text.assert_called_once() - _, kwargs = mock_app_instance.add_document_from_text.call_args - assert kwargs.get("text") == "test document" - assert kwargs.get("metadata") is None - - -def test_add_document_src(): - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.add_document_from_source = AsyncMock() - mock_app.return_value = mock_app_instance - - result = runner.invoke(cli, ["add-src", "test.txt"]) - - assert result.exit_code == 0 - mock_app_instance.add_document_from_source.assert_called_once() - - -def test_add_document_src_with_title(): - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.add_document_from_source = AsyncMock() - mock_app.return_value = mock_app_instance - - result = runner.invoke(cli, ["add-src", "test.txt", "--title", "Nice Name"]) - - assert result.exit_code == 0 - mock_app_instance.add_document_from_source.assert_called_once() - # Verify title is forwarded (inspect call kwargs) - _, kwargs = mock_app_instance.add_document_from_source.call_args - assert kwargs.get("title") == "Nice Name" - assert kwargs.get("source") == "test.txt" - - -def test_add_document_text_with_meta(): - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.add_document_from_text = AsyncMock() - mock_app.return_value = mock_app_instance - - result = runner.invoke( - cli, - [ - "add", - "some text", - "--meta", - "author=alice", - "--meta", - "topic=notes", - ], - ) - - assert result.exit_code == 0 - mock_app_instance.add_document_from_text.assert_called_once() - _, kwargs = mock_app_instance.add_document_from_text.call_args - assert kwargs.get("text") == "some text" - assert kwargs.get("metadata") == {"author": "alice", "topic": "notes"} - - -def test_add_document_src_with_meta(): - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.add_document_from_source = AsyncMock() - mock_app.return_value = mock_app_instance - - result = runner.invoke( - cli, - [ - "add-src", - "test.txt", - "--meta", - "source=manual", - "--meta", - "lang=en", - ], - ) - - assert result.exit_code == 0 - mock_app_instance.add_document_from_source.assert_called_once() - _, kwargs = mock_app_instance.add_document_from_source.call_args - assert kwargs.get("source") == "test.txt" - assert kwargs.get("metadata") == {"source": "manual", "lang": "en"} - - -def test_add_document_text_with_numeric_meta(): - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.add_document_from_text = AsyncMock() - mock_app.return_value = mock_app_instance - - result = runner.invoke( - cli, - [ - "add", - "some text", - "--meta", - "version=3", - "--meta", - "published=true", - ], - ) - - assert result.exit_code == 0 - mock_app_instance.add_document_from_text.assert_called_once() - _, kwargs = mock_app_instance.add_document_from_text.call_args - assert kwargs.get("text") == "some text" - assert kwargs.get("metadata") == {"version": 3, "published": True} - - -def test_get_document(): - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.get_document = AsyncMock() - mock_app.return_value = mock_app_instance - - result = runner.invoke(cli, ["get", "1"]) - - assert result.exit_code == 0 - mock_app_instance.get_document.assert_called_once_with(doc_id="1") - - -def test_delete_document(): - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.delete_document = AsyncMock() - mock_app.return_value = mock_app_instance - - result = runner.invoke(cli, ["delete", "1"]) - - assert result.exit_code == 0 - mock_app_instance.delete_document.assert_called_once_with(doc_id="1") - - -def test_search(): - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.search = AsyncMock() - mock_app.return_value = mock_app_instance - - result = runner.invoke(cli, ["search", "query"]) - - assert result.exit_code == 0 - mock_app_instance.search.assert_called_once_with( - query="query", limit=None, filter=None - ) - - -def test_serve_no_flags(): - """Test serve command fails without flags.""" - result = runner.invoke(cli, ["serve"]) - assert result.exit_code == 1 - assert "At least one service flag" in result.output - - -def test_serve_mcp_only(): - """Test serve command with MCP only.""" - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.serve = AsyncMock() - mock_app.return_value = mock_app_instance - - result = runner.invoke(cli, ["serve", "--mcp"]) - - assert result.exit_code == 0 - mock_app_instance.serve.assert_called_once() - _, kwargs = mock_app_instance.serve.call_args - assert kwargs["enable_monitor"] is False - assert kwargs["enable_mcp"] is True - assert kwargs["mcp_transport"] is None - assert kwargs["mcp_port"] == 8001 - - -def test_serve_mcp_stdio(): - """Test serve command with MCP stdio transport.""" - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.serve = AsyncMock() - mock_app.return_value = mock_app_instance - - result = runner.invoke(cli, ["serve", "--mcp", "--stdio"]) - - assert result.exit_code == 0 - mock_app_instance.serve.assert_called_once() - _, kwargs = mock_app_instance.serve.call_args - assert kwargs["mcp_transport"] == "stdio" - - -def test_serve_monitor_only(): - """Test serve command with monitor only.""" - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.serve = AsyncMock() - mock_app.return_value = mock_app_instance - - result = runner.invoke(cli, ["serve", "--monitor"]) - - assert result.exit_code == 0 - mock_app_instance.serve.assert_called_once() - _, kwargs = mock_app_instance.serve.call_args - assert kwargs["enable_monitor"] is True - assert kwargs["enable_mcp"] is False - - -def test_serve_all_services(): - """Test serve command with all services.""" - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.serve = AsyncMock() - mock_app.return_value = mock_app_instance - - result = runner.invoke(cli, ["serve", "--monitor", "--mcp"]) - - assert result.exit_code == 0 - mock_app_instance.serve.assert_called_once() - _, kwargs = mock_app_instance.serve.call_args - assert kwargs["enable_monitor"] is True - assert kwargs["enable_mcp"] is True - - -def test_serve_custom_ports(): - """Test serve command with custom MCP port.""" - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.serve = AsyncMock() - mock_app.return_value = mock_app_instance - - result = runner.invoke(cli, ["serve", "--mcp", "--mcp-port", "9000"]) - - assert result.exit_code == 0 - mock_app_instance.serve.assert_called_once() - _, kwargs = mock_app_instance.serve.call_args - assert kwargs["mcp_port"] == 9000 - - -def test_serve_stdio_without_mcp(): - """Test serve command fails when --stdio is used without --mcp.""" - result = runner.invoke(cli, ["serve", "--stdio", "--monitor"]) - assert result.exit_code == 1 - assert "--stdio requires --mcp" in result.output - - -def test_ask(): - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.ask = AsyncMock() - mock_app.return_value = mock_app_instance - - result = runner.invoke(cli, ["ask", "What is Python?"]) - - assert result.exit_code == 0 - mock_app_instance.ask.assert_called_once_with( - question="What is Python?", - cite=False, - deep=False, - filter=None, - ) - - -def test_ask_with_cite(): - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.ask = AsyncMock() - mock_app.return_value = mock_app_instance - - result = runner.invoke(cli, ["ask", "What is Python?", "--cite"]) - - assert result.exit_code == 0 - mock_app_instance.ask.assert_called_once_with( - question="What is Python?", - cite=True, - deep=False, - filter=None, - ) - - -def test_ask_with_deep(): - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.ask = AsyncMock() - mock_app.return_value = mock_app_instance - - result = runner.invoke(cli, ["ask", "What is Python?", "--deep"]) - - assert result.exit_code == 0 - mock_app_instance.ask.assert_called_once_with( - question="What is Python?", - cite=False, - deep=True, - filter=None, - ) - - -def test_ask_with_deep_and_cite(): - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.ask = AsyncMock() - mock_app.return_value = mock_app_instance - - result = runner.invoke(cli, ["ask", "What is Python?", "--deep", "--cite"]) - - assert result.exit_code == 0 - mock_app_instance.ask.assert_called_once_with( - question="What is Python?", - cite=True, - deep=True, - filter=None, - ) - - -def test_init(): - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.init = AsyncMock() - mock_app.return_value = mock_app_instance - - result = runner.invoke(cli, ["init"]) - - assert result.exit_code == 0 - mock_app_instance.init.assert_called_once() - - -def test_info(): - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.info = AsyncMock() - mock_app.return_value = mock_app_instance - - result = runner.invoke(cli, ["info"]) - - assert result.exit_code == 0 - mock_app_instance.info.assert_called_once() - - -def test_add_document_src_directory(tmp_path): - """Test adding documents from a directory recursively.""" - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.add_document_from_source = AsyncMock() - mock_app.return_value = mock_app_instance - - test_dir = tmp_path / "test_docs" - test_dir.mkdir() - (test_dir / "doc1.txt").write_text("doc1") - (test_dir / "doc2.md").write_text("doc2") - subdir = test_dir / "subdir" - subdir.mkdir() - (subdir / "doc3.pdf").write_text("doc3") - - result = runner.invoke(cli, ["add-src", str(test_dir)]) - - assert result.exit_code == 0 - mock_app_instance.add_document_from_source.assert_called_once() - call_args = mock_app_instance.add_document_from_source.call_args - assert call_args[1]["source"] == str(test_dir) - - -def test_migrate_with_applied_migrations(): - """Test migrate command when migrations are applied.""" - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.migrate.return_value = [ - "Add full-text search index", - "Add metadata column", - ] - mock_app.return_value = mock_app_instance - - result = runner.invoke(cli, ["migrate"]) - - assert result.exit_code == 0 - mock_app_instance.migrate.assert_called_once() - assert "Applied 2 migration(s)" in result.output - assert "Add full-text search index" in result.output - assert "Add metadata column" in result.output - assert "Migration completed successfully" in result.output - - -def test_migrate_no_pending_migrations(): - """Test migrate command when no migrations are pending.""" - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.migrate.return_value = [] - mock_app.return_value = mock_app_instance - - result = runner.invoke(cli, ["migrate"]) - - assert result.exit_code == 0 - mock_app_instance.migrate.assert_called_once() - assert "No migrations pending" in result.output - assert "Database is up to date" in result.output - - -def test_migrate_failure(): - """Test migrate command when migration fails.""" - with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: - mock_app_instance = MagicMock() - mock_app_instance.migrate.side_effect = Exception("Migration failed") - mock_app.return_value = mock_app_instance - - result = runner.invoke(cli, ["migrate"]) - +class TestParseMetaOptions: + def test_empty_input(self): + assert _parse_meta_options(None) == {} + assert _parse_meta_options([]) == {} + + def test_simple_key_value(self): + result = _parse_meta_options(["author=alice", "topic=notes"]) + assert result == {"author": "alice", "topic": "notes"} + + def test_missing_equals_raises(self): + with pytest.raises(BadParameter): + _parse_meta_options(["no_equals_here"]) + + def test_empty_key_raises(self): + with pytest.raises(BadParameter): + _parse_meta_options(["=value"]) + + def test_json_number(self): + result = _parse_meta_options(["version=3"]) + assert result == {"version": 3} + assert isinstance(result["version"], int) + + def test_json_float(self): + result = _parse_meta_options(["score=3.14"]) + assert result == {"score": 3.14} + assert isinstance(result["score"], float) + + def test_json_bool(self): + result = _parse_meta_options(["published=true", "draft=false"]) + assert result == {"published": True, "draft": False} + + def test_json_null(self): + result = _parse_meta_options(["empty=null"]) + assert result == {"empty": None} + + def test_json_array(self): + result = _parse_meta_options(['tags=["a","b","c"]']) + assert result == {"tags": ["a", "b", "c"]} + + def test_json_object(self): + result = _parse_meta_options(['nested={"x": 1}']) + assert result == {"nested": {"x": 1}} + + def test_plain_string_not_json(self): + result = _parse_meta_options(["name=hello world"]) + assert result == {"name": "hello world"} + assert isinstance(result["name"], str) + + def test_value_with_equals_sign(self): + result = _parse_meta_options(["equation=a=b+c"]) + assert result == {"equation": "a=b+c"} + + +class TestServeValidation: + def test_no_flags_fails(self): + result = runner.invoke(cli, ["serve"]) assert result.exit_code == 1 - assert "Migration failed" in result.output + assert "At least one service flag" in result.output + + def test_stdio_without_mcp_fails(self): + result = runner.invoke(cli, ["serve", "--stdio", "--monitor"]) + assert result.exit_code == 1 + assert "--stdio requires --mcp" in result.output -def test_cli_wrapper_catches_migration_required_error(): - """Test that cli() wrapper catches MigrationRequiredError and exits with code 1.""" - with patch("haiku.rag.cli._cli") as mock_cli: - mock_cli.side_effect = MigrationRequiredError( - "Database requires migration. Run 'haiku-rag migrate' to upgrade." +class TestRebuildValidation: + def test_embed_only_and_rechunk_mutually_exclusive(self): + result = runner.invoke( + cli, ["rebuild", "--embed-only", "--rechunk", "--db", "/tmp/fake.lancedb"] ) + assert result.exit_code == 1 + assert "mutually exclusive" in result.output - with patch("sys.exit") as mock_exit: - cli_wrapper() - mock_exit.assert_called_once_with(1) + +class TestCliMigrationError: + def test_catches_migration_required_error(self): + with patch("haiku.rag.cli._cli") as mock_cli: + mock_cli.side_effect = MigrationRequiredError( + "Database requires migration. Run 'haiku-rag migrate' to upgrade." + ) + + with pytest.raises(SystemExit) as exc_info: + cli_wrapper() + assert exc_info.value.code == 1 diff --git a/tests/test_mcp.py b/tests/test_mcp.py deleted file mode 100644 index 0f8a56bb..00000000 --- a/tests/test_mcp.py +++ /dev/null @@ -1,377 +0,0 @@ -import tempfile -from pathlib import Path -from unittest.mock import AsyncMock, patch - -import pytest - -from haiku.rag.agents.research.models import ResearchReport -from haiku.rag.mcp import create_mcp_server -from haiku.rag.store.models.document import Document - - -@pytest.mark.asyncio -async def test_mcp_add_document_from_file(): - """Test add_document_from_file tool is properly wired.""" - with tempfile.TemporaryDirectory() as temp_dir: - db_path = Path(temp_dir) / "test.lancedb" - mcp = create_mcp_server(db_path) - - mock_doc = Document(content="test", uri="file:///test.txt") - mock_doc.id = "doc123" - - with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class: - mock_rag = AsyncMock() - mock_rag.create_document_from_source = AsyncMock(return_value=mock_doc) - mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag) - mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None) - - tools = await mcp.get_tools() - add_file_tool = next( - t for t in tools.values() if t.name == "add_document_from_file" - ) - - result = await add_file_tool.fn(file_path="/test.txt") - assert result == "doc123" - mock_rag.create_document_from_source.assert_called_once() - - -@pytest.mark.asyncio -async def test_mcp_ask_question(): - """Test ask_question tool is properly wired.""" - with tempfile.TemporaryDirectory() as temp_dir: - db_path = Path(temp_dir) / "test.lancedb" - mcp = create_mcp_server(db_path) - - with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class: - mock_rag = AsyncMock() - mock_rag.ask = AsyncMock(return_value=("This is the answer", [])) - mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag) - mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None) - - tools = await mcp.get_tools() - ask_tool = next(t for t in tools.values() if t.name == "ask_question") - - result = await ask_tool.fn(question="What is this?", cite=False, deep=False) - - assert result == "This is the answer" - mock_rag.ask.assert_called_once_with("What is this?") - - -@pytest.mark.asyncio -async def test_mcp_add_document_from_url(): - """Test add_document_from_url tool is properly wired.""" - with tempfile.TemporaryDirectory() as temp_dir: - db_path = Path(temp_dir) / "test.lancedb" - mcp = create_mcp_server(db_path) - - mock_doc = Document(content="test", uri="https://example.com") - mock_doc.id = "doc456" - - with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class: - mock_rag = AsyncMock() - mock_rag.create_document_from_source = AsyncMock(return_value=mock_doc) - mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag) - mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None) - - tools = await mcp.get_tools() - add_url_tool = next( - t for t in tools.values() if t.name == "add_document_from_url" - ) - - result = await add_url_tool.fn(url="https://example.com") - assert result == "doc456" - mock_rag.create_document_from_source.assert_called_once() - - -@pytest.mark.asyncio -async def test_mcp_add_document_from_text(): - """Test add_document_from_text tool is properly wired.""" - with tempfile.TemporaryDirectory() as temp_dir: - db_path = Path(temp_dir) / "test.lancedb" - mcp = create_mcp_server(db_path) - - mock_doc = Document(content="test content", uri="text://test") - mock_doc.id = "doc789" - - with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class: - mock_rag = AsyncMock() - mock_rag.create_document = AsyncMock(return_value=mock_doc) - mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag) - mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None) - - tools = await mcp.get_tools() - add_text_tool = next( - t for t in tools.values() if t.name == "add_document_from_text" - ) - - result = await add_text_tool.fn(content="test content", uri="text://test") - - assert result == "doc789" - mock_rag.create_document.assert_called_once() - - -@pytest.mark.asyncio -async def test_mcp_search_documents(): - """Test search_documents tool is properly wired.""" - with tempfile.TemporaryDirectory() as temp_dir: - db_path = Path(temp_dir) / "test.lancedb" - mcp = create_mcp_server(db_path) - - from haiku.rag.store.models import SearchResult - - mock_results = [ - SearchResult(content="Result 1", score=0.9, document_id="doc1"), - SearchResult(content="Result 2", score=0.8, document_id="doc2"), - ] - - with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class: - mock_rag = AsyncMock() - mock_rag.search = AsyncMock(return_value=mock_results) - mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag) - mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None) - - tools = await mcp.get_tools() - search_tool = next( - t for t in tools.values() if t.name == "search_documents" - ) - - result = await search_tool.fn(query="test query", limit=5) - assert len(result) == 2 - assert result[0].document_id == "doc1" - assert result[0].content == "Result 1" - assert result[0].score == 0.9 - assert result[1].document_id == "doc2" - mock_rag.search.assert_called_once_with("test query", limit=5) - - -@pytest.mark.asyncio -async def test_mcp_get_document(): - """Test get_document tool is properly wired.""" - with tempfile.TemporaryDirectory() as temp_dir: - db_path = Path(temp_dir) / "test.lancedb" - mcp = create_mcp_server(db_path) - - from datetime import UTC, datetime - - mock_doc = Document(content="test", uri="file:///test.txt", title="Test Doc") - mock_doc.id = "doc123" - mock_doc.created_at = datetime(2024, 1, 1, tzinfo=UTC) - mock_doc.updated_at = datetime(2024, 1, 2, tzinfo=UTC) - - with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class: - mock_rag = AsyncMock() - mock_rag.get_document_by_id = AsyncMock(return_value=mock_doc) - mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag) - mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None) - - tools = await mcp.get_tools() - get_tool = next(t for t in tools.values() if t.name == "get_document") - - result = await get_tool.fn(document_id="doc123") - assert result is not None - assert result.id == "doc123" - assert result.content == "test" - assert result.title == "Test Doc" - mock_rag.get_document_by_id.assert_called_once_with("doc123") - - -@pytest.mark.asyncio -async def test_mcp_list_documents(): - """Test list_documents tool is properly wired.""" - with tempfile.TemporaryDirectory() as temp_dir: - db_path = Path(temp_dir) / "test.lancedb" - mcp = create_mcp_server(db_path) - - from datetime import UTC, datetime - - mock_doc1 = Document(content="test1", uri="file:///test1.txt") - mock_doc1.id = "doc1" - mock_doc1.created_at = datetime(2024, 1, 1, tzinfo=UTC) - mock_doc1.updated_at = datetime(2024, 1, 1, tzinfo=UTC) - - mock_doc2 = Document(content="test2", uri="file:///test2.txt") - mock_doc2.id = "doc2" - mock_doc2.created_at = datetime(2024, 1, 2, tzinfo=UTC) - mock_doc2.updated_at = datetime(2024, 1, 2, tzinfo=UTC) - - with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class: - mock_rag = AsyncMock() - mock_rag.list_documents = AsyncMock(return_value=[mock_doc1, mock_doc2]) - mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag) - mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None) - - tools = await mcp.get_tools() - list_tool = next(t for t in tools.values() if t.name == "list_documents") - - result = await list_tool.fn(limit=10, offset=0) - assert len(result) == 2 - assert result[0].id == "doc1" - assert result[1].id == "doc2" - mock_rag.list_documents.assert_called_once_with(10, 0, None) - - -@pytest.mark.asyncio -async def test_mcp_delete_document(): - """Test delete_document tool is properly wired.""" - with tempfile.TemporaryDirectory() as temp_dir: - db_path = Path(temp_dir) / "test.lancedb" - mcp = create_mcp_server(db_path) - - with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class: - mock_rag = AsyncMock() - mock_rag.delete_document = AsyncMock(return_value=True) - mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag) - mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None) - - tools = await mcp.get_tools() - delete_tool = next(t for t in tools.values() if t.name == "delete_document") - - result = await delete_tool.fn(document_id="doc123") - assert result is True - mock_rag.delete_document.assert_called_once_with("doc123") - - -@pytest.mark.asyncio -async def test_mcp_ask_question_deep(): - """Test ask_question tool with deep=True uses research graph.""" - with tempfile.TemporaryDirectory() as temp_dir: - db_path = Path(temp_dir) / "test.lancedb" - mcp = create_mcp_server(db_path) - - with ( - patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class, - patch( - "haiku.rag.agents.research.graph.build_research_graph" - ) as mock_graph_builder, - ): - mock_rag = AsyncMock() - mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag) - mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None) - - mock_graph = AsyncMock() - mock_result = AsyncMock() - mock_result.executive_summary = "Deep answer from research" - mock_graph.run = AsyncMock(return_value=mock_result) - mock_graph_builder.return_value = mock_graph - - tools = await mcp.get_tools() - ask_tool = next(t for t in tools.values() if t.name == "ask_question") - - # cite=False to avoid citation formatting in output - result = await ask_tool.fn(question="Deep question?", cite=False, deep=True) - - assert result == "Deep answer from research" - mock_graph.run.assert_called_once() - - -@pytest.mark.asyncio -async def test_mcp_research_question(): - """Test research_question tool is properly wired.""" - with tempfile.TemporaryDirectory() as temp_dir: - db_path = Path(temp_dir) / "test.lancedb" - mcp = create_mcp_server(db_path) - - mock_report = ResearchReport( - title="Research Title", - executive_summary="Summary", - main_findings=["Finding 1"], - conclusions=["Conclusion 1"], - recommendations=["Recommendation 1"], - sources_summary="Sources used", - ) - - with ( - patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class, - patch( - "haiku.rag.agents.research.graph.build_research_graph" - ) as mock_graph_builder, - ): - mock_rag = AsyncMock() - mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag) - mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None) - - mock_graph = AsyncMock() - mock_graph.run = AsyncMock(return_value=mock_report) - mock_graph_builder.return_value = mock_graph - - tools = await mcp.get_tools() - research_tool = next( - t for t in tools.values() if t.name == "research_question" - ) - - result = await research_tool.fn( - question="Research question?", - ) - - assert result is not None - assert result.title == "Research Title" - assert result.executive_summary == "Summary" - mock_graph.run.assert_called_once() - - -@pytest.mark.asyncio -async def test_mcp_rlm_question(): - """Test rlm_question tool is properly wired.""" - with tempfile.TemporaryDirectory() as temp_dir: - db_path = Path(temp_dir) / "test.lancedb" - mcp = create_mcp_server(db_path) - - from haiku.rag.agents.rlm.models import RLMResult - - mock_result = RLMResult( - answer="The total is 42.", - program="result = sum(values)\nprint(result)", - ) - - with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class: - mock_rag = AsyncMock() - mock_rag.rlm = AsyncMock(return_value=mock_result) - mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag) - mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None) - - tools = await mcp.get_tools() - rlm_tool = next(t for t in tools.values() if t.name == "rlm_question") - - result = await rlm_tool.fn(question="What is the total?") - - assert result == "The total is 42." - mock_rag.rlm.assert_called_once_with( - "What is the total?", documents=None, filter=None - ) - - -@pytest.mark.asyncio -async def test_mcp_rlm_question_with_document_and_filter(): - """Test rlm_question tool with document and filter parameters.""" - with tempfile.TemporaryDirectory() as temp_dir: - db_path = Path(temp_dir) / "test.lancedb" - mcp = create_mcp_server(db_path) - - from haiku.rag.agents.rlm.models import RLMResult - - mock_result = RLMResult( - answer="Filtered answer", - program="print('filtered')", - ) - - with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class: - mock_rag = AsyncMock() - mock_rag.rlm = AsyncMock(return_value=mock_result) - mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag) - mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None) - - tools = await mcp.get_tools() - rlm_tool = next(t for t in tools.values() if t.name == "rlm_question") - - result = await rlm_tool.fn( - question="Analyze this", - document="doc-123", - filter="uri LIKE '%test%'", - ) - - assert result == "Filtered answer" - mock_rag.rlm.assert_called_once_with( - "Analyze this", - documents=["doc-123"], - filter="uri LIKE '%test%'", - ) diff --git a/tests/test_reranker.py b/tests/test_reranker.py index 8d1d3f9f..6cd8a607 100644 --- a/tests/test_reranker.py +++ b/tests/test_reranker.py @@ -108,52 +108,14 @@ class TestGetReranker: result = get_reranker(config) assert result is None - def test_mxbai_provider(self): - try: - from haiku.rag.reranking.mxbai import MxBAIReranker - - config = AppConfig( - reranking=RerankingConfig( - model=ModelConfig( - provider="mxbai", name="mixedbread-ai/mxbai-rerank-base-v2" - ) - ) - ) - result = get_reranker(config) - assert isinstance(result, MxBAIReranker) - except ImportError: - pytest.skip("MxBAI package not installed") - - def test_cohere_provider(self): - try: - from haiku.rag.reranking.cohere import CohereReranker - - config = AppConfig( - reranking=RerankingConfig( - model=ModelConfig(provider="cohere", name="rerank-v3.5") - ) - ) - result = get_reranker(config) - assert isinstance(result, CohereReranker) - except ImportError: - pytest.skip("Cohere package not installed") - - def test_vllm_provider_with_base_url(self): - from haiku.rag.reranking.vllm import VLLMReranker - + def test_unknown_provider_returns_none(self): config = AppConfig( reranking=RerankingConfig( - model=ModelConfig( - provider="vllm", - name="BAAI/bge-reranker-v2-m3", - base_url="http://localhost:8000", - ) + model=ModelConfig(provider="unknown_provider", name="some-model") ) ) result = get_reranker(config) - assert isinstance(result, VLLMReranker) - assert result._model == "BAAI/bge-reranker-v2-m3" - assert result._base_url == "http://localhost:8000" + assert result is None def test_vllm_provider_without_base_url_raises_error(self): config = AppConfig( @@ -164,75 +126,115 @@ class TestGetReranker: with pytest.raises(ValueError, match="vLLM reranker requires base_url"): get_reranker(config) - def test_zeroentropy_provider(self): - try: - from haiku.rag.reranking.zeroentropy import ZeroEntropyReranker + @pytest.mark.parametrize( + "provider, model_name, class_module, class_name, extra_model_kwargs, expected_attrs, env_vars", + [ + ( + "mxbai", + "mixedbread-ai/mxbai-rerank-base-v2", + "haiku.rag.reranking.mxbai", + "MxBAIReranker", + {}, + {}, + {}, + ), + ( + "cohere", + "rerank-v3.5", + "haiku.rag.reranking.cohere", + "CohereReranker", + {}, + {}, + {}, + ), + ( + "vllm", + "BAAI/bge-reranker-v2-m3", + "haiku.rag.reranking.vllm", + "VLLMReranker", + {"base_url": "http://localhost:8000"}, + { + "_model": "BAAI/bge-reranker-v2-m3", + "_base_url": "http://localhost:8000", + }, + {}, + ), + ( + "zeroentropy", + "zerank-1", + "haiku.rag.reranking.zeroentropy", + "ZeroEntropyReranker", + {}, + {"_model": "zerank-1"}, + {}, + ), + ( + "zeroentropy", + "", + "haiku.rag.reranking.zeroentropy", + "ZeroEntropyReranker", + {}, + {"_model": "zerank-1"}, + {}, + ), + ( + "jina", + "jina-reranker-v3", + "haiku.rag.reranking.jina", + "JinaReranker", + {}, + {"_model": "jina-reranker-v3"}, + {"JINA_API_KEY": "test-api-key"}, + ), + ( + "jina-local", + "jinaai/jina-reranker-v3", + "haiku.rag.reranking.jina_local", + "JinaLocalReranker", + {}, + {"_model": "jinaai/jina-reranker-v3"}, + {}, + ), + ], + ids=[ + "mxbai", + "cohere", + "vllm", + "zeroentropy", + "zeroentropy-default", + "jina", + "jina-local", + ], + ) + def test_provider( + self, + provider, + model_name, + class_module, + class_name, + extra_model_kwargs, + expected_attrs, + env_vars, + monkeypatch, + ): + mod = pytest.importorskip(class_module) + expected_class = getattr(mod, class_name) - config = AppConfig( - reranking=RerankingConfig( - model=ModelConfig(provider="zeroentropy", name="zerank-1") - ) - ) - result = get_reranker(config) - assert isinstance(result, ZeroEntropyReranker) - assert result._model == "zerank-1" - except ImportError: - pytest.skip("Zero Entropy package not installed") - - def test_zeroentropy_provider_default_model(self): - try: - from haiku.rag.reranking.zeroentropy import ZeroEntropyReranker - - config = AppConfig( - reranking=RerankingConfig( - model=ModelConfig(provider="zeroentropy", name="") - ) - ) - result = get_reranker(config) - assert isinstance(result, ZeroEntropyReranker) - assert result._model == "zerank-1" - except ImportError: - pytest.skip("Zero Entropy package not installed") - - def test_unknown_provider_returns_none(self): - config = AppConfig( - reranking=RerankingConfig( - model=ModelConfig(provider="unknown_provider", name="some-model") - ) - ) - result = get_reranker(config) - assert result is None - - def test_jina_provider(self, monkeypatch): - monkeypatch.setenv("JINA_API_KEY", "test-api-key") - - from haiku.rag.reranking.jina import JinaReranker + for key, value in env_vars.items(): + monkeypatch.setenv(key, value) config = AppConfig( reranking=RerankingConfig( - model=ModelConfig(provider="jina", name="jina-reranker-v3") + model=ModelConfig( + provider=provider, name=model_name, **extra_model_kwargs + ) ) ) result = get_reranker(config) - assert isinstance(result, JinaReranker) - assert result._model == "jina-reranker-v3" + assert isinstance(result, expected_class) - def test_jina_local_provider(self): - try: - from haiku.rag.reranking.jina_local import JinaLocalReranker - - config = AppConfig( - reranking=RerankingConfig( - model=ModelConfig( - provider="jina-local", name="jinaai/jina-reranker-v3" - ) - ) - ) - result = get_reranker(config) - assert isinstance(result, JinaLocalReranker) - assert result._model == "jinaai/jina-reranker-v3" - except ImportError: - pytest.skip("Jina local dependencies not installed") + for attr, value in expected_attrs.items(): + assert getattr(result, attr) == value def test_jina_reranker_missing_api_key(monkeypatch):