Coverage for intelligence_toolkit/tests/unit/query_text_data/test_answer_builder.py: 100%

1import json 

2 

3import pytest 

4 

5from intelligence_toolkit.query_text_data.answer_builder import answer_query 

6from intelligence_toolkit.query_text_data.classes import AnswerObject 

7 

8 

9class _DummyCommentary: 

10 def __init__(self): 

11 self.structure = { 

12 "themes": {"Theme Alpha": ["point-1"]}, 

13 "points": {"point-1": "Point One"}, 

14 "point_sources": {"point-1": [1]}, 

15 } 

16 

17 def get_clustered_cids(self): 

18 return {"Theme Alpha": [1]} 

19 

20 

21class _DummyProcessedChunks: 

22 cid_to_text = { 

23 1: json.dumps( 

24 {"title": "Doc", "chunk_id": 1, "text_chunk": "Chunk body"} 

25 ) 

26 } 

27 

28 

29@pytest.mark.asyncio 

30async def test_answer_query_returns_report_with_sources(mocker): 

31 mocker.patch( 

32 "intelligence_toolkit.query_text_data.answer_builder.utils.prepare_messages", 

33 side_effect=lambda *args, **kwargs: {"args": args, "kwargs": kwargs}, 

34 ) 

35 mocker.patch( 

36 "intelligence_toolkit.query_text_data.answer_builder.utils.map_generate_text", 

37 new=mocker.AsyncMock( 

38 return_value=[ 

39 json.dumps( 

40 { 

41 "theme_title": "Theme Alpha", 

42 "theme_points": [ 

43 { 

44 "point_title": "Point One", 

45 "point_evidence": "**Source evidence**: Point One [source: 1]", 

46 "point_commentary": "**AI commentary**: Point One", 

47 } 

48 ], 

49 } 

50 ) 

51 ] 

52 ), 

53 ) 

54 mocker.patch( 

55 "intelligence_toolkit.query_text_data.answer_builder.utils.generate_text", 

56 return_value=json.dumps( 

57 { 

58 "answer": "Some answer", 

59 "report_title": "Aligned Themes", 

60 "report_overview": "Overview text", 

61 "report_implications": "Implication text", 

62 } 

63 ), 

64 ) 

65 

66 result = await answer_query( 

67 ai_configuration={"model": "test"}, 

68 query="What happened?", 

69 expanded_query="Detailed question", 

70 processed_chunks=_DummyProcessedChunks(), 

71 commentary=_DummyCommentary(), 

72 ) 

73 

74 assert isinstance(result, AnswerObject) 

75 assert result.references == [1] 

76 assert "[source: [1](#source-1)]" in result.extended_answer 

77 assert "#### Source 1" in result.extended_answer 

78 assert "Doc (1)" in result.referenced_chunks