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

53 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-25 07:19 +0800

1"""Rule-based routing strategy.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import Callable 

6from dataclasses import dataclass 

7 

8from lexigram.ai.rag.multimodal.types import Modality 

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

10 DataSource, 

11 DataSourceType, 

12 QueryFeatures, 

13 QueryIntent, 

14 RoutingDecision, 

15) 

16 

17 

18@dataclass 

19class RoutingRule: 

20 """A routing rule for rule-based routing. 

21 

22 Attributes: 

23 name: Unique identifier for the rule. 

24 condition: Function that checks if rule applies to query features. 

25 data_source_types: Preferred data source types when rule matches. 

26 strategy: Retrieval strategy to use when rule matches. 

27 priority: Priority of the rule (higher = checked first). 

28 description: Human-readable description of the rule. 

29 """ 

30 

31 name: str 

32 condition: Callable[[QueryFeatures], bool] 

33 data_source_types: list[DataSourceType] 

34 strategy: str 

35 priority: int = 0 

36 description: str = "" 

37 

38 

39class RuleBasedRouter: 

40 """Rule-based routing strategy using if-then rules. 

41 

42 Routes queries based on configurable rules that match query features 

43 to appropriate data sources and retrieval strategies. 

44 

45 Example: 

46 ```python 

47 router = RuleBasedRouter.with_defaults() 

48 

49 # Add custom rule 

50 router.add_rule(RoutingRule( 

51 name="multimodal_images", 

52 condition=lambda f: Modality.IMAGE in f.modalities, 

53 data_source_types=[DataSourceType.MULTIMODAL_STORE], 

54 strategy="multimodal", 

55 priority=10, 

56 description="Route image queries to multimodal store" 

57 )) 

58 

59 # Route query 

60 decision = await router.route(features, available_sources) 

61 ``` 

62 """ 

63 

64 def __init__(self) -> None: 

65 """Initialize the rule-based router with an empty rules list. 

66 

67 Use `with_defaults()` classmethod to create a router with default rules. 

68 """ 

69 self.rules: list[RoutingRule] = [] 

70 

71 @classmethod 

72 def with_defaults(cls) -> RuleBasedRouter: 

73 """Create a router pre-populated with default routing rules. 

74 

75 Returns: 

76 A RuleBasedRouter with all default rules registered. 

77 """ 

78 instance = cls() 

79 instance._load_default_rules() 

80 return instance 

81 

82 def add_rule(self, rule: RoutingRule) -> None: 

83 """Add a routing rule. 

84 

85 Args: 

86 rule: Routing rule to add. 

87 """ 

88 self.rules.append(rule) 

89 # Sort rules by priority (highest first) 

90 self.rules.sort(key=lambda r: r.priority, reverse=True) 

91 

92 def remove_rule(self, name: str) -> bool: 

93 """Remove a routing rule by name. 

94 

95 Args: 

96 name: Name of the rule to remove. 

97 

98 Returns: 

99 True if rule was found and removed, False otherwise. 

100 """ 

101 initial_count = len(self.rules) 

102 self.rules = list(filter(lambda r: r.name != name, self.rules)) 

103 return len(self.rules) < initial_count 

104 

105 async def route( 

106 self, 

107 features: QueryFeatures, 

108 available_sources: list[DataSource], 

109 ) -> RoutingDecision: 

110 """Route query using rule-based logic. 

111 

112 Args: 

113 features: Extracted query features. 

114 available_sources: List of available data sources. 

115 

116 Returns: 

117 Routing decision based on matched rules. 

118 """ 

119 # Try each rule in priority order 

120 for rule in self.rules: 

121 if rule.condition(features): 

122 # Find matching data sources 

123 matching_sources = [ 

124 source 

125 for source in available_sources 

126 if source.type in rule.data_source_types 

127 ] 

128 

129 if matching_sources: 

130 # Sort by priority 

131 matching_sources.sort(key=lambda s: s.priority, reverse=True) 

132 

133 return RoutingDecision( 

134 query=features.text, 

135 data_sources=matching_sources, 

136 strategy=rule.strategy, 

137 confidence=0.9, # High confidence for rule-based 

138 reasoning=f"Matched rule: {rule.description or rule.name}", 

139 features=features, 

140 metadata={"rule": rule.name}, 

141 ) 

142 

143 # Fallback: use first available source with default strategy 

144 if available_sources: 

145 # Prefer vector stores for general queries 

146 vector_stores = [ 

147 s for s in available_sources if s.type == DataSourceType.VECTOR_STORE 

148 ] 

149 

150 fallback_sources = vector_stores or available_sources 

151 fallback_sources.sort(key=lambda s: s.priority, reverse=True) 

152 

153 return RoutingDecision( 

154 query=features.text, 

155 data_sources=[fallback_sources[0]], 

156 strategy="dense_search", 

157 confidence=0.5, 

158 reasoning="No matching rules, using default fallback", 

159 features=features, 

160 metadata={"fallback": True}, 

161 ) 

162 

163 # No sources available 

164 return RoutingDecision( 

165 query=features.text, 

166 data_sources=[], 

167 strategy="none", 

168 confidence=0.0, 

169 reasoning="No data sources available", 

170 features=features, 

171 metadata={"error": "no_sources"}, 

172 ) 

173 

174 def _load_default_rules(self) -> None: 

175 """Load default routing rules.""" 

176 

177 # Rule 1: Multimodal queries with images 

178 self.add_rule( 

179 RoutingRule( 

180 name="multimodal_image", 

181 condition=lambda f: Modality.IMAGE in f.modalities, 

182 data_source_types=[DataSourceType.MULTIMODAL_STORE], 

183 strategy="multimodal", 

184 priority=100, 

185 description="Route image queries to multimodal store", 

186 ), 

187 ) 

188 

189 # Rule 2: Multimodal queries with video 

190 self.add_rule( 

191 RoutingRule( 

192 name="multimodal_video", 

193 condition=lambda f: Modality.VIDEO in f.modalities, 

194 data_source_types=[DataSourceType.MULTIMODAL_STORE], 

195 strategy="multimodal", 

196 priority=95, 

197 description="Route video queries to multimodal store", 

198 ), 

199 ) 

200 

201 # Rule 3: Multimodal queries with audio 

202 self.add_rule( 

203 RoutingRule( 

204 name="multimodal_audio", 

205 condition=lambda f: Modality.AUDIO in f.modalities, 

206 data_source_types=[DataSourceType.MULTIMODAL_STORE], 

207 strategy="multimodal", 

208 priority=90, 

209 description="Route audio queries to multimodal store", 

210 ), 

211 ) 

212 

213 # Rule 4: Knowledge graph for analytical queries 

214 self.add_rule( 

215 RoutingRule( 

216 name="analytical_graph", 

217 condition=lambda f: f.intent == QueryIntent.ANALYTICAL, 

218 data_source_types=[ 

219 DataSourceType.KNOWLEDGE_GRAPH, 

220 DataSourceType.VECTOR_STORE, 

221 ], 

222 strategy="hybrid", 

223 priority=80, 

224 description="Route analytical queries to knowledge graph + vector store", 

225 ), 

226 ) 

227 

228 # Rule 5: Keyword search for navigational queries 

229 self.add_rule( 

230 RoutingRule( 

231 name="navigational_keyword", 

232 condition=lambda f: f.intent == QueryIntent.NAVIGATIONAL, 

233 data_source_types=[ 

234 DataSourceType.KEYWORD_INDEX, 

235 DataSourceType.VECTOR_STORE, 

236 ], 

237 strategy="sparse", 

238 priority=70, 

239 description="Route navigational queries to keyword index", 

240 ), 

241 ) 

242 

243 # Rule 6: SQL database for structured queries 

244 self.add_rule( 

245 RoutingRule( 

246 name="structured_sql", 

247 condition=lambda f: ( 

248 any(kw in f.keywords for kw in ["count", "total", "average", "sum"]) 

249 or "data" in f.domain 

250 if f.domain 

251 else False 

252 ), 

253 data_source_types=[ 

254 DataSourceType.SQL_DATABASE, 

255 DataSourceType.VECTOR_STORE, 

256 ], 

257 strategy="structured", 

258 priority=60, 

259 description="Route structured queries to SQL database", 

260 ), 

261 ) 

262 

263 # Rule 7: Keyword-rich queries use sparse retrieval 

264 self.add_rule( 

265 RoutingRule( 

266 name="keyword_rich", 

267 condition=lambda f: len(f.keywords) > 7, 

268 data_source_types=[ 

269 DataSourceType.KEYWORD_INDEX, 

270 DataSourceType.VECTOR_STORE, 

271 ], 

272 strategy="sparse", 

273 priority=50, 

274 description="Route keyword-rich queries to keyword index", 

275 ), 

276 ) 

277 

278 # Rule 8: Technical domain prefers vector stores 

279 self.add_rule( 

280 RoutingRule( 

281 name="technical_vector", 

282 condition=lambda f: f.domain == "technical", 

283 data_source_types=[DataSourceType.VECTOR_STORE], 

284 strategy="dense", 

285 priority=40, 

286 description="Route technical queries to vector store", 

287 ), 

288 ) 

289 

290 # Rule 9: Long queries use dense retrieval 

291 self.add_rule( 

292 RoutingRule( 

293 name="long_dense", 

294 condition=lambda f: f.is_long, 

295 data_source_types=[DataSourceType.VECTOR_STORE], 

296 strategy="dense", 

297 priority=30, 

298 description="Route long queries to dense retrieval", 

299 ), 

300 ) 

301 

302 # Rule 10: Complex queries use hybrid search 

303 self.add_rule( 

304 RoutingRule( 

305 name="complex_hybrid", 

306 condition=lambda f: f.is_complex, 

307 data_source_types=[ 

308 DataSourceType.VECTOR_STORE, 

309 DataSourceType.KEYWORD_INDEX, 

310 ], 

311 strategy="hybrid", 

312 priority=20, 

313 description="Route complex queries to hybrid search", 

314 ), 

315 ) 

316 

317 # Rule 11: Simple factual queries use vector store 

318 self.add_rule( 

319 RoutingRule( 

320 name="simple_factual", 

321 condition=lambda f: f.intent == QueryIntent.FACTUAL and f.is_simple, 

322 data_source_types=[DataSourceType.VECTOR_STORE], 

323 strategy="dense", 

324 priority=10, 

325 description="Route simple factual queries to vector store", 

326 ), 

327 )