77 lines
3.0 KiB
Python
77 lines
3.0 KiB
Python
"""Tests for the Phase 1.3 QueryDecomposer component."""
|
|
|
|
import json
|
|
from typing import List
|
|
|
|
import pytest
|
|
|
|
from app.services.query_decomposer import QueryDecomposer
|
|
|
|
|
|
class MockLLMClient:
|
|
def __init__(self, response: str):
|
|
self._response = response
|
|
self.last_prompt = None
|
|
self.last_step_name = None
|
|
|
|
async def complete(self, prompt: str, temperature: float = 0.7, step_name: str = "LLM") -> str:
|
|
self.last_prompt = prompt
|
|
self.last_step_name = step_name
|
|
return self._response
|
|
|
|
|
|
async def test_decompose_valid_json(mock_prompt_service):
|
|
llm = MockLLMClient('["alpha", "beta", "gamma"]')
|
|
decomposer = QueryDecomposer(llm, prompt_service=mock_prompt_service)
|
|
questions, prompt = await decomposer.decompose("What are keywords for X?")
|
|
assert questions == ["alpha", "beta", "gamma"]
|
|
assert prompt == "Given this question: 'What are keywords for X?'\n\nBreak it down into 2-5 simplified sub-questions that would help search for relevant information. Each sub-question should be short and focused on one aspect. Return as a JSON array of strings."
|
|
assert llm.last_prompt == prompt
|
|
|
|
|
|
async def test_decompose_empty_question_returns_empty(mock_prompt_service):
|
|
llm = MockLLMClient('["should_not_be_used"]')
|
|
decomposer = QueryDecomposer(llm, prompt_service=mock_prompt_service)
|
|
questions, prompt = await decomposer.decompose("")
|
|
assert questions == []
|
|
assert prompt == ""
|
|
assert llm.last_prompt is None
|
|
|
|
|
|
async def test_decompose_invalid_json_returns_empty(mock_prompt_service):
|
|
llm = MockLLMClient("not-json")
|
|
decomposer = QueryDecomposer(llm, prompt_service=mock_prompt_service)
|
|
questions, prompt = await decomposer.decompose("Question?")
|
|
assert questions == []
|
|
assert prompt != ""
|
|
|
|
|
|
async def test_decompose_non_list_json_returns_empty(mock_prompt_service):
|
|
llm = MockLLMClient("{\"a\": 1}")
|
|
decomposer = QueryDecomposer(llm, prompt_service=mock_prompt_service)
|
|
questions, prompt = await decomposer.decompose("Question?")
|
|
assert questions == []
|
|
assert prompt != ""
|
|
|
|
|
|
async def test_decompose_mixed_types_coerced_to_strings(mock_prompt_service):
|
|
llm = MockLLMClient('["a", 2, null]')
|
|
decomposer = QueryDecomposer(llm, prompt_service=mock_prompt_service)
|
|
questions, prompt = await decomposer.decompose("Question?")
|
|
assert questions == ["a", "2", "None"]
|
|
assert prompt != ""
|
|
|
|
|
|
async def test_decompose_json_in_markdown_code_block(mock_prompt_service):
|
|
llm = MockLLMClient('```json\n["project", "manager", "limits"]\n```')
|
|
decomposer = QueryDecomposer(llm, prompt_service=mock_prompt_service)
|
|
questions, prompt = await decomposer.decompose("What are the limits?")
|
|
assert questions == ["project", "manager", "limits"]
|
|
|
|
|
|
async def test_decompose_json_in_plain_code_block(mock_prompt_service):
|
|
llm = MockLLMClient('```\n["alpha", "beta"]\n```')
|
|
decomposer = QueryDecomposer(llm, prompt_service=mock_prompt_service)
|
|
questions, prompt = await decomposer.decompose("Keywords?")
|
|
assert questions == ["alpha", "beta"]
|