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 )