From 004cee8bc8521951a4e2e315591bb4e663f295e4 Mon Sep 17 00:00:00 2001 From: nuwangeek Date: Mon, 22 Jun 2026 15:35:29 +0530 Subject: [PATCH] unit tests for api tool calling --- tests/test_agentic_loop.py | 138 +++++++- tests/test_api_caller.py | 6 +- tests/test_api_response_formatter.py | 2 +- tests/test_api_semantic_searcher.py | 65 +++- tests/test_api_tool_session_store.py | 46 ++- tests/test_api_tool_workflow.py | 34 +- tests/test_api_tool_workflow_integration.py | 340 +++++++++++++++++++- tests/test_atc_cache.py | 12 +- tests/test_atc_cache_store.py | 104 ++---- tests/test_direct_step_executor.py | 4 +- tests/test_follow_up_detector.py | 2 +- tests/test_multi_api_caller.py | 8 +- tests/test_multi_response_formatter.py | 2 +- tests/test_param_extractor.py | 2 +- tests/test_qdrant_manager.py | 2 - tests/test_tool_classifier.py | 282 ++++++++++++++++ 16 files changed, 903 insertions(+), 146 deletions(-) diff --git a/tests/test_agentic_loop.py b/tests/test_agentic_loop.py index b8903788..d52fb3ce 100644 --- a/tests/test_agentic_loop.py +++ b/tests/test_agentic_loop.py @@ -5,9 +5,9 @@ import pytest -from src.tool_classifier.agentic_loop import AgenticLoop -from src.tool_classifier.enums import AgenticLoopStatus -from src.tool_classifier.param_extractor import ParamExtractionResult +from tool_classifier.agentic_loop import AgenticLoop +from tool_classifier.enums import AgenticLoopStatus +from tool_classifier.param_extractor import ParamExtractionResult # --------------------------------------------------------------------------- @@ -540,6 +540,7 @@ async def fake_to_thread(fn: Any, *args: Any, **kwargs: Any) -> Any: _HISTORY, {"validFrom": "2026-01-01"}, "en", + 1, ) @@ -963,3 +964,134 @@ async def test_user_exit_during_stream_returns_empty_tokens(self) -> None: assert tokens == [] # Collected params returned unchanged on exit assert result.collected_params == {"validFrom": "2026-01-01"} + + +# --------------------------------------------------------------------------- +# seeded_params — L2 param_update pre-population at turn 0 +# --------------------------------------------------------------------------- + + +class TestSeededParamsTurn0: + """Verify that seeded_params from L2 follow-up routing are merged into + collected_params at turn 0 only, with collected_params taking priority.""" + + @pytest.mark.asyncio + async def test_seeded_params_merged_at_turn_0(self) -> None: + """seeded_params are prepended to collected_params when turn_count=0.""" + # Extractor returns only validFrom as newly extracted; countryIsoCode comes + # from seeded_params. + extractor_mock = _make_extractor_mock( + _extraction( + {"validFrom": "2026-01-01"}, + [], # nothing missing — both params will be present after seed merge + "none", + ) + ) + loop = _make_loop(extractor_mock) + + result = await loop.run_turn( + chat_id=_CHAT_ID, + user_message="January 2026", + conversation_history=[], + params_schema=_SCHEMA_TWO_REQUIRED, + collected_params={}, + turn_count=0, + seeded_params={"countryIsoCode": "EE"}, + ) + + # Both params present → COMPLETED + assert result.status == AgenticLoopStatus.COMPLETED + assert result.collected_params.get("countryIsoCode") == "EE" + assert result.collected_params.get("validFrom") == "2026-01-01" + + @pytest.mark.asyncio + async def test_collected_params_override_seeded_params(self) -> None: + """collected_params values beat seeded_params when the key overlaps.""" + extractor_mock = _make_extractor_mock( + _extraction( + {"validFrom": "2026-06-01"}, + [], + "none", + ) + ) + loop = _make_loop(extractor_mock) + + result = await loop.run_turn( + chat_id=_CHAT_ID, + user_message="June 2026", + conversation_history=[], + params_schema=_SCHEMA_TWO_REQUIRED, + collected_params={"countryIsoCode": "LV"}, # explicit value takes priority + turn_count=0, + seeded_params={"countryIsoCode": "EE"}, # seeded value must be overridden + ) + + assert result.collected_params.get("countryIsoCode") == "LV" + + @pytest.mark.asyncio + async def test_seeded_params_not_applied_on_subsequent_turns(self) -> None: + """seeded_params are ignored when turn_count > 0.""" + extractor_mock = _make_extractor_mock( + _extraction({}, ["countryIsoCode", "validFrom"], "Which country and date?") + ) + loop = _make_loop(extractor_mock) + + result = await loop.run_turn( + chat_id=_CHAT_ID, + user_message="hello", + conversation_history=[], + params_schema=_SCHEMA_TWO_REQUIRED, + collected_params={}, + turn_count=1, # NOT turn 0 → seeded_params must be ignored + seeded_params={"countryIsoCode": "EE", "validFrom": "2026-01-01"}, + ) + + # Even though seeded_params would satisfy all required params, they should + # not be applied after turn 0 → still NEEDS_INPUT + assert result.status == AgenticLoopStatus.NEEDS_INPUT + # seeded values not present in collected_params + assert "countryIsoCode" not in result.collected_params + assert "validFrom" not in result.collected_params + + @pytest.mark.asyncio + async def test_seeded_params_none_does_not_raise(self) -> None: + """Passing seeded_params=None (default) at turn 0 behaves normally.""" + extractor_mock = _make_extractor_mock( + _extraction({}, ["countryIsoCode", "validFrom"], "Which country?") + ) + loop = _make_loop(extractor_mock) + + result = await loop.run_turn( + chat_id=_CHAT_ID, + user_message="hello", + conversation_history=[], + params_schema=_SCHEMA_TWO_REQUIRED, + collected_params={}, + turn_count=0, + seeded_params=None, + ) + + assert result.status == AgenticLoopStatus.NEEDS_INPUT + + @pytest.mark.asyncio + async def test_seeded_params_partial_fill_still_asks_for_missing(self) -> None: + """seeded_params satisfy only one of two required params → still NEEDS_INPUT.""" + extractor_mock = _make_extractor_mock( + _extraction({}, ["validFrom"], "From which date?") + ) + loop = _make_loop(extractor_mock) + + result = await loop.run_turn( + chat_id=_CHAT_ID, + user_message="Estonia", + conversation_history=[], + params_schema=_SCHEMA_TWO_REQUIRED, + collected_params={}, + turn_count=0, + seeded_params={"countryIsoCode": "EE"}, # only one param seeded + ) + + # validFrom still missing → NEEDS_INPUT + assert result.status == AgenticLoopStatus.NEEDS_INPUT + # But seeded countryIsoCode should be present + assert result.collected_params.get("countryIsoCode") == "EE" diff --git a/tests/test_api_caller.py b/tests/test_api_caller.py index a6713fdb..de7a535c 100644 --- a/tests/test_api_caller.py +++ b/tests/test_api_caller.py @@ -7,8 +7,8 @@ import httpx import pytest -from src.tool_classifier.api_caller import APICaller, CircuitBreaker -from src.tool_classifier.constants import ( +from tool_classifier.api_caller import APICaller, CircuitBreaker +from tool_classifier.constants import ( CB_STATE_CLOSED, CB_STATE_HALF_OPEN, CB_STATE_OPEN, @@ -16,7 +16,7 @@ SERVICE_TIMEOUT_MESSAGES, SERVICE_UNAVAILABLE_MESSAGES, ) -from src.tool_classifier.models import APICallResult +from tool_classifier.models import APICallResult # --------------------------------------------------------------------------- diff --git a/tests/test_api_response_formatter.py b/tests/test_api_response_formatter.py index 752c57a2..05fe683f 100644 --- a/tests/test_api_response_formatter.py +++ b/tests/test_api_response_formatter.py @@ -9,7 +9,7 @@ import dspy.streaming import pytest -from src.tool_classifier.api_response_formatter import ( +from tool_classifier.api_response_formatter import ( APIResponseFormatterModule, _FORMATTER_ERROR_MESSAGES, ) diff --git a/tests/test_api_semantic_searcher.py b/tests/test_api_semantic_searcher.py index 15711269..f1b44040 100644 --- a/tests/test_api_semantic_searcher.py +++ b/tests/test_api_semantic_searcher.py @@ -448,13 +448,6 @@ async def test_multiple_medium_triggers_disambiguation(self) -> None: client.get = AsyncMock(return_value=count_resp) client.post = AsyncMock(side_effect=[dense_resp, hybrid_resp]) - mock_disambiguator = MagicMock() - mock_disambiguator.return_value = None # "forward" returns string or None - # Wrap in a module-like object that has a forward() callable via __call__ - mock_disambiguator_module = MagicMock() - mock_disambiguator_module.forward = MagicMock(return_value="ep-holidays") - mock_disambiguator_module.__call__ = MagicMock(return_value="ep-holidays") - # Inject our disambiguator — searcher calls self._disambiguator(query, candidates) # which in turn calls forward() via __call__ async_disambiguator = MagicMock() @@ -507,6 +500,64 @@ async def test_disambiguation_rejects_all_returns_empty(self) -> None: assert results == [] + @pytest.mark.asyncio + async def test_disambiguation_rejects_all_multi_candidates_returns_top_with_hint( + self, + ) -> None: + """Disambiguator returns None for >1 medium candidates → top candidate returned + with multi_intent_hint=True and llm_validated=False so IntentDecomposer gate + can run in the classifier.""" + cos_a = API_TOOL_MIN_THRESHOLD + 0.08 # higher cosine → becomes 'top' + cos_b = API_TOOL_MIN_THRESHOLD + 0.02 + + dense_points = [ + _point({**_EP_HOLIDAYS}, cos_a), + _point({**_EP_WEATHER}, cos_b), + ] + hybrid_points = [ + _point({**_EP_HOLIDAYS}, 0.012), + _point({**_EP_WEATHER}, 0.009), + ] + + dense_resp = _make_qdrant_dense_response(dense_points) + hybrid_resp = _make_qdrant_hybrid_response(hybrid_points) + count_resp = _make_count_response(10) + + client = AsyncMock() + client.get = AsyncMock(return_value=count_resp) + client.post = AsyncMock(side_effect=[dense_resp, hybrid_resp]) + + searcher = _make_searcher(client) + + # asyncio.to_thread is called twice: + # 1st call → _get_query_embedding → must return a valid embedding vector + # 2nd call → disambiguator.forward → must return None ("none" response) + precomputed = [0.1] * 10 + _call_count = 0 + + async def _to_thread_side_effect(fn: Any, *args: Any, **kwargs: Any) -> Any: + nonlocal _call_count + _call_count += 1 + if _call_count == 1: + return precomputed # embedding call + return None # disambiguator call → rejects all candidates + + with patch( + "tool_classifier.api_semantic_searcher.asyncio.to_thread", + side_effect=_to_thread_side_effect, + ): + results = await searcher.search("holidays AND weather") + + # Must return exactly one result — the top cosine candidate + assert len(results) == 1 + top = results[0] + # Top candidate by cosine score is ep-holidays + assert top.endpoint_id == "ep-holidays" + # NOT llm_validated — disambiguator explicitly rejected it + assert top.llm_validated is False + # multi_intent_hint signals the classifier to try IntentDecomposer + assert top.multi_intent_hint is True + class TestSearchBelowThreshold: @pytest.mark.asyncio diff --git a/tests/test_api_tool_session_store.py b/tests/test_api_tool_session_store.py index f6e5fb2d..9b877e0b 100644 --- a/tests/test_api_tool_session_store.py +++ b/tests/test_api_tool_session_store.py @@ -5,8 +5,8 @@ import pytest from pydantic import ValidationError -from src.models.session_models import APIToolSession, EndpointSessionState -from src.utils.api_tool_session_store import ( +from models.session_models import APIToolSession, EndpointSessionState +from utils.api_tool_session_store import ( APIToolSessionStore, _key, require_session_store, @@ -185,7 +185,7 @@ async def test_get_returns_none_when_key_missing(self): redis_mock.get = AsyncMock(return_value=None) with patch( - "src.utils.api_tool_session_store.get_redis_client", return_value=redis_mock + "utils.api_tool_session_store.get_redis_client", return_value=redis_mock ): result = await store.get("missing-chat") @@ -199,7 +199,7 @@ async def test_get_returns_session_when_key_exists(self): redis_mock.get = AsyncMock(return_value=session.model_dump_json()) with patch( - "src.utils.api_tool_session_store.get_redis_client", return_value=redis_mock + "utils.api_tool_session_store.get_redis_client", return_value=redis_mock ): result = await store.get(session.chat_id) @@ -210,9 +210,7 @@ async def test_get_returns_session_when_key_exists(self): @pytest.mark.asyncio async def test_get_returns_none_when_redis_unavailable(self): store = APIToolSessionStore() - with patch( - "src.utils.api_tool_session_store.get_redis_client", return_value=None - ): + with patch("utils.api_tool_session_store.get_redis_client", return_value=None): result = await store.get("any-chat") assert result is None @@ -231,7 +229,7 @@ async def test_save_calls_redis_set_with_correct_key_and_ttl(self): redis_mock = _make_redis_mock() with patch( - "src.utils.api_tool_session_store.get_redis_client", return_value=redis_mock + "utils.api_tool_session_store.get_redis_client", return_value=redis_mock ): await store.save(session) @@ -245,9 +243,7 @@ async def test_save_skips_when_redis_unavailable(self): store = APIToolSessionStore() session = _make_session() - with patch( - "src.utils.api_tool_session_store.get_redis_client", return_value=None - ): + with patch("utils.api_tool_session_store.get_redis_client", return_value=None): # Should not raise await store.save(session) @@ -284,7 +280,7 @@ async def test_update_merges_fields_and_resets_ttl(self): redis_mock.pipeline = MagicMock(return_value=pipe_mock) with patch( - "src.utils.api_tool_session_store.get_redis_client", return_value=redis_mock + "utils.api_tool_session_store.get_redis_client", return_value=redis_mock ): result = await store.update( original.chat_id, @@ -313,7 +309,7 @@ async def test_update_returns_none_when_session_missing(self): redis_mock.pipeline = MagicMock(return_value=pipe_mock) with patch( - "src.utils.api_tool_session_store.get_redis_client", return_value=redis_mock + "utils.api_tool_session_store.get_redis_client", return_value=redis_mock ): result = await store.update("ghost-chat", turn_count=3) @@ -338,7 +334,7 @@ async def test_update_resets_ttl(self): redis_mock.pipeline = MagicMock(return_value=pipe_mock) with patch( - "src.utils.api_tool_session_store.get_redis_client", return_value=redis_mock + "utils.api_tool_session_store.get_redis_client", return_value=redis_mock ): await store.update(session.chat_id, state="ready") @@ -360,7 +356,7 @@ async def test_delete_calls_redis_delete(self): redis_mock = _make_redis_mock() with patch( - "src.utils.api_tool_session_store.get_redis_client", return_value=redis_mock + "utils.api_tool_session_store.get_redis_client", return_value=redis_mock ): await store.delete("chat-to-delete") @@ -369,9 +365,7 @@ async def test_delete_calls_redis_delete(self): @pytest.mark.asyncio async def test_delete_skips_when_redis_unavailable(self): store = APIToolSessionStore() - with patch( - "src.utils.api_tool_session_store.get_redis_client", return_value=None - ): + with patch("utils.api_tool_session_store.get_redis_client", return_value=None): await store.delete("any-chat") # Should not raise @@ -388,7 +382,7 @@ async def test_exists_returns_true_when_key_present(self): redis_mock.exists = AsyncMock(return_value=1) with patch( - "src.utils.api_tool_session_store.get_redis_client", return_value=redis_mock + "utils.api_tool_session_store.get_redis_client", return_value=redis_mock ): result = await store.exists("chat-123") @@ -401,7 +395,7 @@ async def test_exists_returns_false_when_key_absent(self): redis_mock.exists = AsyncMock(return_value=0) with patch( - "src.utils.api_tool_session_store.get_redis_client", return_value=redis_mock + "utils.api_tool_session_store.get_redis_client", return_value=redis_mock ): result = await store.exists("chat-123") @@ -410,9 +404,7 @@ async def test_exists_returns_false_when_key_absent(self): @pytest.mark.asyncio async def test_exists_returns_false_when_redis_unavailable(self): store = APIToolSessionStore() - with patch( - "src.utils.api_tool_session_store.get_redis_client", return_value=None - ): + with patch("utils.api_tool_session_store.get_redis_client", return_value=None): result = await store.exists("any-chat") assert result is False @@ -431,7 +423,7 @@ async def test_get_returns_none_on_redis_error(self): redis_mock.get = AsyncMock(side_effect=ConnectionError("timeout")) with patch( - "src.utils.api_tool_session_store.get_redis_client", return_value=redis_mock + "utils.api_tool_session_store.get_redis_client", return_value=redis_mock ): result = await store.get("chat-xyz") @@ -445,7 +437,7 @@ async def test_save_does_not_raise_on_redis_error(self): redis_mock.set = AsyncMock(side_effect=ConnectionError("timeout")) with patch( - "src.utils.api_tool_session_store.get_redis_client", return_value=redis_mock + "utils.api_tool_session_store.get_redis_client", return_value=redis_mock ): await store.save(session) # Should not raise @@ -456,7 +448,7 @@ async def test_delete_does_not_raise_on_redis_error(self): redis_mock.delete = AsyncMock(side_effect=ConnectionError("timeout")) with patch( - "src.utils.api_tool_session_store.get_redis_client", return_value=redis_mock + "utils.api_tool_session_store.get_redis_client", return_value=redis_mock ): await store.delete("chat-xyz") # Should not raise @@ -467,7 +459,7 @@ async def test_exists_returns_false_on_redis_error(self): redis_mock.exists = AsyncMock(side_effect=ConnectionError("timeout")) with patch( - "src.utils.api_tool_session_store.get_redis_client", return_value=redis_mock + "utils.api_tool_session_store.get_redis_client", return_value=redis_mock ): result = await store.exists("chat-xyz") diff --git a/tests/test_api_tool_workflow.py b/tests/test_api_tool_workflow.py index 6fc640bb..55483e00 100644 --- a/tests/test_api_tool_workflow.py +++ b/tests/test_api_tool_workflow.py @@ -135,6 +135,9 @@ def _format_sse(chat_id: str, content: str) -> str: return f'data: {{"chatId":"{chat_id}","payload":{{"content":"{content}"}}}}\n\n' svc.format_sse = _format_sse + svc.handle_output_guardrails = AsyncMock( + side_effect=lambda _adapter, response, _req, _costs: response + ) return svc @@ -494,19 +497,25 @@ async def _fake_stream(**kwargs: Any) -> AsyncIterator[str]: for token in ["Holiday", " info", " here."]: yield token - executor._formatter.stream_forward = _fake_stream + mock_formatter = MagicMock() + mock_formatter.stream_forward = _fake_stream - frames = [ - frame - async for frame in executor._stream_api_and_format( - chat_id=_CHAT_ID, - endpoint=_ENDPOINT_HOLIDAYS, - collected_params={"countryIsoCode": "EE"}, - user_query="holidays", - detected_language="en", - orchestration_service=svc, - ) - ] + with patch( + "tool_classifier.workflows.api_tool_workflow.APIResponseFormatterModule", + return_value=mock_formatter, + ): + frames = [ + frame + async for frame in executor._stream_api_and_format( + chat_id=_CHAT_ID, + endpoint=_ENDPOINT_HOLIDAYS, + collected_params={"countryIsoCode": "EE"}, + user_query="holidays", + detected_language="en", + orchestration_service=svc, + request=_make_request(), + ) + ] # 3 token frames + 1 END frame assert len(frames) == 4 @@ -536,6 +545,7 @@ async def test_api_failure_streams_error_frame(self) -> None: user_query="holidays", detected_language="et", orchestration_service=svc, + request=_make_request(), ) ] diff --git a/tests/test_api_tool_workflow_integration.py b/tests/test_api_tool_workflow_integration.py index 4d97fd81..7db8a5a9 100644 --- a/tests/test_api_tool_workflow_integration.py +++ b/tests/test_api_tool_workflow_integration.py @@ -6,6 +6,8 @@ Covers: - Phase 2: Full multi-turn workflow, fast-path, streaming, cost tracking - Phase 4: Fallback chain regression tests +- Parallel execution mode (ExecutionMode.PARALLEL end-to-end) +- Test-endpoint session wipe guard """ import json @@ -17,13 +19,14 @@ import pytest from models.request_models import OrchestrationRequest, OrchestrationResponse -from models.session_models import APIToolSession +from models.session_models import APIToolSession, EndpointSessionState from tool_classifier.classifier import ToolClassifier -from tool_classifier.enums import AgenticLoopStatus, WorkflowType +from tool_classifier.enums import AgenticLoopStatus, ExecutionMode, WorkflowType from tool_classifier.models import ( AgenticLoopResult, APICallResult, ClassificationResult, + MultiAPICallResult, ) @@ -127,6 +130,9 @@ async def _mock_rag(**kwargs: Any) -> OrchestrationResponse: svc._execute_orchestration_pipeline = AsyncMock(side_effect=_mock_rag) svc._initialize_service_components = MagicMock(return_value={}) + svc.handle_output_guardrails = AsyncMock( + side_effect=lambda _adapter, response, _req, _costs: response + ) async def _mock_rag_stream(**kwargs: Any) -> AsyncGenerator[str, None]: yield 'data: {"chatId":"test","payload":{"content":"RAG stream answer"}}\n\n' @@ -484,6 +490,9 @@ async def _fake_stream_forward(**kwargs: Any) -> AsyncIterator[str]: for token in ["It is ", "15°C ", "in Tallinn."]: yield token + mock_formatter = MagicMock() + mock_formatter.stream_forward = _fake_stream_forward + with ( patch.object( classifier.api_tool_workflow._api_caller, @@ -491,10 +500,9 @@ async def _fake_stream_forward(**kwargs: Any) -> AsyncIterator[str]: new_callable=AsyncMock, return_value=api_call_result, ), - patch.object( - classifier.api_tool_workflow._formatter, - "stream_forward", - side_effect=_fake_stream_forward, + patch( + "tool_classifier.workflows.api_tool_workflow.APIResponseFormatterModule", + return_value=mock_formatter, ), ): stream = await classifier.route_to_workflow( @@ -896,3 +904,323 @@ def _make_mock_loop( loop = MagicMock() loop.stream_run_turn = AsyncMock(return_value=(result, question_tokens)) return loop + + +# --------------------------------------------------------------------------- +# TestParallelExecutionMode +# --------------------------------------------------------------------------- + + +class TestParallelExecutionMode: + """Full parallel path: classify → ExecutionMode.PARALLEL → MultiEndpointAgenticLoop + → MultiAPICaller → MultiResponseFormatterModule. + + Only DSPy (formatter/extractor), Qdrant HTTP, and Redis are mocked. + """ + + @pytest.mark.asyncio + async def test_parallel_fast_path_no_required_params_both_apis_called( + self, + classifier: ToolClassifier, + mock_session_store: AsyncMock, + ) -> None: + """Both endpoints have no required params → immediate parallel API calls, no session.""" + + # Two endpoints with no required params + ep_weather_no_params = {**_ENDPOINT_WEATHER, "params": []} + ep_holidays_no_params = {**_ENDPOINT_HOLIDAYS, "params": []} + + classification = ClassificationResult( + workflow=WorkflowType.API_TOOL_CALLING, + confidence=0.68, + metadata={ + "execution_mode": ExecutionMode.PARALLEL, + "matched_endpoints": [ep_weather_no_params, ep_holidays_no_params], + }, + ) + request = _make_request("holidays AND weather") + + weather_result = APICallResult( + success=True, status_code=200, response_data={"temp": 22}, error=None + ) + holidays_result = APICallResult( + success=True, + status_code=200, + response_data={"holidays": ["Jõulupüha"]}, + error=None, + ) + multi_result = MultiAPICallResult( + results=[weather_result, holidays_result], + endpoints=[ + {**ep_weather_no_params, "call_params": {}}, + {**ep_holidays_no_params, "call_params": {}}, + ], + ) + + with ( + patch.object( + classifier.api_tool_workflow._api_caller.__class__, + "__init__", + return_value=None, + ), + patch( + "tool_classifier.workflows.api_tool_workflow.MultiAPICaller", + ) as mock_multi_caller_cls, + patch( + "tool_classifier.workflows.api_tool_workflow.asyncio.to_thread", + new_callable=AsyncMock, + return_value="It is 22°C and there are public holidays.", + ), + ): + mock_multi_caller_inst = AsyncMock() + mock_multi_caller_inst.call_all = AsyncMock(return_value=multi_result) + mock_multi_caller_cls.return_value = mock_multi_caller_inst + + response = await classifier.route_to_workflow( + classification=classification, + request=request, + is_streaming=False, + ) + + assert isinstance(response, OrchestrationResponse) + assert response.content != "" + # No session created (fast path) + assert await mock_session_store.get(_CHAT_ID) is None + + @pytest.mark.asyncio + async def test_parallel_session_created_when_params_needed( + self, + classifier: ToolClassifier, + mock_session_store: AsyncMock, + ) -> None: + """When endpoints have required params, a PARALLEL session is created and a + clarifying question is returned for the first turn.""" + + # Both endpoints have required params + classification = ClassificationResult( + workflow=WorkflowType.API_TOOL_CALLING, + confidence=0.68, + metadata={ + "execution_mode": ExecutionMode.PARALLEL, + "matched_endpoints": [_ENDPOINT_HOLIDAYS, _ENDPOINT_WEATHER], + }, + ) + request = _make_request("holidays AND weather please") + + loop_result = AgenticLoopResult( + status=AgenticLoopStatus.NEEDS_INPUT, + collected_params={}, + clarifying_question="Which country for holidays?", + turn_count=1, + ) + + with patch( + "tool_classifier.workflows.api_tool_workflow.MultiEndpointAgenticLoop", + ) as mock_multi_loop_cls: + mock_multi_loop_inst = MagicMock() + mock_multi_loop_inst.stream_run_turn = AsyncMock( + return_value=(loop_result, ["Which", " country", "?"]) + ) + mock_multi_loop_cls.return_value = mock_multi_loop_inst + + response = await classifier.route_to_workflow( + classification=classification, + request=request, + is_streaming=False, + ) + + assert isinstance(response, OrchestrationResponse) + assert response.content != "" + # Session created in Redis with parallel execution mode + session = await mock_session_store.get(_CHAT_ID) + assert session is not None + assert session.execution_mode == ExecutionMode.PARALLEL.value + + @pytest.mark.asyncio + async def test_parallel_max_turns_falls_back_to_rag( + self, + classifier: ToolClassifier, + mock_session_store: AsyncMock, + ) -> None: + """Parallel loop MAX_TURNS_REACHED → session deleted → RAG fallback.""" + # Seed a parallel session + session = APIToolSession( + chat_id=_CHAT_ID, + state="collecting_params", + selected_endpoint=_ENDPOINT_HOLIDAYS, + collected_params={}, + turn_count=6, + max_turns=6, + awaiting_continuation=False, + detected_language="en", + original_query="holidays AND weather", + execution_mode=ExecutionMode.PARALLEL.value, + parallel_endpoints=[ + EndpointSessionState(endpoint=_ENDPOINT_HOLIDAYS), + EndpointSessionState(endpoint=_ENDPOINT_WEATHER), + ], + ) + await mock_session_store.save(session) + + classification = ClassificationResult( + workflow=WorkflowType.API_TOOL_CALLING, + confidence=1.0, + metadata={"reason": "active_session_resume"}, + ) + request = _make_request("I give up") + + max_turns_result = AgenticLoopResult( + status=AgenticLoopStatus.MAX_TURNS_REACHED, + collected_params={}, + clarifying_question="", + turn_count=7, + ) + + with patch( + "tool_classifier.workflows.api_tool_workflow.MultiEndpointAgenticLoop", + ) as mock_multi_loop_cls: + mock_multi_loop_inst = MagicMock() + mock_multi_loop_inst.stream_run_turn = AsyncMock( + return_value=(max_turns_result, []) + ) + mock_multi_loop_cls.return_value = mock_multi_loop_inst + + response = await classifier.route_to_workflow( + classification=classification, + request=request, + is_streaming=False, + ) + + # Session deleted + assert await mock_session_store.get(_CHAT_ID) is None + # Falls back to RAG + assert isinstance(response, OrchestrationResponse) + + @pytest.mark.asyncio + async def test_parallel_streaming_question_yields_sse_frames( + self, + classifier: ToolClassifier, + mock_session_store: AsyncMock, + ) -> None: + """Streaming parallel path: first turn → clarifying question → SSE frames.""" + + classification = ClassificationResult( + workflow=WorkflowType.API_TOOL_CALLING, + confidence=0.68, + metadata={ + "execution_mode": ExecutionMode.PARALLEL, + "matched_endpoints": [_ENDPOINT_HOLIDAYS, _ENDPOINT_WEATHER], + }, + ) + request = _make_request("holidays AND weather") + + loop_result = AgenticLoopResult( + status=AgenticLoopStatus.NEEDS_INPUT, + collected_params={}, + clarifying_question="Which country for holidays?", + turn_count=1, + ) + + with patch( + "tool_classifier.workflows.api_tool_workflow.MultiEndpointAgenticLoop", + ) as mock_multi_loop_cls: + mock_multi_loop_inst = MagicMock() + mock_multi_loop_inst.stream_run_turn = AsyncMock( + return_value=(loop_result, ["Which", " country", "?"]) + ) + mock_multi_loop_cls.return_value = mock_multi_loop_inst + + stream = await classifier.route_to_workflow( + classification=classification, + request=request, + is_streaming=True, + ) + frames = [frame async for frame in stream] + + assert len(frames) >= 1 + for frame in frames: + assert frame.startswith("data: ") or frame.strip() == "" + + +# --------------------------------------------------------------------------- +# TestTestEndpointSessionWipe +# --------------------------------------------------------------------------- + + +class TestTestEndpointSessionWipe: + """Verify that the /orchestrate/test endpoint deletes any stale 'test-session' + in the API tool session store before each request so multi-turn state never + leaks between consecutive test API calls.""" + + @pytest.mark.asyncio + async def test_session_store_delete_called_with_test_session_key(self) -> None: + """The endpoint must call session_store.delete('test-session') on every request + regardless of whether a session exists.""" + from httpx import AsyncClient, ASGITransport + from llm_orchestration_service_api import app + + session_store_mock = AsyncMock() + session_store_mock.delete = AsyncMock(return_value=None) + + # Minimal orchestration service mock that returns a valid response + orch_mock = AsyncMock() + orch_mock.process_orchestration_request = AsyncMock( + return_value=OrchestrationResponse( + chatId="test-session", + llmServiceActive=True, + questionOutOfLLMScope=False, + inputGuardFailed=False, + content="Test answer.", + ) + ) + + app.state.orchestration_service = orch_mock + app.state.session_store = session_store_mock + + async with AsyncClient( + transport=ASGITransport(app=app), base_url="http://test" + ) as client: + await client.post( + "/orchestrate/test", + json={"message": "hello", "environment": "production"}, + ) + + # The endpoint must have called delete("test-session") before processing + session_store_mock.delete.assert_awaited_with("test-session") + + @pytest.mark.asyncio + async def test_stale_session_cleared_before_request_not_after(self) -> None: + """If session_store.delete raises, the endpoint propagates the error as HTTP 500 + because the wipe-guard is not wrapped in try/except.""" + from httpx import AsyncClient, ASGITransport + from llm_orchestration_service_api import app + + session_store_mock = AsyncMock() + session_store_mock.delete = AsyncMock( + side_effect=RuntimeError("Redis unavailable") + ) + + orch_mock = AsyncMock() + orch_mock.process_orchestration_request = AsyncMock( + return_value=OrchestrationResponse( + chatId="test-session", + llmServiceActive=True, + questionOutOfLLMScope=False, + inputGuardFailed=False, + content="Fallback answer.", + ) + ) + + app.state.orchestration_service = orch_mock + app.state.session_store = session_store_mock + + async with AsyncClient( + transport=ASGITransport(app=app), base_url="http://test" + ) as client: + resp = await client.post( + "/orchestrate/test", + json={"message": "hello", "environment": "production"}, + ) + + # delete() raised → the unguarded await bubbles up as HTTP 500 + assert resp.status_code == 500 diff --git a/tests/test_atc_cache.py b/tests/test_atc_cache.py index f8390f5e..50ae95b4 100644 --- a/tests/test_atc_cache.py +++ b/tests/test_atc_cache.py @@ -19,10 +19,18 @@ import pytest from models.request_models import OrchestrationRequest -from models.session_models import APIToolSession, EndpointSessionState, LastCallContext +from models.session_models import ( + APIToolSession, + EndpointSessionState, + LastCallContext, +) from tool_classifier.classifier import ToolClassifier from tool_classifier.enums import AgenticLoopStatus, WorkflowType -from tool_classifier.models import AgenticLoopResult, APICallResult, MultiAPICallResult +from tool_classifier.models import ( + AgenticLoopResult, + APICallResult, + MultiAPICallResult, +) from tool_classifier.workflows.api_tool_workflow import APIToolWorkflowExecutor from utils.atc_cache_store import ATCCacheStore diff --git a/tests/test_atc_cache_store.py b/tests/test_atc_cache_store.py index 28e9532b..38593f47 100644 --- a/tests/test_atc_cache_store.py +++ b/tests/test_atc_cache_store.py @@ -5,14 +5,14 @@ import pytest -from src.models.session_models import LastCallContext -from src.tool_classifier.constants import ( +from models.session_models import LastCallContext +from tool_classifier.constants import ( ATC_CACHE_DEFAULT_TTL_SECONDS, ATC_CACHE_KEY_PREFIX, ATC_LAST_CALL_KEY_PREFIX, ATC_LAST_CALL_TTL_SECONDS, ) -from src.utils.atc_cache_store import ATCCacheStore +from utils.atc_cache_store import ATCCacheStore # --------------------------------------------------------------------------- # Shared fixtures @@ -170,9 +170,7 @@ async def test_returns_deserialised_value_on_hit(self): redis_mock = _make_redis_mock() redis_mock.get = AsyncMock(return_value=json.dumps(RESPONSE)) - with patch( - "src.utils.atc_cache_store.get_redis_client", return_value=redis_mock - ): + with patch("utils.atc_cache_store.get_redis_client", return_value=redis_mock): result = await store.get_l1(CHAT_ID, API_NAME, PARAMS) assert result == RESPONSE @@ -183,9 +181,7 @@ async def test_returns_none_on_miss(self): redis_mock = _make_redis_mock() redis_mock.get = AsyncMock(return_value=None) - with patch( - "src.utils.atc_cache_store.get_redis_client", return_value=redis_mock - ): + with patch("utils.atc_cache_store.get_redis_client", return_value=redis_mock): result = await store.get_l1(CHAT_ID, API_NAME, PARAMS) assert result is None @@ -193,7 +189,7 @@ async def test_returns_none_on_miss(self): @pytest.mark.asyncio async def test_returns_none_when_redis_unavailable(self): store = ATCCacheStore() - with patch("src.utils.atc_cache_store.get_redis_client", return_value=None): + with patch("utils.atc_cache_store.get_redis_client", return_value=None): result = await store.get_l1(CHAT_ID, API_NAME, PARAMS) assert result is None @@ -204,9 +200,7 @@ async def test_returns_none_on_redis_exception(self): redis_mock = _make_redis_mock() redis_mock.get = AsyncMock(side_effect=RuntimeError("connection lost")) - with patch( - "src.utils.atc_cache_store.get_redis_client", return_value=redis_mock - ): + with patch("utils.atc_cache_store.get_redis_client", return_value=redis_mock): result = await store.get_l1(CHAT_ID, API_NAME, PARAMS) assert result is None @@ -224,9 +218,7 @@ async def selective_get(key): redis_mock = _make_redis_mock() redis_mock.get = AsyncMock(side_effect=selective_get) - with patch( - "src.utils.atc_cache_store.get_redis_client", return_value=redis_mock - ): + with patch("utils.atc_cache_store.get_redis_client", return_value=redis_mock): result = await store.get_l1(CHAT_ID, API_NAME, different_params) assert result is None @@ -244,9 +236,7 @@ async def test_calls_redis_set_with_correct_key_and_value(self): redis_mock = _make_redis_mock() expected_key = ATCCacheStore._l1_key(CHAT_ID, API_NAME, PARAMS) - with patch( - "src.utils.atc_cache_store.get_redis_client", return_value=redis_mock - ): + with patch("utils.atc_cache_store.get_redis_client", return_value=redis_mock): await store.set_l1(CHAT_ID, API_NAME, PARAMS, RESPONSE) redis_mock.set.assert_called_once_with( @@ -260,9 +250,7 @@ async def test_uses_default_ttl_when_not_specified(self): store = ATCCacheStore() redis_mock = _make_redis_mock() - with patch( - "src.utils.atc_cache_store.get_redis_client", return_value=redis_mock - ): + with patch("utils.atc_cache_store.get_redis_client", return_value=redis_mock): await store.set_l1(CHAT_ID, API_NAME, PARAMS, RESPONSE) assert redis_mock.set.call_args[1]["ex"] == ATC_CACHE_DEFAULT_TTL_SECONDS @@ -272,9 +260,7 @@ async def test_uses_custom_ttl_when_provided(self): store = ATCCacheStore() redis_mock = _make_redis_mock() - with patch( - "src.utils.atc_cache_store.get_redis_client", return_value=redis_mock - ): + with patch("utils.atc_cache_store.get_redis_client", return_value=redis_mock): await store.set_l1(CHAT_ID, API_NAME, PARAMS, RESPONSE, ttl=120) assert redis_mock.set.call_args[1]["ex"] == 120 @@ -282,7 +268,7 @@ async def test_uses_custom_ttl_when_provided(self): @pytest.mark.asyncio async def test_no_op_when_redis_unavailable(self): store = ATCCacheStore() - with patch("src.utils.atc_cache_store.get_redis_client", return_value=None): + with patch("utils.atc_cache_store.get_redis_client", return_value=None): await store.set_l1(CHAT_ID, API_NAME, PARAMS, RESPONSE) # must not raise @pytest.mark.asyncio @@ -291,9 +277,7 @@ async def test_no_op_on_redis_exception(self): redis_mock = _make_redis_mock() redis_mock.set = AsyncMock(side_effect=RuntimeError("write failed")) - with patch( - "src.utils.atc_cache_store.get_redis_client", return_value=redis_mock - ): + with patch("utils.atc_cache_store.get_redis_client", return_value=redis_mock): await store.set_l1(CHAT_ID, API_NAME, PARAMS, RESPONSE) # must not raise @@ -308,9 +292,7 @@ async def test_set_then_get_returns_same_response(self): store = ATCCacheStore() _, redis_mock = _fake_redis_store() - with patch( - "src.utils.atc_cache_store.get_redis_client", return_value=redis_mock - ): + with patch("utils.atc_cache_store.get_redis_client", return_value=redis_mock): await store.set_l1(CHAT_ID, API_NAME, PARAMS, RESPONSE) result = await store.get_l1(CHAT_ID, API_NAME, PARAMS) @@ -324,9 +306,7 @@ async def test_string_year_hits_entry_stored_with_int_year(self): params_int = {"country": "EE", "year": 2026} params_str = {"country": "EE", "year": "2026"} - with patch( - "src.utils.atc_cache_store.get_redis_client", return_value=redis_mock - ): + with patch("utils.atc_cache_store.get_redis_client", return_value=redis_mock): await store.set_l1(CHAT_ID, API_NAME, params_int, RESPONSE) result = await store.get_l1(CHAT_ID, API_NAME, params_str) @@ -339,9 +319,7 @@ async def test_list_response_survives_round_trip(self): _, redis_mock = _fake_redis_store() list_response = [{"date": "2026-02-24"}, {"date": "2026-06-23"}] - with patch( - "src.utils.atc_cache_store.get_redis_client", return_value=redis_mock - ): + with patch("utils.atc_cache_store.get_redis_client", return_value=redis_mock): await store.set_l1(CHAT_ID, API_NAME, PARAMS, list_response) result = await store.get_l1(CHAT_ID, API_NAME, PARAMS) @@ -360,9 +338,7 @@ async def test_round_trip_returns_correct_context_list(self): ctx = _make_last_call_ctx() _, redis_mock = _fake_redis_store() - with patch( - "src.utils.atc_cache_store.get_redis_client", return_value=redis_mock - ): + with patch("utils.atc_cache_store.get_redis_client", return_value=redis_mock): await store.set_l2(CHAT_ID, [ctx]) result = await store.get_l2(CHAT_ID) @@ -381,9 +357,7 @@ async def test_multi_intent_stores_all_entries(self): ctx2 = _make_last_call_ctx("get_electricity_prices") _, redis_mock = _fake_redis_store() - with patch( - "src.utils.atc_cache_store.get_redis_client", return_value=redis_mock - ): + with patch("utils.atc_cache_store.get_redis_client", return_value=redis_mock): await store.set_l2(CHAT_ID, [ctx1, ctx2]) result = await store.get_l2(CHAT_ID) @@ -400,9 +374,7 @@ async def test_set_l2_uses_correct_ttl(self): redis_mock = _make_redis_mock() ctx = _make_last_call_ctx() - with patch( - "src.utils.atc_cache_store.get_redis_client", return_value=redis_mock - ): + with patch("utils.atc_cache_store.get_redis_client", return_value=redis_mock): await store.set_l2(CHAT_ID, [ctx]) assert redis_mock.set.call_args[1]["ex"] == ATC_LAST_CALL_TTL_SECONDS @@ -414,9 +386,7 @@ async def test_set_l2_writes_to_correct_key(self): ctx = _make_last_call_ctx() expected_key = ATCCacheStore._l2_key(CHAT_ID) - with patch( - "src.utils.atc_cache_store.get_redis_client", return_value=redis_mock - ): + with patch("utils.atc_cache_store.get_redis_client", return_value=redis_mock): await store.set_l2(CHAT_ID, [ctx]) assert redis_mock.set.call_args[0][0] == expected_key @@ -426,9 +396,7 @@ async def test_get_l2_returns_none_on_miss(self): store = ATCCacheStore() redis_mock = _make_redis_mock() - with patch( - "src.utils.atc_cache_store.get_redis_client", return_value=redis_mock - ): + with patch("utils.atc_cache_store.get_redis_client", return_value=redis_mock): result = await store.get_l2(CHAT_ID) assert result is None @@ -436,7 +404,7 @@ async def test_get_l2_returns_none_on_miss(self): @pytest.mark.asyncio async def test_get_l2_returns_none_when_redis_unavailable(self): store = ATCCacheStore() - with patch("src.utils.atc_cache_store.get_redis_client", return_value=None): + with patch("utils.atc_cache_store.get_redis_client", return_value=None): result = await store.get_l2(CHAT_ID) assert result is None @@ -447,9 +415,7 @@ async def test_get_l2_returns_none_on_exception(self): redis_mock = _make_redis_mock() redis_mock.get = AsyncMock(side_effect=RuntimeError("connection reset")) - with patch( - "src.utils.atc_cache_store.get_redis_client", return_value=redis_mock - ): + with patch("utils.atc_cache_store.get_redis_client", return_value=redis_mock): result = await store.get_l2(CHAT_ID) assert result is None @@ -458,7 +424,7 @@ async def test_get_l2_returns_none_on_exception(self): async def test_set_l2_no_op_when_redis_unavailable(self): store = ATCCacheStore() ctx = _make_last_call_ctx() - with patch("src.utils.atc_cache_store.get_redis_client", return_value=None): + with patch("utils.atc_cache_store.get_redis_client", return_value=None): await store.set_l2(CHAT_ID, [ctx]) # must not raise @pytest.mark.asyncio @@ -468,9 +434,7 @@ async def test_set_l2_no_op_on_redis_exception(self): redis_mock.set = AsyncMock(side_effect=RuntimeError("write error")) ctx = _make_last_call_ctx() - with patch( - "src.utils.atc_cache_store.get_redis_client", return_value=redis_mock - ): + with patch("utils.atc_cache_store.get_redis_client", return_value=redis_mock): await store.set_l2(CHAT_ID, [ctx]) # must not raise @@ -486,9 +450,7 @@ async def test_calls_delete_with_correct_l2_key(self): redis_mock = _make_redis_mock() expected_key = ATCCacheStore._l2_key(CHAT_ID) - with patch( - "src.utils.atc_cache_store.get_redis_client", return_value=redis_mock - ): + with patch("utils.atc_cache_store.get_redis_client", return_value=redis_mock): await store.invalidate_l2(CHAT_ID) redis_mock.delete.assert_called_once_with(expected_key) @@ -499,9 +461,7 @@ async def test_get_l2_returns_none_after_invalidate(self): _, redis_mock = _fake_redis_store() ctx = _make_last_call_ctx() - with patch( - "src.utils.atc_cache_store.get_redis_client", return_value=redis_mock - ): + with patch("utils.atc_cache_store.get_redis_client", return_value=redis_mock): await store.set_l2(CHAT_ID, [ctx]) await store.invalidate_l2(CHAT_ID) result = await store.get_l2(CHAT_ID) @@ -514,9 +474,7 @@ async def test_invalidate_only_deletes_l2_key_not_l1(self): store = ATCCacheStore() redis_mock = _make_redis_mock() - with patch( - "src.utils.atc_cache_store.get_redis_client", return_value=redis_mock - ): + with patch("utils.atc_cache_store.get_redis_client", return_value=redis_mock): await store.invalidate_l2(CHAT_ID) deleted_key: str = redis_mock.delete.call_args[0][0] @@ -526,7 +484,7 @@ async def test_invalidate_only_deletes_l2_key_not_l1(self): @pytest.mark.asyncio async def test_no_op_when_redis_unavailable(self): store = ATCCacheStore() - with patch("src.utils.atc_cache_store.get_redis_client", return_value=None): + with patch("utils.atc_cache_store.get_redis_client", return_value=None): await store.invalidate_l2(CHAT_ID) # must not raise @pytest.mark.asyncio @@ -535,7 +493,5 @@ async def test_no_op_on_redis_exception(self): redis_mock = _make_redis_mock() redis_mock.delete = AsyncMock(side_effect=RuntimeError("gone")) - with patch( - "src.utils.atc_cache_store.get_redis_client", return_value=redis_mock - ): + with patch("utils.atc_cache_store.get_redis_client", return_value=redis_mock): await store.invalidate_l2(CHAT_ID) # must not raise diff --git a/tests/test_direct_step_executor.py b/tests/test_direct_step_executor.py index 95bc582e..be351ecf 100644 --- a/tests/test_direct_step_executor.py +++ b/tests/test_direct_step_executor.py @@ -10,8 +10,8 @@ import pytest -from src.models.request_models import OrchestrationRequest -from src.tool_classifier.workflows.service_workflow import ServiceWorkflowExecutor +from models.request_models import OrchestrationRequest +from tool_classifier.workflows.service_workflow import ServiceWorkflowExecutor def _make_request( diff --git a/tests/test_follow_up_detector.py b/tests/test_follow_up_detector.py index 6877aa2e..135f8305 100644 --- a/tests/test_follow_up_detector.py +++ b/tests/test_follow_up_detector.py @@ -7,7 +7,7 @@ import dspy import pytest -from src.tool_classifier.follow_up_detector import ( +from tool_classifier.follow_up_detector import ( FollowUpDetectorModule, _validate_updated_params, ) diff --git a/tests/test_multi_api_caller.py b/tests/test_multi_api_caller.py index 906faf4f..1a115100 100644 --- a/tests/test_multi_api_caller.py +++ b/tests/test_multi_api_caller.py @@ -4,16 +4,16 @@ import pytest -from src.tool_classifier.api_caller import APICaller -from src.tool_classifier.constants import ( +from tool_classifier.api_caller import APICaller +from tool_classifier.constants import ( CIRCUIT_BREAKER_OPEN_MESSAGES, MULTI_API_BATCH_TIMEOUT, MULTI_API_PARTIAL_FAILURE_MESSAGES, SERVICE_TIMEOUT_MESSAGES, SERVICE_UNAVAILABLE_MESSAGES, ) -from src.tool_classifier.models import APICallResult -from src.tool_classifier.multi_api_caller import MultiAPICaller +from tool_classifier.models import APICallResult +from tool_classifier.multi_api_caller import MultiAPICaller # --------------------------------------------------------------------------- diff --git a/tests/test_multi_response_formatter.py b/tests/test_multi_response_formatter.py index c5eaaea0..92562168 100644 --- a/tests/test_multi_response_formatter.py +++ b/tests/test_multi_response_formatter.py @@ -8,7 +8,7 @@ import dspy.streaming import pytest -from src.tool_classifier.multi_response_formatter import ( +from tool_classifier.multi_response_formatter import ( MultiResponseFormatterModule, _MULTI_FORMATTER_ERROR_MESSAGES, _MAX_TOTAL_RESPONSE_BYTES, diff --git a/tests/test_param_extractor.py b/tests/test_param_extractor.py index 500fa2d7..f47d2450 100644 --- a/tests/test_param_extractor.py +++ b/tests/test_param_extractor.py @@ -8,7 +8,7 @@ import dspy.streaming import pytest -from src.tool_classifier.param_extractor import ( +from tool_classifier.param_extractor import ( ParamExtractionModule, strip_format_hints, ) diff --git a/tests/test_qdrant_manager.py b/tests/test_qdrant_manager.py index 58b96778..73554999 100644 --- a/tests/test_qdrant_manager.py +++ b/tests/test_qdrant_manager.py @@ -103,8 +103,6 @@ def _make_qdrant_client( """Build a mock QdrantClient.""" client = MagicMock() - col_mock = MagicMock() - col_mock.name = "some_collection" collections_result = MagicMock() collections_result.collections = [MagicMock(name=n) for n in collection_names] # Fix: MagicMock(name=n) doesn't work as expected — set attribute explicitly diff --git a/tests/test_tool_classifier.py b/tests/test_tool_classifier.py index a9158b42..5bda698a 100644 --- a/tests/test_tool_classifier.py +++ b/tests/test_tool_classifier.py @@ -12,6 +12,8 @@ - Qdrant timeout during classification → fallback """ +from __future__ import annotations + from typing import Any, Dict, List, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -21,6 +23,7 @@ from models.request_models import OrchestrationRequest from tool_classifier.classifier import ToolClassifier from tool_classifier.enums import WorkflowType +from tool_classifier.intent_decomposer import IntentDecomposerModule # --------------------------------------------------------------------------- @@ -476,3 +479,282 @@ async def test_qdrant_timeout_falls_back_to_context(self) -> None: ) assert result.workflow == WorkflowType.CONTEXT + + +# --------------------------------------------------------------------------- +# IntentDecomposerModule +# --------------------------------------------------------------------------- + + +class TestIntentDecomposer: + """Unit tests for IntentDecomposerModule.forward() and .decompose(). + + The DSPy predictor is replaced with a MagicMock so no real LLM is called. + """ + + def _make_module_with_prediction( + self, + mode: str, + sub_queries: str, + ) -> IntentDecomposerModule: + module = IntentDecomposerModule() + mock_pred = MagicMock() + mock_pred.mode = mode + mock_pred.sub_queries = sub_queries + module.predictor = MagicMock(return_value=mock_pred) + return module + + def test_forward_single_mode_returns_single(self) -> None: + from tool_classifier.intent_decomposer import DecompositionResult + + module = self._make_module_with_prediction("single", "[]") + result = module.forward("What are the public holidays in Estonia?") + + assert isinstance(result, DecompositionResult) + assert result.mode == "single" + assert result.sub_queries == [] + + def test_forward_parallel_mode_returns_sub_queries(self) -> None: + from tool_classifier.intent_decomposer import DecompositionResult + + module = self._make_module_with_prediction( + "parallel", + '["public holidays in Estonia", "weather in Tallinn"]', + ) + result = module.forward( + "What are the public holidays in Estonia AND the weather in Tallinn?" + ) + + assert isinstance(result, DecompositionResult) + assert result.mode == "parallel" + assert len(result.sub_queries) == 2 + assert "public holidays in Estonia" in result.sub_queries + + def test_forward_caps_sub_queries_at_max_endpoints(self) -> None: + """sub_queries exceeding MULTI_API_MAX_ENDPOINTS are truncated.""" + from tool_classifier.intent_decomposer import ( + DecompositionResult, + IntentDecomposerModule, + ) + from tool_classifier.constants import MULTI_API_MAX_ENDPOINTS + + module = IntentDecomposerModule() + # Build a prediction with more sub-queries than the cap allows + over_cap = ["query " + str(i) for i in range(MULTI_API_MAX_ENDPOINTS + 2)] + import json + + mock_pred = MagicMock() + mock_pred.mode = "parallel" + mock_pred.sub_queries = json.dumps(over_cap) + module.predictor = MagicMock(return_value=mock_pred) + + result = module.forward("many intents query") + + assert isinstance(result, DecompositionResult) + assert result.mode == "parallel" + assert len(result.sub_queries) == MULTI_API_MAX_ENDPOINTS + + def test_forward_unexpected_mode_falls_back_to_single(self) -> None: + from tool_classifier.intent_decomposer import DecompositionResult + + module = self._make_module_with_prediction("unknown_value", "[]") + result = module.forward("some query") + + assert isinstance(result, DecompositionResult) + assert result.mode == "single" + assert result.sub_queries == [] + + def test_forward_parallel_with_fewer_than_2_sub_queries_falls_back(self) -> None: + """mode=parallel but only 1 sub-query parsed → conservative fallback to single.""" + from tool_classifier.intent_decomposer import DecompositionResult + + module = self._make_module_with_prediction("parallel", '["only one query"]') + result = module.forward("something") + + assert isinstance(result, DecompositionResult) + assert result.mode == "single" + + def test_forward_predictor_exception_falls_back_to_single(self) -> None: + """Any exception from the DSPy predictor → conservative single fallback.""" + from tool_classifier.intent_decomposer import ( + DecompositionResult, + IntentDecomposerModule, + ) + + module = IntentDecomposerModule() + module.predictor = MagicMock(side_effect=RuntimeError("LLM unavailable")) + + result = module.forward("multi intent query") + + assert isinstance(result, DecompositionResult) + assert result.mode == "single" + + def test_forward_parallel_with_invalid_json_falls_back_to_single(self) -> None: + """Invalid JSON in sub_queries → falls back to mode=single.""" + from tool_classifier.intent_decomposer import DecompositionResult + + module = self._make_module_with_prediction("parallel", "not valid json") + result = module.forward("holidays and weather") + + assert isinstance(result, DecompositionResult) + assert result.mode == "single" + + def test_forward_markdown_fenced_json_parsed_correctly(self) -> None: + """sub_queries wrapped in markdown code fences are unwrapped before JSON parse.""" + from tool_classifier.intent_decomposer import DecompositionResult + + fenced = '```json\n["query A", "query B"]\n```' + module = self._make_module_with_prediction("parallel", fenced) + result = module.forward("query") + + assert isinstance(result, DecompositionResult) + assert result.mode == "parallel" + assert result.sub_queries == ["query A", "query B"] + + @pytest.mark.asyncio + async def test_decompose_async_wraps_forward(self) -> None: + """.decompose() is the async wrapper — result matches .forward() output.""" + from tool_classifier.intent_decomposer import DecompositionResult + + module = self._make_module_with_prediction("parallel", '["sub A", "sub B"]') + + with patch( + "tool_classifier.intent_decomposer.asyncio.to_thread", + new_callable=AsyncMock, + ) as mock_thread: + mock_thread.return_value = DecompositionResult( + mode="parallel", sub_queries=["sub A", "sub B"] + ) + result = await module.decompose("two intents") + + assert isinstance(result, DecompositionResult) + assert result.mode == "parallel" + assert result.sub_queries == ["sub A", "sub B"] + + +# --------------------------------------------------------------------------- +# classify() — MULTI_INTENT_ENABLED feature flag toggle +# --------------------------------------------------------------------------- + + +class TestClassifyMultiIntentFeatureFlag: + """Verify the MULTI_INTENT_ENABLED flag gates the parallel decomposition path.""" + + @pytest.mark.asyncio + async def test_multi_intent_disabled_suppresses_hint_result(self) -> None: + """When MULTI_INTENT_ENABLED=False a multi_intent_hint result is suppressed + and the classifier falls through to CONTEXT/RAG.""" + from tool_classifier.api_semantic_searcher import APIToolSearchResult + + svc = _make_orchestration_service(session_store=None) + classifier = _make_classifier(svc) + + # Build a result with multi_intent_hint=True (disambiguator rejected all) + hint_result = MagicMock(spec=APIToolSearchResult) + hint_result.endpoint_id = "ep-holidays" + hint_result.name = "get_public_holidays" + hint_result.description = "Returns public holidays" + hint_result.method = "GET" + hint_result.url = "https://openholidaysapi.org/PublicHolidays" + hint_result.params = [] + hint_result.cosine_score = 0.55 + hint_result.rrf_score = 0.01 + hint_result.confidence = "medium" + hint_result.llm_validated = False + hint_result.multi_intent_hint = True + hint_result.to_dict.return_value = {"endpoint_id": "ep-holidays"} + + classifier.api_tool_searcher.search = AsyncMock(return_value=[hint_result]) + + with ( + patch( + "tool_classifier.classifier.FeatureFlags.API_TOOL_CALLING_WORKFLOW_ENABLED", + True, + ), + patch( + "tool_classifier.classifier.FeatureFlags.MULTI_INTENT_ENABLED", + False, + ), + patch( + "tool_classifier.classifier.FeatureFlags.SERVICE_WORKFLOW_ENABLED", + False, + ), + ): + result = await classifier.classify( + query="public holidays AND weather", + conversation_history=[], + language="en", + request=_make_request("public holidays AND weather"), + ) + + assert result.workflow == WorkflowType.CONTEXT + + @pytest.mark.asyncio + async def test_multi_intent_enabled_triggers_decomposer_on_ambiguous_band( + self, + ) -> None: + """MULTI_INTENT_ENABLED=True + cosine in ambiguous band → IntentDecomposer runs.""" + from tool_classifier.api_semantic_searcher import APIToolSearchResult + from tool_classifier.intent_decomposer import DecompositionResult + + svc = _make_orchestration_service(session_store=None) + classifier = _make_classifier(svc) + + # A result in the ambiguous band (not llm_validated, not multi_intent_hint) + ambiguous_result = MagicMock(spec=APIToolSearchResult) + ambiguous_result.endpoint_id = "ep-holidays" + ambiguous_result.name = "get_public_holidays" + ambiguous_result.description = "Returns public holidays" + ambiguous_result.method = "GET" + ambiguous_result.url = "https://openholidaysapi.org/PublicHolidays" + ambiguous_result.params = [] + ambiguous_result.cosine_score = 0.50 # in [0.40, 0.60) band + ambiguous_result.rrf_score = 0.01 + ambiguous_result.confidence = "medium" + ambiguous_result.llm_validated = False + ambiguous_result.multi_intent_hint = False + ambiguous_result.to_dict.return_value = { + "endpoint_id": "ep-holidays", + "name": "get_public_holidays", + "description": "Returns public holidays", + "method": "GET", + "url": "https://openholidaysapi.org/PublicHolidays", + "params": [], + "cosine_score": 0.50, + "rrf_score": 0.01, + "confidence": "medium", + } + + classifier.api_tool_searcher.search = AsyncMock(return_value=[ambiguous_result]) + + # IntentDecomposer returns single → falls through to single-endpoint path + decomposer_result = DecompositionResult(mode="single", sub_queries=[]) + classifier.intent_decomposer.decompose = AsyncMock( + return_value=decomposer_result + ) + + with ( + patch( + "tool_classifier.classifier.FeatureFlags.API_TOOL_CALLING_WORKFLOW_ENABLED", + True, + ), + patch( + "tool_classifier.classifier.FeatureFlags.MULTI_INTENT_ENABLED", + True, + ), + patch( + "tool_classifier.classifier.FeatureFlags.SERVICE_WORKFLOW_ENABLED", + False, + ), + ): + result = await classifier.classify( + query="public holidays AND weather", + conversation_history=[], + language="en", + request=_make_request("public holidays AND weather"), + ) + + # Decomposer was consulted + classifier.intent_decomposer.decompose.assert_awaited_once() + # Single mode → normal API_TOOL_CALLING result + assert result.workflow == WorkflowType.API_TOOL_CALLING