Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/routing/strategies/llm.py: 26%
69 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""LLM-based routing strategy."""
3from __future__ import annotations
5from collections.abc import Callable
6from typing import Any
8from lexigram.ai.rag.routing.types import (
9 DataSource,
10 DataSourceType,
11 QueryFeatures,
12 RoutingDecision,
13)
14from lexigram.contracts.ai import LLMClientProtocol
15from lexigram.contracts.ai.llm import ChatMessage, Role
16from lexigram.di.decorators import inject
17from lexigram.logging import (
18 get_logger,
19)
20from lexigram.serialization import JSONDecodeError, loads
22logger = get_logger(__name__)
25@inject
26class LLMRouter:
27 """LLM-based routing strategy using language models.
29 Uses an LLM to analyze queries and make routing decisions based on
30 understanding of query intent and data source capabilities.
32 Example:
33 ```python
34 from lexigram.ai.rag import LLMRouter
36 async def my_llm_call(prompt):
37 # Call your LLM
38 return llm.generate(prompt)
40 router = LLMRouter(llm_client=my_llm_client)
41 decision = await router.route(features, available_sources)
42 ```
43 """
45 ROUTING_PROMPT_TEMPLATE = """You are a query routing expert. Analyze the following query and determine which data sources should be used.
47Query: {query}
49Query Features:
50- Intent: {intent}
51- Keywords: {keywords}
52- Domain: {domain}
53- Complexity: {complexity}
54- Modalities: {modalities}
56Available Data Sources:
57{data_sources}
59Instructions:
601. Choose the most appropriate data source(s) from the list above
612. Select a retrieval strategy: dense, sparse, hybrid, multimodal, structured, graph
623. Provide a confidence score (0-1)
634. Explain your reasoning
65Respond with a JSON object:
66{{
67 "data_source_names": ["name1", "name2"],
68 "strategy": "dense",
69 "confidence": 0.9,
70 "reasoning": "explanation here"
71}}
73Response:"""
75 def __init__(
76 self,
77 *,
78 llm_client: LLMClientProtocol | None = None,
79 llm_fn: Callable[[str], Any] | None = None,
80 temperature: float = 0.1,
81 max_tokens: int = 500,
82 ):
83 """Initialize the LLM router.
85 Args:
86 llm_client: Platform LLM client to use for routing.
87 llm_fn: Async function to call LLM (fallback if no client).
88 temperature: LLM temperature for routing decisions.
89 max_tokens: Maximum tokens for LLM response.
90 """
91 self.llm_client = llm_client
92 self.llm_fn = llm_fn
93 self.temperature = temperature
94 self.max_tokens = max_tokens
96 async def route(
97 self,
98 features: QueryFeatures,
99 available_sources: list[DataSource],
100 ) -> RoutingDecision:
101 """Route query using LLM-based decision making.
103 Args:
104 features: Extracted query features.
105 available_sources: List of available data sources.
107 Returns:
108 Routing decision from LLM analysis.
109 """
110 if not self.llm_client and not self.llm_fn:
111 return self._fallback_routing(
112 features,
113 available_sources,
114 "No LLM client or function configured",
115 )
117 if not available_sources:
118 return RoutingDecision(
119 query=features.text,
120 data_sources=[],
121 strategy="none",
122 confidence=0.0,
123 reasoning="No data sources available",
124 features=features,
125 metadata={"error": "no_sources"},
126 )
128 # Build prompt
129 prompt = self._build_prompt(features, available_sources)
131 # Call LLM
132 try:
133 if self.llm_client:
134 # Use platform LLM client
135 messages = [
136 ChatMessage(
137 role=Role.USER,
138 content=prompt,
139 ),
140 ]
141 response_obj = await self.llm_client.complete(
142 messages=messages,
143 temperature=self.temperature,
144 max_tokens=self.max_tokens,
145 )
146 # Unwrap Result
147 if hasattr(response_obj, "is_err"):
148 if response_obj.is_err():
149 raise response_obj.unwrap_err()
150 response_obj = response_obj.unwrap() # type: ignore[assignment]
151 # Handle both Completion object and raw string
152 if hasattr(response_obj, "content"):
153 response = response_obj.content
154 elif hasattr(response_obj, "text"):
155 response = response_obj.text
156 else:
157 response = str(response_obj)
158 else:
159 # Use fallback function
160 llm_fn = self.llm_fn
161 assert llm_fn is not None # noqa: S101 # fallback branch entered only when llm_fn set
162 response = await llm_fn(prompt)
164 # Parse response
165 routing_data = self._parse_llm_response(response)
167 # Find selected data sources
168 selected_sources = [
169 source
170 for source in available_sources
171 if source.name in routing_data.get("data_source_names", [])
172 ]
174 if not selected_sources:
175 # Use first available if LLM didn't select valid sources
176 selected_sources = [available_sources[0]]
177 routing_data["reasoning"] += " (Using fallback source)"
179 return RoutingDecision(
180 query=features.text,
181 data_sources=selected_sources,
182 strategy=routing_data.get("strategy", "dense"),
183 confidence=routing_data.get("confidence", 0.5),
184 reasoning=routing_data.get("reasoning", "LLM routing decision"),
185 features=features,
186 metadata={"llm_routing": True},
187 )
189 except (ConnectionError, TimeoutError, RuntimeError, ValueError, OSError) as e:
190 return self._fallback_routing(
191 features,
192 available_sources,
193 f"LLM routing failed: {e}",
194 )
196 def _build_prompt(
197 self,
198 features: QueryFeatures,
199 available_sources: list[DataSource],
200 ) -> str:
201 """Build prompt for LLM routing.
203 Args:
204 features: Query features.
205 available_sources: Available data sources.
207 Returns:
208 Formatted prompt for LLM.
209 """
210 # Format data sources
211 sources_text = "\n".join(
212 [
213 f"- {source.name}: {source.description} (type: {source.type.value}, capabilities: {', '.join(source.capabilities)})"
214 for source in available_sources
215 ],
216 )
218 # Format modalities
219 modalities_text = ", ".join([m.value for m in features.modalities])
221 return self.ROUTING_PROMPT_TEMPLATE.format(
222 query=features.text,
223 intent=features.intent.value,
224 keywords=", ".join(features.keywords) if features.keywords else "none",
225 domain=features.domain or "general",
226 complexity=f"{features.complexity:.2f}",
227 modalities=modalities_text,
228 data_sources=sources_text,
229 )
231 def _parse_llm_response(self, response: str) -> dict:
232 """Parse LLM response to extract routing decision.
234 Args:
235 response: LLM response text.
237 Returns:
238 Parsed routing data.
239 """
240 # Try to extract JSON from response
241 try:
242 # Look for JSON in response
243 start = response.find("{")
244 end = response.rfind("}") + 1
246 if start != -1 and end > start:
247 json_str = response[start:end]
248 return loads(json_str)
249 except (JSONDecodeError, ValueError, TypeError) as e:
250 logger.debug("Failed to parse LLM response: %s", e)
251 # Continue with fallback
253 # Fallback: return default
254 return {
255 "data_source_names": [],
256 "strategy": "dense",
257 "confidence": 0.3,
258 "reasoning": "Failed to parse LLM response",
259 }
261 def _fallback_routing(
262 self,
263 features: QueryFeatures,
264 available_sources: list[DataSource],
265 reason: str,
266 ) -> RoutingDecision:
267 """Fallback routing when LLM fails.
269 Args:
270 features: Query features.
271 available_sources: Available data sources.
272 reason: Reason for fallback.
274 Returns:
275 Fallback routing decision.
276 """
277 if available_sources:
278 vector_stores = [
279 s for s in available_sources if s.type == DataSourceType.VECTOR_STORE
280 ]
282 fallback_sources = vector_stores or available_sources
283 fallback_sources.sort(key=lambda s: s.priority, reverse=True)
285 return RoutingDecision(
286 query=features.text,
287 data_sources=[fallback_sources[0]],
288 strategy="dense",
289 confidence=0.3,
290 reasoning=f"LLM fallback: {reason}",
291 features=features,
292 metadata={"fallback": True, "reason": reason},
293 )
295 return RoutingDecision(
296 query=features.text,
297 data_sources=[],
298 strategy="none",
299 confidence=0.0,
300 reasoning=f"No sources available: {reason}",
301 features=features,
302 metadata={"error": "no_sources", "reason": reason},
303 )