Coverage for src/qdrant_loader_mcp_server/search/hybrid/orchestration/topic_chain.py: 83%

1from __future__ import annotations 

2 

3import inspect 

4from typing import Any 

5 

6from ...components.search_result_models import HybridSearchResult 

7from ...enhanced.topic_search_chain import ChainStrategy, TopicSearchChain 

8 

9 

10async def generate_topic_search_chain( 

11 engine: Any, 

12 query: str, 

13 strategy: ChainStrategy = ChainStrategy.MIXED_EXPLORATION, 

14 max_links: int = 5, 

15 initialize_from_search: bool = True, 

16) -> TopicSearchChain: 

17 # Use public accessor instead of private attribute 

18 if initialize_from_search: 

19 try: 

20 init_attr = getattr(engine, "is_topic_chains_initialized", False) 

21 is_initialized: bool 

22 if callable(init_attr): 

23 init_result = init_attr() 

24 if inspect.isawaitable(init_result): 

25 init_result = await init_result 

26 is_initialized = bool(init_result) 

27 else: 

28 is_initialized = bool(init_attr) 

29 except Exception: 

30 # Be conservative: if we cannot determine, assume not initialized 

31 is_initialized = False 

32 if not is_initialized: 

33 await _initialize_topic_relationships(engine, query) 

34 result = engine.topic_chain_generator.generate_search_chain( 

35 original_query=query, strategy=strategy, max_links=max_links 

36 ) 

37 if inspect.isawaitable(result): 

38 return await result 

39 return result 

40 

41 

42async def execute_topic_chain_search( 

43 engine: Any, 

44 topic_chain: TopicSearchChain, 

45 results_per_link: int = 3, 

46 source_types: list[str] | None = None, 

47 project_ids: list[str] | None = None, 

48) -> dict[str, list[HybridSearchResult]]: 

49 chain_results: dict[str, list[HybridSearchResult]] = {} 

50 

51 original_results = await engine.search( 

52 query=topic_chain.original_query, 

53 limit=results_per_link, 

54 source_types=source_types, 

55 project_ids=project_ids, 

56 ) 

57 chain_results[topic_chain.original_query] = original_results 

58 

59 for link in topic_chain.chain_links: 

60 try: 

61 link_results = await engine.search( 

62 query=link.query, 

63 limit=results_per_link, 

64 source_types=source_types, 

65 project_ids=project_ids, 

66 ) 

67 chain_results[link.query] = link_results 

68 except Exception: 

69 # Log the exception with context; include traceback 

70 engine.logger.exception( 

71 "Error running topic chain for query=%s", link.query 

72 ) 

73 chain_results[link.query] = [] 

74 

75 return chain_results 

76 

77 

78async def _initialize_topic_relationships(engine: Any, sample_query: str) -> None: 

79 sample_results = await engine.search( 

80 query=sample_query, limit=20, source_types=None, project_ids=None 

81 ) 

82 if sample_results: 

83 engine.topic_chain_generator.initialize_from_results(sample_results) 

84 # Mark initialization via public API instead of touching private attribute 

85 if hasattr(engine, "set_topic_chains_initialized"): 

86 engine.set_topic_chains_initialized(True) 

87 else: 

88 engine.mark_topic_chains_initialized()