« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Semantic routing strategy using embeddings."""
3from __future__ import annotations
5from collections.abc import Callable
6from dataclasses import dataclass, field
8import numpy as np
10from lexigram.ai.rag.routing.types import (
11 DataSource,
12 DataSourceType,
13 QueryFeatures,
14 RoutingDecision,
15)
18@dataclass
19class RoutingPattern:
20 """A routing pattern for semantic routing.
22 Attributes:
23 name: Unique identifier for the pattern.
24 examples: Example queries that match this pattern.
25 data_source_types: Preferred data source types for this pattern.
26 strategy: Retrieval strategy to use for this pattern.
27 description: Human-readable description.
28 embedding: Pre-computed embedding of the pattern (computed from examples).
29 """
31 name: str
32 examples: list[str]
33 data_source_types: list[DataSourceType]
34 strategy: str
35 description: str = ""
36 embedding: np.ndarray | None = field(default=None, repr=False)
39class SemanticRouter:
40 """Semantic routing strategy using embedding similarity.
42 Routes queries by comparing query embeddings to pre-defined routing
43 patterns and selecting the most similar pattern.
45 Example:
46 ```python
47 from lexigram.ai.rag import SemanticRouter, RoutingPattern
49 router = SemanticRouter(embed_fn=my_embed_function)
51 # Add routing pattern
52 router.add_pattern(RoutingPattern(
53 name="technical_docs",
54 examples=[
55 "How do I configure authentication?",
56 "API reference for user management",
57 "What are the deployment steps?"
58 ],
59 data_source_types=[DataSourceType.VECTOR_STORE],
60 strategy="dense",
61 description="Technical documentation queries"
62 ))
64 # Route query
65 decision = await router.route(features, available_sources)
66 ```
67 """
69 def __init__(
70 self,
71 *,
72 embed_fn: Callable | None = None,
73 similarity_threshold: float = 0.7,
74 ):
75 """Initialize the semantic router.
77 Args:
78 embed_fn: Function to embed text (async callable).
79 similarity_threshold: Minimum similarity for pattern matching.
81 Use `with_defaults()` classmethod to create a router with default patterns.
82 """
83 self.embed_fn = embed_fn
84 self.similarity_threshold = similarity_threshold
85 self.patterns: list[RoutingPattern] = []
87 @classmethod
88 def with_defaults(cls, *args, **kwargs) -> SemanticRouter:
89 """Create a router pre-populated with default routing patterns.
91 Args:
92 *args: Positional arguments passed to __init__.
93 **kwargs: Keyword arguments passed to __init__.
95 Returns:
96 A SemanticRouter with all default patterns registered.
97 """
98 instance = cls(*args, **kwargs)
99 instance._load_default_patterns()
100 return instance
102 async def add_pattern(self, pattern: RoutingPattern) -> None:
103 """Add a routing pattern and compute its embedding.
105 Args:
106 pattern: Routing pattern to add.
107 """
108 # Compute pattern embedding (average of example embeddings)
109 if self.embed_fn and pattern.embedding is None:
110 example_embeddings = []
111 for example in pattern.examples:
112 embedding = await self.embed_fn(example)
113 example_embeddings.append(embedding)
115 # Average embeddings to create pattern embedding
116 pattern.embedding = np.mean(example_embeddings, axis=0)
118 self.patterns.append(pattern)
120 def remove_pattern(self, name: str) -> bool:
121 """Remove a routing pattern by name.
123 Args:
124 name: Name of the pattern to remove.
126 Returns:
127 True if pattern was found and removed, False otherwise.
128 """
129 initial_count = len(self.patterns)
130 self.patterns = list(filter(lambda p: p.name != name, self.patterns))
131 return len(self.patterns) < initial_count
133 async def route(
134 self,
135 features: QueryFeatures,
136 available_sources: list[DataSource],
137 ) -> RoutingDecision:
138 """Route query using semantic similarity.
140 Args:
141 features: Extracted query features.
142 available_sources: List of available data sources.
144 Returns:
145 Routing decision based on most similar pattern.
146 """
147 # Need embedding function to route
148 if not self.embed_fn:
149 return self._fallback_routing(
150 features,
151 available_sources,
152 "No embedding function configured",
153 )
155 # Compute query embedding
156 try:
157 query_embedding = await self.embed_fn(features.text)
158 except (ConnectionError, TimeoutError, RuntimeError, ValueError, OSError) as e:
159 return self._fallback_routing(
160 features,
161 available_sources,
162 f"Failed to embed query: {e}",
163 )
165 # Find most similar pattern
166 best_pattern = None
167 best_similarity = -1.0
169 for pattern in self.patterns:
170 if pattern.embedding is None:
171 continue
173 # Compute cosine similarity
174 similarity = self._cosine_similarity(query_embedding, pattern.embedding)
176 if similarity > best_similarity:
177 best_similarity = similarity
178 best_pattern = pattern
180 # Check if similarity meets threshold
181 if best_pattern and best_similarity >= self.similarity_threshold:
182 # Find matching data sources
183 matching_sources = [
184 source
185 for source in available_sources
186 if source.type in best_pattern.data_source_types
187 ]
189 if matching_sources:
190 matching_sources.sort(key=lambda s: s.priority, reverse=True)
192 return RoutingDecision(
193 query=features.text,
194 data_sources=matching_sources,
195 strategy=best_pattern.strategy,
196 confidence=float(best_similarity),
197 reasoning=f"Matched pattern: {best_pattern.description or best_pattern.name} (similarity: {best_similarity:.3f})",
198 features=features,
199 metadata={
200 "pattern": best_pattern.name,
201 "similarity": float(best_similarity),
202 },
203 )
205 # Fallback if no pattern matched
206 return self._fallback_routing(
207 features,
208 available_sources,
209 f"No pattern above threshold {self.similarity_threshold} (best: {best_similarity:.3f})",
210 )
212 def _fallback_routing(
213 self,
214 features: QueryFeatures,
215 available_sources: list[DataSource],
216 reason: str,
217 ) -> RoutingDecision:
218 """Fallback routing when no pattern matches.
220 Args:
221 features: Query features.
222 available_sources: Available data sources.
223 reason: Reason for fallback.
225 Returns:
226 Fallback routing decision.
227 """
228 if available_sources:
229 # Prefer vector stores
230 vector_stores = [
231 s for s in available_sources if s.type == DataSourceType.VECTOR_STORE
232 ]
234 fallback_sources = vector_stores or available_sources
235 fallback_sources.sort(key=lambda s: s.priority, reverse=True)
237 return RoutingDecision(
238 query=features.text,
239 data_sources=[fallback_sources[0]],
240 strategy="dense",
241 confidence=0.3,
242 reasoning=f"Fallback routing: {reason}",
243 features=features,
244 metadata={"fallback": True, "reason": reason},
245 )
247 return RoutingDecision(
248 query=features.text,
249 data_sources=[],
250 strategy="none",
251 confidence=0.0,
252 reasoning=f"No sources available: {reason}",
253 features=features,
254 metadata={"error": "no_sources", "reason": reason},
255 )
257 def _cosine_similarity(self, a: np.ndarray, b: np.ndarray) -> float:
258 """Compute cosine similarity between two vectors.
260 Args:
261 a: First vector.
262 b: Second vector.
264 Returns:
265 Cosine similarity (-1 to 1).
266 """
267 dot_product = np.dot(a, b)
268 norm_a = np.linalg.norm(a)
269 norm_b = np.linalg.norm(b)
271 if norm_a == 0 or norm_b == 0:
272 return 0.0
274 return float(dot_product / (norm_a * norm_b))
276 def _load_default_patterns(self) -> None:
277 """Load default routing patterns.
279 Note: Embeddings will be computed when embed_fn is available.
280 """
281 # Pattern 1: Technical documentation
282 self.patterns.append(
283 RoutingPattern(
284 name="technical_docs",
285 examples=[
286 "How do I configure authentication?",
287 "API reference for user management",
288 "What are the deployment steps?",
289 "How to setup the database connection?",
290 ],
291 data_source_types=[DataSourceType.VECTOR_STORE],
292 strategy="dense",
293 description="Technical documentation queries",
294 ),
295 )
297 # Pattern 2: General Q&A
298 self.patterns.append(
299 RoutingPattern(
300 name="general_qa",
301 examples=[
302 "What is the capital of France?",
303 "Who invented the telephone?",
304 "When did World War II end?",
305 "What is photosynthesis?",
306 ],
307 data_source_types=[DataSourceType.VECTOR_STORE],
308 strategy="dense",
309 description="General knowledge questions",
310 ),
311 )
313 # Pattern 3: Multi-modal content
314 self.patterns.append(
315 RoutingPattern(
316 name="multimodal_content",
317 examples=[
318 "Show me images of sunset beaches",
319 "Find videos about machine learning",
320 "Pictures of classic cars",
321 "Audio recordings of bird songs",
322 ],
323 data_source_types=[DataSourceType.MULTIMODAL_STORE],
324 strategy="multimodal",
325 description="Multi-modal content queries",
326 ),
327 )
329 # Pattern 4: Data analysis
330 self.patterns.append(
331 RoutingPattern(
332 name="data_analysis",
333 examples=[
334 "What is the average sales by region?",
335 "Count total orders in Q4",
336 "Show revenue trends over time",
337 "Compare performance metrics",
338 ],
339 data_source_types=[
340 DataSourceType.SQL_DATABASE,
341 DataSourceType.VECTOR_STORE,
342 ],
343 strategy="structured",
344 description="Data analysis and aggregation queries",
345 ),
346 )
348 # Pattern 5: Navigation and search
349 self.patterns.append(
350 RoutingPattern(
351 name="navigation",
352 examples=[
353 "Find the pricing page",
354 "Locate documentation for API",
355 "Where is the user guide?",
356 "Link to terms of service",
357 ],
358 data_source_types=[DataSourceType.KEYWORD_INDEX],
359 strategy="sparse",
360 description="Navigation and page finding queries",
361 ),
362 )