Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-llm/src/lexigram/ai/llm/routing/orchestrator.py: 28%

1"""Intelligent LLM orchestration and routing.""" 

2 

3from __future__ import annotations 

4 

5from dataclasses import dataclass 

6from typing import TYPE_CHECKING 

7 

8from lexigram.ai.llm.exceptions import LLMError 

9from lexigram.ai.llm.types import ChatMessage, Role 

10from lexigram.contracts.ai.models import ModelRequest, ModelResponse 

11from lexigram.contracts.ai.providers import ModelCapability, SelectionStrategy 

12from lexigram.logging import ( 

13 get_logger, 

14) 

15from lexigram.result import Err, Ok, Result 

16 

17if TYPE_CHECKING: 

18 from lexigram.ai.llm.registry.core import ProviderRegistry 

19 

20logger = get_logger(__name__) 

21 

22 

23class OrchestratorError(LLMError): 

24 """Base orchestration error.""" 

25 

26 _code: str = "LEX_ERR_LLM_025" 

27 

28 

29class NoSuitableModelError(OrchestratorError): 

30 """No model found that meets the requirements.""" 

31 

32 _code: str = "LEX_ERR_LLM_026" 

33 

34 

35@dataclass(frozen=True) 

36class _ModelSelection: 

37 """Internal result of model selection. 

38 

39 Attributes: 

40 provider: Provider name key in the registry. 

41 model_id: Selected model identifier. 

42 """ 

43 

44 provider: str 

45 model_id: str 

46 

47 

48class LLMOrchestrator: 

49 """Intelligent routing and orchestration of multiple LLM providers. 

50 

51 Selects the optimal provider/model for a given request based on: 

52 - Required capabilities 

53 - Token/cost constraints 

54 - Latency requirements 

55 - Provider availability 

56 

57 The selection strategy is read from ``request.extra_params["strategy"]`` 

58 (a :class:`~lexigram.contracts.ai.providers.SelectionStrategy` value or 

59 its string equivalent). Defaults to ``CAPABILITY_MATCH``. 

60 """ 

61 

62 def __init__(self, registry: ProviderRegistry) -> None: 

63 """Initialize orchestrator with a provider registry. 

64 

65 Args: 

66 registry: The provider registry containing all models. 

67 """ 

68 self.registry = registry 

69 self._round_robin_index: int = 0 

70 

71 async def execute( 

72 self, request: ModelRequest 

73 ) -> Result[ModelResponse, OrchestratorError]: 

74 """Execute a request using the best available model. 

75 

76 Args: 

77 request: The model request. 

78 

79 Returns: 

80 ``Ok(ModelResponse)`` on success, ``Err(OrchestratorError)`` on failure. 

81 """ 

82 selection = self._select_model(request) 

83 if not selection: 

84 error: OrchestratorError = NoSuitableModelError( 

85 f"No model found matching requirements: {request.model_id!r}" 

86 ) 

87 logger.error( 

88 "no_suitable_model", 

89 required_capabilities=sorted( 

90 c.value for c in (request.required_capabilities or set()) 

91 ), 

92 ) 

93 return Err(error) 

94 

95 try: 

96 self.registry.get_provider(selection.provider) 

97 except KeyError as e: 

98 logger.exception("provider_not_registered", provider=selection.provider) 

99 raise OrchestratorError( 

100 f"Provider '{selection.provider}' not registered" 

101 ) from e 

102 

103 logger.info( 

104 "executing_request", 

105 provider=selection.provider, 

106 model_id=selection.model_id, 

107 ) 

108 

109 # Client must be supplied pre-configured via the DI container. Callers can 

110 # pass a ready-made client in ``request.extra_params["client"]`` for testing 

111 # or for use cases where instantiation is managed externally. 

112 client = request.extra_params.get("client") 

113 if client is None: 

114 error = OrchestratorError( 

115 "No pre-configured client found in request.extra_params['client']. " 

116 "Resolve the client through the DI container before calling execute()." 

117 ) 

118 logger.error("no_client_in_request", provider=selection.provider) 

119 return Err(error) 

120 

121 messages = [ChatMessage(role=Role.USER, content=request.prompt)] 

122 try: 

123 result = await client.complete( 

124 messages, 

125 model=selection.model_id, 

126 temperature=request.temperature, 

127 max_tokens=request.max_tokens, 

128 ) 

129 except (LLMError, OSError, RuntimeError, ValueError) as e: 

130 error = OrchestratorError(f"Provider error: {e!s}") 

131 logger.exception("provider_error", provider=selection.provider) 

132 return Err(error) 

133 

134 if result.is_ok(): 

135 completion = result.unwrap() 

136 else: 

137 llm_error = result.unwrap_err() 

138 return Err(OrchestratorError(str(llm_error))) 

139 tokens_used = completion.usage.total_tokens if completion.usage else 0 

140 logger.info( 

141 "request_completed", 

142 provider=selection.provider, 

143 tokens_used=tokens_used, 

144 ) 

145 return Ok( 

146 ModelResponse( 

147 content=completion.content, 

148 tokens_used=tokens_used, 

149 stop_reason=completion.finish_reason, 

150 ) 

151 ) 

152 

153 def _select_model(self, request: ModelRequest) -> _ModelSelection | None: 

154 """Select the best model for a request using the configured strategy. 

155 

156 Strategy is read from ``request.extra_params["strategy"]`` and defaults 

157 to :attr:`~lexigram.contracts.ai.providers.SelectionStrategy.CAPABILITY_MATCH`. 

158 

159 Priority: 

160 1. Explicit ``model_id`` — find provider that owns this model in its default set. 

161 2. ``ROUND_ROBIN`` — rotate through capability-filtered providers. 

162 3. ``COST_OPTIMAL`` / ``LATENCY_OPTIMAL`` / ``CAPABILITY_MATCH`` / ``PREFERRED`` 

163 — return the first capability-filtered provider (cost/latency data is not 

164 yet available in :class:`~lexigram.ai.llm.registry.core.ProviderRegistry`). 

165 

166 Args: 

167 request: The model request. 

168 

169 Returns: 

170 A :class:`_ModelSelection` with provider name and model ID, or ``None`` 

171 if no suitable provider is found. 

172 """ 

173 # Resolve selection strategy 

174 raw_strategy = request.extra_params.get( 

175 "strategy", SelectionStrategy.CAPABILITY_MATCH 

176 ) 

177 try: 

178 strategy = SelectionStrategy(raw_strategy) 

179 except ValueError: 

180 logger.warning( 

181 "unknown_strategy", strategy=raw_strategy, fallback="capability_match" 

182 ) 

183 strategy = SelectionStrategy.CAPABILITY_MATCH 

184 

185 # Explicit model ID takes precedence: find which provider owns it 

186 if request.model_id: 

187 for name in self.registry.list_providers(): 

188 info = self.registry.get_provider(name) 

189 if request.model_id in info.default_models: 

190 return _ModelSelection(provider=name, model_id=request.model_id) 

191 logger.warning("model_not_found", model_id=request.model_id) 

192 return None 

193 

194 # Derive capability requirements from the request 

195 caps: set[ModelCapability] = request.required_capabilities or set() 

196 needs_tools = ModelCapability.FUNCTION_CALLING in caps 

197 needs_vision = ModelCapability.VISION in caps 

198 needs_streaming = ModelCapability.STREAMING in caps 

199 

200 candidates = self.registry.search_providers( 

201 supports_streaming=True if needs_streaming else None, 

202 supports_tools=True if needs_tools else None, 

203 supports_vision=True if needs_vision else None, 

204 ) 

205 

206 available = [p for p in candidates if p.default_models] 

207 if not available: 

208 return None 

209 

210 # Apply selection strategy 

211 if strategy == SelectionStrategy.ROUND_ROBIN: 

212 chosen = available[self._round_robin_index % len(available)] 

213 self._round_robin_index = (self._round_robin_index + 1) % len(available) 

214 else: 

215 # COST_OPTIMAL, LATENCY_OPTIMAL, CAPABILITY_MATCH, PREFERRED: 

216 # Without real cost/latency data in the registry, fall back to first match. 

217 chosen = available[0] 

218 

219 return _ModelSelection(provider=chosen.name, model_id=chosen.default_models[0])