From 5cf2de5f74678409e6bc87fa9c217288cfcbaef7 Mon Sep 17 00:00:00 2001 From: krypticmouse Date: Thu, 26 Mar 2026 03:11:53 +0000 Subject: [PATCH] feat: add end-to-end integration test for Deep Research pipeline Tests the full path from multi-source Document ingestion through IngestionPipeline -> KnowledgeStore -> TwoStageRetriever -> KnowledgeSearchTool -> DeepResearchAgent to a cited report, and verifies cross-platform retrieval returns results from >= 2 sources. Co-Authored-By: Claude Sonnet 4.6 --- .../agents/test_deep_research_integration.py | 221 ++++++++++++++++++ 1 file changed, 221 insertions(+) create mode 100644 tests/agents/test_deep_research_integration.py diff --git a/tests/agents/test_deep_research_integration.py b/tests/agents/test_deep_research_integration.py new file mode 100644 index 00000000..0d1b9f9d --- /dev/null +++ b/tests/agents/test_deep_research_integration.py @@ -0,0 +1,221 @@ +"""End-to-end integration tests for the Deep Research pipeline. + +Covers the full path: + multi-source Documents + -> IngestionPipeline + -> KnowledgeStore + -> TwoStageRetriever + -> KnowledgeSearchTool + -> DeepResearchAgent + -> cited report +""" + +from __future__ import annotations + +import json +from pathlib import Path +from unittest.mock import MagicMock + +from openjarvis.agents.deep_research import DeepResearchAgent +from openjarvis.connectors._stubs import Document +from openjarvis.connectors.pipeline import IngestionPipeline +from openjarvis.connectors.retriever import TwoStageRetriever +from openjarvis.connectors.store import KnowledgeStore +from openjarvis.tools.knowledge_search import KnowledgeSearchTool + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +_K8S_DOCS = [ + Document( + doc_id="slack-001", + source="slack", + doc_type="message", + content=( + "Hey team, we should migrate to Kubernetes for better" + " orchestration. The new cluster is ready." + ), + title="K8s migration proposal", + author="sarah", + ), + Document( + doc_id="gmail-001", + source="gmail", + doc_type="email", + content=( + "Cost analysis for the Kubernetes migration shows a 40% reduction" + " in infrastructure spend over 12 months." + ), + title="Cost analysis for K8s migration", + author="mike", + ), + Document( + doc_id="gdrive-001", + source="gdrive", + doc_type="document", + content=( + "Kubernetes Migration Proposal v2: This document outlines the" + " phased approach for moving all services to the new K8s cluster." + ), + title="K8s Migration Proposal v2", + author="sarah", + ), + Document( + doc_id="gcalendar-001", + source="gcalendar", + doc_type="event", + content=( + "Infrastructure Sync meeting to review the Kubernetes rollout" + " timeline and assign owners for each microservice." + ), + title="Infrastructure Sync", + author="", + ), + Document( + doc_id="granola-001", + source="granola", + doc_type="meeting_notes", + content=( + "Meeting notes: Sarah presented the Kubernetes migration plan." + " Action items: Mike to run cost analysis, team to review" + " proposal doc by Friday." + ), + title="Infra meeting notes", + author="sarah", + ), +] + + +def _make_engine_response(content, tool_calls=None): + result = { + "content": content, + "usage": { + "prompt_tokens": 50, + "completion_tokens": 100, + "total_tokens": 150, + }, + "model": "test-model", + "finish_reason": "stop", + } + if tool_calls: + result["tool_calls"] = tool_calls + result["finish_reason"] = "tool_calls" + return result + + +def _build_populated_store(tmp_path: Path) -> KnowledgeStore: + """Return a KnowledgeStore populated with the 5 test documents.""" + store = KnowledgeStore(db_path=str(tmp_path / "integration_test.db")) + pipeline = IngestionPipeline(store) + chunks_stored = pipeline.ingest(_K8S_DOCS) + assert chunks_stored >= 5, f"Expected >= 5 chunks, got {chunks_stored}" + return store + + +# --------------------------------------------------------------------------- +# Test 1 — Full research pipeline +# --------------------------------------------------------------------------- + + +def test_full_research_pipeline(tmp_path): + """Full path: IngestionPipeline -> KnowledgeStore -> TwoStageRetriever + -> KnowledgeSearchTool -> DeepResearchAgent -> cited report. + """ + # 1. Populate store via pipeline + store = _build_populated_store(tmp_path) + + # 2. Two-stage retriever wrapping the store (no reranker = BM25 only) + retriever = TwoStageRetriever(store) + + # 3. Knowledge search tool with both store and retriever + ks_tool = KnowledgeSearchTool(store=store, retriever=retriever) + + # 4. Mock engine: first call returns tool_call, second returns final answer + mock_engine = MagicMock() + mock_engine.engine_id = "mock" + mock_engine.health.return_value = True + + tool_call_response = _make_engine_response( + "", + tool_calls=[ + { + "id": "call_k8s_1", + "type": "function", + "function": { + "name": "knowledge_search", + "arguments": json.dumps({"query": "Kubernetes migration"}), + }, + } + ], + ) + final_response = _make_engine_response( + "Based on my research across Slack, email, and documents, the" + " Kubernetes migration was proposed by Sarah and supported by a cost" + " analysis from Mike. [slack] K8s migration proposal -- sarah\n" + "[gmail] Cost analysis for K8s migration -- mike\n" + "[gdrive] K8s Migration Proposal v2 -- sarah\n" + "\n## Sources\n" + "- [slack] K8s migration proposal -- sarah\n" + "- [gmail] Cost analysis for K8s migration -- mike\n" + "- [gdrive] K8s Migration Proposal v2 -- sarah" + ) + mock_engine.generate.side_effect = [tool_call_response, final_response] + + # 5. Create agent with tool + agent = DeepResearchAgent(mock_engine, "test-model", tools=[ks_tool]) + + # 6. Run the agent + result = agent.run("What is the status of the Kubernetes migration?") + + # 7. Assertions + assert result.content, "Result should have non-empty content" + assert "Kubernetes" in result.content, "Result should mention Kubernetes" + assert result.turns >= 1, f"Expected at least 1 turn, got {result.turns}" + + # At least one successful knowledge_search tool result + ks_results = [ + tr + for tr in result.tool_results + if tr.tool_name == "knowledge_search" and tr.success + ] + assert len(ks_results) >= 1, ( + f"Expected at least 1 successful knowledge_search call, " + f"got {len(ks_results)}: {result.tool_results}" + ) + + +# --------------------------------------------------------------------------- +# Test 2 — Cross-platform search finds data from multiple sources +# --------------------------------------------------------------------------- + + +def test_search_finds_cross_platform_data(tmp_path): + """KnowledgeSearchTool via TwoStageRetriever returns results from + at least 2 different sources when searching for 'Kubernetes migration'. + """ + # 1. Populate the same store + store = _build_populated_store(tmp_path) + + # 2. Build search tool with retriever + retriever = TwoStageRetriever(store) + ks_tool = KnowledgeSearchTool(store=store, retriever=retriever) + + # 3. Execute the search directly (no agent) + tool_result = ks_tool.execute(query="Kubernetes migration", top_k=10) + + assert tool_result.success, f"Search failed: {tool_result.content}" + assert tool_result.content, "Search returned empty content" + + # 4. Identify which sources appear in the formatted output + content = tool_result.content + sources_found = { + label + for label in ("slack", "gmail", "gdrive", "granola", "gcalendar") + if f"[{label}]" in content + } + + assert len(sources_found) >= 2, ( + f"Expected results from at least 2 different sources, " + f"but only found: {sources_found}\n\nFull output:\n{content}" + )