Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-llm/src/lexigram/ai/llm/selection/core.py: 97%
110 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Model selection and routing for intelligent LLM usage.
3This module provides intelligent model selection based on task characteristics,
4allowing you to optimize for cost, quality, and latency.
6Example:
7 >>> from lexigram.ai.llm import ModelSelector, SelectionStrategy
8 >>>
9 >>> selector = ModelSelector(
10 ... default_model="gpt-3.5-turbo",
11 ... strategies=[
12 ... SelectionStrategy(
13 ... name="complex",
14 ... model="gpt-4-turbo",
15 ... conditions={"min_tokens": 1000}
16 ... )
17 ... ],
18 ... fallback_chain=["gpt-4-turbo", "gpt-3.5-turbo", "ollama/llama3"]
19 ... )
20 >>>
21 >>> # Automatically selects appropriate model
22 >>> model = selector.select("Write a complex analysis...")
23"""
25from __future__ import annotations
27from dataclasses import dataclass
28from enum import StrEnum
29from typing import Any
31from lexigram.ai.llm.selection._capabilities import (
32 DEFAULT_MODEL_CAPABILITIES,
33 ModelCapabilities,
34)
35from lexigram.contracts.ai.llm import TokenCounterProtocol
36from lexigram.domain import DomainModel
37from lexigram.validation import Field
39__all__ = [
40 "DEFAULT_MODEL_CAPABILITIES",
41 "ModelCapabilities",
42 "ModelSelector",
43 "SelectionCriteria",
44 "SelectionStrategy",
45 "create_balanced_selector",
46 "create_cost_optimized_selector",
47 "create_quality_optimized_selector",
48]
51class SelectionCriteria(StrEnum):
52 """Criteria for model selection."""
54 TOKEN_COUNT = "token_count" # noqa: S105 # criterion name, not a credential
55 COST = "cost"
56 LATENCY = "latency"
57 QUALITY = "quality"
58 CUSTOM = "custom"
61@dataclass(init=False)
62class SelectionStrategy(DomainModel):
63 """Strategy for selecting models based on conditions.
65 Example:
66 >>> strategy = SelectionStrategy(
67 ... name="long_context",
68 ... model="gpt-4-turbo-preview",
69 ... conditions={
70 ... "min_tokens": 2000,
71 ... "max_tokens": 100000
72 ... }
73 ... )
74 """
76 name: str = Field(..., description="Strategy name")
77 model: str = Field(..., description="Model to use for this strategy")
78 conditions: dict[str, Any] = Field(
79 default_factory=dict,
80 description="Conditions that trigger this strategy",
81 )
82 priority: int = Field(
83 default=0,
84 description="Priority (higher = evaluated first)",
85 )
86 description: str | None = Field(
87 None,
88 description="Human-readable description",
89 )
91 def matches(self, context: dict[str, Any]) -> bool:
92 """Check if this strategy matches the given context.
94 Args:
95 context: Context dictionary with prompt info
97 Returns:
98 True if all conditions are met
100 Example:
101 >>> context = {"token_count": 2500, "has_code": True}
102 >>> strategy.matches(context)
103 True
104 """
105 for key, value in self.conditions.items():
106 # Handle different condition types
107 if key.startswith("min_"):
108 actual_key = key[4:] # Remove "min_" prefix
109 if actual_key not in context:
110 return False
111 if context[actual_key] < value:
112 return False
113 elif key.startswith("max_"):
114 actual_key = key[4:] # Remove "max_" prefix
115 if actual_key not in context:
116 return False
117 if context[actual_key] > value:
118 return False
119 elif key.startswith("has_"):
120 # Boolean flag check
121 if key not in context:
122 return False
123 if context[key] != value:
124 return False
125 else:
126 # Exact match
127 if key not in context:
128 return False
129 if context[key] != value:
130 return False
132 return True
135class ModelSelector:
136 """Intelligent model selector with fallback support.
138 Automatically selects the best model based on prompt characteristics
139 and provides fallback chains for reliability.
141 Example:
142 >>> selector = ModelSelector(
143 ... default_model="gpt-3.5-turbo",
144 ... strategies=[
145 ... SelectionStrategy(
146 ... name="complex",
147 ... model="gpt-4-turbo",
148 ... conditions={"min_tokens": 1000}
149 ... ),
150 ... SelectionStrategy(
151 ... name="simple",
152 ... model="claude-3-haiku-20240307",
153 ... conditions={"max_tokens": 500}
154 ... )
155 ... ],
156 ... fallback_chain=["gpt-4-turbo", "gpt-3.5-turbo"]
157 ... )
158 >>>
159 >>> # Select model for a prompt
160 >>> model = selector.select("Long prompt here...")
161 >>> print(model)
162 'gpt-4-turbo'
163 >>>
164 >>> # Get next fallback on error
165 >>> fallback = selector.get_fallback("gpt-4-turbo")
166 >>> print(fallback)
167 'gpt-3.5-turbo'
168 """
170 def __init__(
171 self,
172 default_model: str | None = None,
173 strategies: list[SelectionStrategy] | None = None,
174 fallback_chain: list[str] | None = None,
175 model_capabilities: dict[str, ModelCapabilities] | None = None,
176 token_counter: TokenCounterProtocol | None = None,
177 ):
178 """Initialize model selector.
180 Args:
181 default_model: Default model to use
182 strategies: List of selection strategies
183 fallback_chain: Ordered list of fallback models
184 model_capabilities: Custom model capabilities
185 token_counter: Token counter for prompt analysis
187 Example:
188 >>> selector = ModelSelector(
189 ... default_model="gpt-3.5-turbo",
190 ... fallback_chain=["gpt-4", "claude-3-sonnet-20240229"]
191 ... )
192 """
193 if default_model is None:
194 default_model = "gpt-3.5-turbo"
196 self.default_model = default_model
197 self.strategies = sorted(
198 strategies or [],
199 key=lambda s: s.priority,
200 reverse=True,
201 )
202 self.fallback_chain = fallback_chain or [default_model]
203 self.model_capabilities = model_capabilities or DEFAULT_MODEL_CAPABILITIES
204 self.token_counter = token_counter
206 def select(
207 self,
208 prompt: str,
209 context: dict[str, Any] | None = None,
210 required_capabilities: list[str] | None = None,
211 ) -> str:
212 """Select the best model for the given prompt.
214 Args:
215 prompt: The prompt text
216 context: Additional context for selection
217 required_capabilities: Required capabilities (e.g., ["supports_functions"])
219 Returns:
220 Selected model name
222 Example:
223 >>> model = selector.select(
224 ... "Analyze this image...",
225 ... required_capabilities=["supports_vision"]
226 ... )
227 >>> print(model)
228 'gpt-4-turbo'
229 """
230 # Build context
231 ctx = self._build_context(prompt, context or {})
233 # Filter by required capabilities
234 available_models = self._filter_by_capabilities(required_capabilities)
236 # Try each strategy in priority order
237 for strategy in self.strategies:
238 if strategy.model not in available_models:
239 continue
241 if strategy.matches(ctx):
242 return strategy.model
244 # Return default if no strategy matched
245 return (
246 self.default_model
247 if self.default_model in available_models
248 else available_models[0]
249 )
251 def get_fallback(self, failed_model: str) -> str | None:
252 """Get the next model in the fallback chain.
254 Args:
255 failed_model: The model that failed
257 Returns:
258 Next fallback model, or None if no fallback available
260 Example:
261 >>> fallback = selector.get_fallback("gpt-4-turbo")
262 >>> print(fallback)
263 'gpt-3.5-turbo'
264 """
265 try:
266 idx = self.fallback_chain.index(failed_model)
267 if idx + 1 < len(self.fallback_chain):
268 return self.fallback_chain[idx + 1]
269 except ValueError:
270 # Model not in fallback chain, return first fallback
271 if self.fallback_chain:
272 return self.fallback_chain[0]
274 return None
276 def get_capabilities(self, model: str) -> ModelCapabilities | None:
277 """Get capabilities for a model.
279 Args:
280 model: Model name
282 Returns:
283 Model capabilities or None if unknown
285 Example:
286 >>> caps = selector.get_capabilities("gpt-4-turbo")
287 >>> print(caps.max_tokens)
288 128000
289 """
290 return self.model_capabilities.get(model)
292 def estimate_cost(
293 self,
294 model: str,
295 input_tokens: int,
296 output_tokens: int,
297 ) -> float:
298 """Estimate cost for a model call.
300 Args:
301 model: Model name
302 input_tokens: Number of input tokens
303 output_tokens: Number of output tokens
305 Returns:
306 Estimated cost in USD
308 Example:
309 >>> cost = selector.estimate_cost("gpt-4-turbo", 1000, 500)
310 >>> print(f"${cost:.4f}")
311 $0.0250
312 """
313 caps = self.get_capabilities(model)
314 if not caps:
315 return 0.0
317 input_cost = (input_tokens / 1000) * caps.cost_per_1k_input
318 output_cost = (output_tokens / 1000) * caps.cost_per_1k_output
320 return input_cost + output_cost
322 def _build_context(self, prompt: str, context: dict[str, Any]) -> dict[str, Any]:
323 """Build context for strategy matching.
325 Args:
326 prompt: The prompt text
327 context: User-provided context
329 Returns:
330 Complete context dictionary
331 """
332 ctx = context.copy()
334 # Add token count if not provided
335 if "tokens" not in ctx and "token_count" not in ctx:
336 # Use simple estimation to avoid async call
337 # Count tokens as roughly 4 characters per token
338 estimated_tokens = len(prompt) // 4 + 1
339 ctx["tokens"] = estimated_tokens
340 ctx["token_count"] = estimated_tokens
342 # Detect prompt characteristics
343 ctx["has_code"] = ctx.get(
344 "has_code",
345 "```" in prompt or "def " in prompt or "function " in prompt,
346 )
347 ctx["has_url"] = ctx.get("has_url", "http://" in prompt or "https://" in prompt)
348 ctx["is_question"] = ctx.get("is_question", "?" in prompt)
349 ctx["prompt_length"] = len(prompt)
351 return ctx
353 def _filter_by_capabilities(
354 self,
355 required_capabilities: list[str] | None,
356 ) -> list[str]:
357 """Filter models by required capabilities.
359 Args:
360 required_capabilities: List of required capability names
362 Returns:
363 List of model names that meet requirements
364 """
365 if not required_capabilities:
366 return list(self.model_capabilities.keys())
368 available = []
369 for model, caps in self.model_capabilities.items():
370 if all(getattr(caps, cap, False) for cap in required_capabilities):
371 available.append(model)
373 return available
376def create_cost_optimized_selector(
377 budget_per_1k_tokens: float = 2.0,
378) -> ModelSelector:
379 """Create a cost-optimized model selector."""
380 from lexigram.ai.llm.selection._scoring import (
381 create_cost_optimized_selector as _create,
382 )
384 return _create(budget_per_1k_tokens=budget_per_1k_tokens)
387def create_quality_optimized_selector() -> ModelSelector:
388 """Create a quality-optimized model selector."""
389 from lexigram.ai.llm.selection._scoring import (
390 create_quality_optimized_selector as _create,
391 )
393 return _create()
396def create_balanced_selector() -> ModelSelector:
397 """Create a balanced model selector."""
398 from lexigram.ai.llm.selection._scoring import create_balanced_selector as _create
400 return _create()