File size: 6,626 Bytes
243b4bc 24efe34 243b4bc 24efe34 243b4bc 24efe34 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | import pytest
import json
from unittest.mock import MagicMock, AsyncMock, patch
import torch
from langchain_core.messages import AIMessage
from src.pipelines.graph_runner_pipeline import GraphRunnerPipeline
from src.models.agent_models import Orchastrator_Output
@pytest.mark.asyncio
async def test_graph_runner_pipeline_casual_chat():
"""Test that GraphRunnerPipeline correctly routes to casual chat and returns response."""
# Configure mock LLM response
mock_structured_output = AsyncMock()
mock_structured_output.ainvoke.return_value = Orchastrator_Output(
redirect_to="chat_node",
querie=""
)
mock_chat_output = AsyncMock()
mock_chat_output.ainvoke.return_value = AIMessage(content="Hello! How can I help you today?")
mock_llm = MagicMock()
mock_llm.with_structured_output.return_value = mock_structured_output
mock_llm.bind_tools.return_value = mock_chat_output
# Patch the llm inside agents_nodes directly
with patch("src.nodes.agents_nodes.llm", mock_llm):
pipeline = GraphRunnerPipeline()
chunks = []
async for chunk in pipeline.initiate(thread_id="test_thread_1", query="Hello"):
chunks.append(chunk)
assert len(chunks) > 0
first_chunk_str = chunks[0].replace("data: ", "").strip()
first_chunk = json.loads(first_chunk_str)
# Verify execution path
assert "orchestrator" in first_chunk or "chat" in first_chunk
mock_llm.with_structured_output.assert_called_once_with(Orchastrator_Output)
mock_structured_output.ainvoke.assert_called_once()
mock_llm.bind_tools.assert_called_once()
mock_chat_output.ainvoke.assert_called_once()
@pytest.mark.skip(
reason=(
"Written against the deprecated retreiver_node which used "
"vectorizer.get_similar_data() + MyModel.predict_emb(). "
"The active node is now retriever_node_v2 which calls "
"vectorizer.invoke() directly. Update this test for the new node."
)
)
@pytest.mark.asyncio
async def test_graph_runner_pipeline_retrieval():
"""Test that GraphRunnerPipeline correctly routes to retriever, queries vector db, and returns recommendations."""
# Configure mock LLM response for orchestrator (routes to retreiver)
mock_structured_output = AsyncMock()
mock_structured_output.ainvoke.return_value = Orchastrator_Output(
redirect_to="retreiver_node",
querie="blue jeans"
)
# Configure mock LLM response for chat (recommendation response)
mock_chat_output = AsyncMock()
mock_chat_output.ainvoke.return_value = AIMessage(content="I found these Sleek Blue Jeans for you.")
mock_llm = MagicMock()
mock_llm.with_structured_output.return_value = mock_structured_output
mock_llm.bind_tools.return_value = mock_chat_output
# Configure MyModel instance predict_emb mock
mock_mymodel_inst = MagicMock()
mock_mymodel_inst.predict_emb.return_value = torch.zeros((1, 512))
mock_mymodel_inst.load_model = MagicMock()
# Configure Vectorizer (vector database) mock query result
mock_vectorizer_inst = MagicMock()
mock_match_1 = MagicMock()
mock_match_1.id = "101"
mock_match_1.score = 0.95
mock_match_1.metadata = {"name": "Sleek Blue Jeans", "price": 1200.0}
mock_vectorizer_inst.get_similar_data = AsyncMock(return_value={
"matches": [mock_match_1]
})
# Patch all the dependencies in agents_nodes directly
with patch("src.nodes.agents_nodes.llm", mock_llm), \
patch("src.nodes.agents_nodes.my_model", return_value=mock_mymodel_inst), \
patch("src.nodes.agents_nodes.vectorizer", return_value=mock_vectorizer_inst), \
patch("src.nodes.agents_nodes._get_image_feat", return_value=torch.zeros((1, 2048))), \
patch("src.nodes.agents_nodes._get_text_feat", return_value=torch.zeros((1, 768))):
pipeline = GraphRunnerPipeline()
chunks = []
async for chunk in pipeline.initiate(thread_id="test_thread_2", query="Show me blue jeans"):
chunks.append(chunk)
assert len(chunks) > 0
# Verify vector db call and model predictions
mock_vectorizer_inst.get_similar_data.assert_called_once()
mock_mymodel_inst.predict_emb.assert_called_once()
mock_chat_output.ainvoke.assert_called_once()
@pytest.mark.asyncio
async def test_graph_runner_pipeline_retrieval_v2():
"""Test that GraphRunnerPipeline correctly routes to retriever_node_v2,
calls vectorizer.invoke(), and returns product recommendations."""
# ── Orchestrator: route to retriever ────────────────────────────────
mock_structured_output = AsyncMock()
mock_structured_output.ainvoke.return_value = Orchastrator_Output(
redirect_to="retreiver_node",
querie="blue jeans"
)
# ── Chat: final recommendation response ─────────────────────────────
mock_chat_output = AsyncMock()
mock_chat_output.ainvoke.return_value = AIMessage(
content="I found these Sleek Blue Jeans for you."
)
mock_llm = MagicMock()
mock_llm.with_structured_output.return_value = mock_structured_output
mock_llm.bind_tools.return_value = mock_chat_output
# ── Vectorizer mock: invoke() returns list of plain dicts (retriever_node_v2 format) ──
mock_vectorizer_inst = MagicMock()
mock_vectorizer_inst.invoke = AsyncMock(return_value=[
{
"id": "101",
"score": 0.95,
"metadata": {"name": "Sleek Blue Jeans", "price": 1200.0}
}
])
with patch("src.nodes.agents_nodes.llm", mock_llm), \
patch("src.nodes.agents_nodes.vectorizer", return_value=mock_vectorizer_inst), \
patch("src.nodes.agents_nodes._get_text_feat", return_value=torch.zeros((1, 768))), \
patch("src.nodes.agents_nodes._get_image_feat", return_value=torch.zeros((1, 2048))):
pipeline = GraphRunnerPipeline()
chunks = []
async for chunk in pipeline.initiate(thread_id="test_thread_v2", query="Show me blue jeans"):
chunks.append(chunk)
assert len(chunks) > 0
# retriever_node_v2 must have called invoke() once
mock_vectorizer_inst.invoke.assert_called_once()
# chat LLM must have been invoked for the final response
mock_chat_output.ainvoke.assert_called_once()
|