Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/reasoning/strategy_registry.py: 98%

1"""Reasoning strategy registry for multi-hop reasoning.""" 

2 

3from __future__ import annotations 

4 

5from typing import Any, Protocol 

6 

7from lexigram.ai.rag.reasoning.base import ReasoningStrategy 

8from lexigram.ai.rag.reasoning.chain_of_thought import ChainOfThoughtReasoner 

9from lexigram.ai.rag.reasoning.decomposition import QueryDecomposer 

10from lexigram.ai.rag.reasoning.iterative import IterativeRefinementReasoner 

11from lexigram.ai.rag.reasoning.multi_hop import MultiHopReasoner 

12from lexigram.contracts.ai import LLMClientProtocol 

13from lexigram.contracts.ai.vector import DocumentVectorStoreProtocol 

14 

15VectorStoreProtocol = DocumentVectorStoreProtocol 

16 

17 

18class ReasoningStrategyHandler(Protocol): 

19 """Protocol for reasoning strategy handlers.""" 

20 

21 def can_handle(self, strategy: ReasoningStrategy) -> bool: 

22 """Check if this handler can handle the strategy.""" 

23 ... 

24 

25 async def create_and_reason( 

26 self, 

27 strategy: ReasoningStrategy, 

28 llm_client: LLMClientProtocol, 

29 vector_store: VectorStoreProtocol, 

30 query: str, 

31 kwargs: dict[str, Any], 

32 ) -> Any: 

33 """Create reasoner and execute reasoning.""" 

34 ... 

35 

36 

37class MultiHopReasoningHandler: 

38 """Handler for MULTI_HOP reasoning strategy.""" 

39 

40 def can_handle(self, strategy: ReasoningStrategy) -> bool: 

41 return strategy == ReasoningStrategy.MULTI_HOP 

42 

43 async def create_and_reason( 

44 self, 

45 strategy: ReasoningStrategy, 

46 llm_client: LLMClientProtocol, 

47 vector_store: VectorStoreProtocol, 

48 query: str, 

49 kwargs: dict[str, Any], 

50 ) -> Any: 

51 reasoner = MultiHopReasoner(llm_client, vector_store, **kwargs) # type: ignore[arg-type] 

52 return await reasoner.reason(query, **kwargs) 

53 

54 

55class ChainOfThoughtReasoningHandler: 

56 """Handler for CHAIN_OF_THOUGHT reasoning strategy.""" 

57 

58 def can_handle(self, strategy: ReasoningStrategy) -> bool: 

59 return strategy == ReasoningStrategy.CHAIN_OF_THOUGHT 

60 

61 async def create_and_reason( 

62 self, 

63 strategy: ReasoningStrategy, 

64 llm_client: LLMClientProtocol, 

65 vector_store: VectorStoreProtocol, 

66 query: str, 

67 kwargs: dict[str, Any], 

68 ) -> Any: 

69 reasoner = ChainOfThoughtReasoner(llm_client, **kwargs) 

70 return await reasoner.reason(query, **kwargs) 

71 

72 

73class DecompositionReasoningHandler: 

74 """Handler for DECOMPOSITION reasoning strategy.""" 

75 

76 def can_handle(self, strategy: ReasoningStrategy) -> bool: 

77 return strategy == ReasoningStrategy.DECOMPOSITION 

78 

79 async def create_and_reason( 

80 self, 

81 strategy: ReasoningStrategy, 

82 llm_client: LLMClientProtocol, 

83 vector_store: VectorStoreProtocol, 

84 query: str, 

85 kwargs: dict[str, Any], 

86 ) -> Any: 

87 reasoner = QueryDecomposer(llm_client, vector_store, **kwargs) # type: ignore[arg-type] 

88 return await reasoner.reason(query, **kwargs) 

89 

90 

91class IterativeRefinementReasoningHandler: 

92 """Handler for ITERATIVE_REFINEMENT reasoning strategy.""" 

93 

94 def can_handle(self, strategy: ReasoningStrategy) -> bool: 

95 return strategy == ReasoningStrategy.ITERATIVE_REFINEMENT 

96 

97 async def create_and_reason( 

98 self, 

99 strategy: ReasoningStrategy, 

100 llm_client: LLMClientProtocol, 

101 vector_store: VectorStoreProtocol, 

102 query: str, 

103 kwargs: dict[str, Any], 

104 ) -> Any: 

105 reasoner = IterativeRefinementReasoner(llm_client, vector_store, **kwargs) # type: ignore[arg-type] 

106 return await reasoner.reason(query, **kwargs) 

107 

108 

109class ReasoningStrategyRegistry: 

110 """Central registry for reasoning strategy handlers.""" 

111 

112 def __init__(self) -> None: 

113 self._handlers: list[ReasoningStrategyHandler] = [] 

114 

115 @classmethod 

116 def with_defaults(cls) -> ReasoningStrategyRegistry: 

117 """Create a registry pre-populated with all built-in strategy handlers.""" 

118 registry = cls() 

119 registry._handlers = [ 

120 MultiHopReasoningHandler(), 

121 ChainOfThoughtReasoningHandler(), 

122 DecompositionReasoningHandler(), 

123 IterativeRefinementReasoningHandler(), 

124 ] 

125 return registry 

126 

127 def register(self, handler: ReasoningStrategyHandler) -> None: 

128 """Register a new strategy handler.""" 

129 self._handlers.insert(0, handler) 

130 

131 async def reason( 

132 self, 

133 strategy: ReasoningStrategy, 

134 llm_client: LLMClientProtocol, 

135 vector_store: VectorStoreProtocol, 

136 query: str, 

137 kwargs: dict[str, Any], 

138 ) -> Any: 

139 """Execute reasoning using the appropriate strategy.""" 

140 for handler in self._handlers: 

141 if handler.can_handle(strategy): 

142 return await handler.create_and_reason( 

143 strategy, 

144 llm_client, 

145 vector_store, 

146 query, 

147 kwargs, 

148 ) 

149 msg = f"Unknown reasoning strategy: {strategy}" 

150 raise ValueError(msg)