Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/routing/router.py: 40%

1"""Main query router for RAG systems.""" 

2 

3from __future__ import annotations 

4 

5from dataclasses import dataclass, field 

6 

7from lexigram.ai.rag.routing.analyzer import QueryAnalyzer 

8from lexigram.ai.rag.routing.strategies.base import RoutingStrategy 

9from lexigram.ai.rag.routing.strategies.rule_based import RuleBasedRouter 

10from lexigram.ai.rag.routing.types import ( 

11 DataSource, 

12 QueryFeatures, 

13 RoutingDecision, 

14) 

15 

16 

17@dataclass 

18class RoutingStatistics: 

19 """Statistics for query routing. 

20 

21 Attributes: 

22 total_queries: Total number of queries routed. 

23 by_strategy: Count of queries by strategy. 

24 by_data_source: Count of queries by data source. 

25 avg_confidence: Average confidence score. 

26 high_confidence_count: Number of high-confidence decisions (>0.7). 

27 low_confidence_count: Number of low-confidence decisions (<0.3). 

28 """ 

29 

30 total_queries: int = 0 

31 by_strategy: dict[str, int] = field(default_factory=dict) 

32 by_data_source: dict[str, int] = field(default_factory=dict) 

33 avg_confidence: float = 0.0 

34 high_confidence_count: int = 0 

35 low_confidence_count: int = 0 

36 

37 def update(self, decision: RoutingDecision) -> None: 

38 """Update statistics with a routing decision. 

39 

40 Args: 

41 decision: Routing decision to record. 

42 """ 

43 self.total_queries += 1 

44 

45 # Track strategy 

46 self.by_strategy[decision.strategy] = ( 

47 self.by_strategy.get(decision.strategy, 0) + 1 

48 ) 

49 

50 # Track data sources 

51 for source in decision.data_sources: 

52 self.by_data_source[source.name] = ( 

53 self.by_data_source.get(source.name, 0) + 1 

54 ) 

55 

56 # Update confidence stats 

57 old_avg = self.avg_confidence 

58 self.avg_confidence = ( 

59 old_avg * (self.total_queries - 1) + decision.confidence 

60 ) / self.total_queries 

61 

62 if decision.confidence > 0.7: 

63 self.high_confidence_count += 1 

64 elif decision.confidence < 0.3: 

65 self.low_confidence_count += 1 

66 

67 def to_dict(self) -> dict: 

68 """Convert statistics to dictionary. 

69 

70 Returns: 

71 Dictionary representation of statistics. 

72 """ 

73 return { 

74 "total_queries": self.total_queries, 

75 "by_strategy": self.by_strategy, 

76 "by_data_source": self.by_data_source, 

77 "avg_confidence": round(self.avg_confidence, 3), 

78 "high_confidence_count": self.high_confidence_count, 

79 "low_confidence_count": self.low_confidence_count, 

80 "high_confidence_rate": round( 

81 self.high_confidence_count / max(self.total_queries, 1), 

82 3, 

83 ), 

84 } 

85 

86 

87class QueryRouter: 

88 """Main query router for RAG systems. 

89 

90 Orchestrates query analysis and routing to appropriate data sources 

91 using configurable routing strategies. 

92 

93 Example: 

94 ```python 

95 from lexigram.ai.rag import ( 

96 QueryRouter, 

97 DataSourceProtocol, 

98 DataSourceType, 

99 RuleBasedRouter 

100 ) 

101 

102 # Create router 

103 router = QueryRouter( 

104 analyzer=QueryAnalyzer(), 

105 strategy=RuleBasedRouter.with_defaults() 

106 ) 

107 

108 # Register data sources 

109 router.register_source(DataSourceProtocol( 

110 name="docs_vector", 

111 type=DataSourceType.VECTOR_STORE, 

112 description="Documentation vector store", 

113 capabilities=["dense_search", "semantic_search"], 

114 priority=10 

115 )) 

116 

117 # Route query 

118 decision = await router.route("How do I configure authentication?") 

119 logger.info(f"Route to: {decision.data_sources[0].name}") 

120 logger.info(f"Strategy: {decision.strategy}") 

121 logger.info(f"Confidence: {decision.confidence}") 

122 

123 # Get statistics 

124 stats = router.get_statistics() 

125 logger.info(f"Total queries: {stats.total_queries}") 

126 logger.info(f"Avg confidence: {stats.avg_confidence}") 

127 ``` 

128 """ 

129 

130 def __init__( 

131 self, 

132 *, 

133 analyzer: QueryAnalyzer | None = None, 

134 strategy: RoutingStrategy | None = None, 

135 ): 

136 """Initialize the query router. 

137 

138 Args: 

139 analyzer: Query analyzer for feature extraction. 

140 strategy: Routing strategy to use. 

141 """ 

142 self.analyzer = analyzer or QueryAnalyzer() 

143 self.strategy = strategy or RuleBasedRouter.with_defaults() 

144 self.data_sources: list[DataSource] = [] 

145 self.statistics = RoutingStatistics() 

146 

147 def register_source(self, source: DataSource) -> None: 

148 """Register a data source. 

149 

150 Args: 

151 source: Data source to register. 

152 """ 

153 # Check if source already exists 

154 existing = list(filter(lambda s: s.name == source.name, self.data_sources)) 

155 if existing: 

156 # Update existing source 

157 self.data_sources.remove(existing[0]) 

158 

159 self.data_sources.append(source) 

160 

161 # Sort by priority (highest first) 

162 self.data_sources.sort(key=lambda s: s.priority, reverse=True) 

163 

164 def unregister_source(self, name: str) -> bool: 

165 """Unregister a data source. 

166 

167 Args: 

168 name: Name of the data source to unregister. 

169 

170 Returns: 

171 True if source was found and removed, False otherwise. 

172 """ 

173 initial_count = len(self.data_sources) 

174 self.data_sources = list(filter(lambda s: s.name != name, self.data_sources)) 

175 return len(self.data_sources) < initial_count 

176 

177 def get_source(self, name: str) -> DataSource | None: 

178 """Get a data source by name. 

179 

180 Args: 

181 name: Name of the data source. 

182 

183 Returns: 

184 Data source if found, None otherwise. 

185 """ 

186 for source in self.data_sources: 

187 if source.name == name: 

188 return source 

189 return None 

190 

191 def list_sources(self) -> list[DataSource]: 

192 """List all registered data sources. 

193 

194 Returns: 

195 List of registered data sources. 

196 """ 

197 return self.data_sources.copy() 

198 

199 async def route( 

200 self, 

201 query: str, 

202 *, 

203 features: QueryFeatures | None = None, 

204 ) -> RoutingDecision: 

205 """Route a query to appropriate data sources. 

206 

207 Args: 

208 query: Query text to route. 

209 features: Pre-extracted query features (optional). 

210 

211 Returns: 

212 Routing decision with selected data sources and strategy. 

213 """ 

214 # Analyze query if features not provided 

215 if features is None: 

216 features = await self.analyzer.analyze(query) 

217 

218 # Route using strategy 

219 decision = await self.strategy.route(features, self.data_sources) 

220 

221 # Update statistics 

222 self.statistics.update(decision) 

223 

224 return decision 

225 

226 async def route_batch( 

227 self, 

228 queries: list[str], 

229 ) -> list[RoutingDecision]: 

230 """Route multiple queries. 

231 

232 Args: 

233 queries: List of query texts to route. 

234 

235 Returns: 

236 List of routing decisions. 

237 """ 

238 decisions = [] 

239 for query in queries: 

240 decision = await self.route(query) 

241 decisions.append(decision) 

242 return decisions 

243 

244 def get_statistics(self) -> RoutingStatistics: 

245 """Get routing statistics. 

246 

247 Returns: 

248 Current routing statistics. 

249 """ 

250 return self.statistics 

251 

252 def reset_statistics(self) -> None: 

253 """Reset routing statistics.""" 

254 self.statistics = RoutingStatistics() 

255 

256 def set_strategy(self, strategy: RoutingStrategy) -> None: 

257 """Set the routing strategy. 

258 

259 Args: 

260 strategy: New routing strategy to use. 

261 """ 

262 self.strategy = strategy 

263 

264 def set_analyzer(self, analyzer: QueryAnalyzer) -> None: 

265 """Set the query analyzer. 

266 

267 Args: 

268 analyzer: New query analyzer to use. 

269 """ 

270 self.analyzer = analyzer