Coverage for src/chat_limiter/adapters.py: 92%

1""" 

2Provider-specific adapters for converting between our unified types and provider APIs. 

3""" 

4 

5import time 

6import warnings 

7from abc import ABC, abstractmethod 

8from typing import Any 

9 

10from .providers import Provider 

11from .types import ( 

12 ChatCompletionRequest, 

13 ChatCompletionResponse, 

14 Choice, 

15 Message, 

16 MessageRole, 

17 Usage, 

18) 

19 

20 

21class ProviderAdapter(ABC): 

22 """Abstract base class for provider-specific adapters.""" 

23 

24 def is_reasoning_model(self, model_name: str) -> bool: 

25 """Check if the model is a reasoning model (o1, o3, o4 series).""" 

26 # Handle prefixed models (e.g., "openai/o3-mini") 

27 if "/" in model_name: 

28 # Extract the base model name after the "/" 

29 base_model = model_name.split("/", 1)[1] 

30 return base_model.startswith(("o1", "o3", "o4")) 

31 

32 # Handle non-prefixed models 

33 return model_name.startswith(("o1", "o3", "o4")) 

34 

35 @abstractmethod 

36 def format_request(self, request: ChatCompletionRequest) -> dict[str, Any]: 

37 """Convert our request format to provider-specific format.""" 

38 pass 

39 

40 @abstractmethod 

41 def parse_response( 

42 self, 

43 response_data: dict[str, Any], 

44 original_request: ChatCompletionRequest 

45 ) -> ChatCompletionResponse: 

46 """Convert provider response to our unified format.""" 

47 pass 

48 

49 @abstractmethod 

50 def get_endpoint(self) -> str: 

51 """Get the API endpoint for this provider.""" 

52 pass 

53 

54 

55class OpenAIAdapter(ProviderAdapter): 

56 """Adapter for OpenAI API.""" 

57 

58 def format_request(self, request: ChatCompletionRequest) -> dict[str, Any]: 

59 """Convert to OpenAI format.""" 

60 # Convert messages 

61 messages: list[dict[str, Any]] = [] 

62 for msg in request.messages: 

63 messages.append({ 

64 "role": msg.role.value, 

65 "content": msg.content 

66 }) 

67 

68 # Build request 

69 openai_request: dict[str, Any] = { 

70 "model": request.model, 

71 "messages": messages, 

72 } 

73 

74 # Add optional parameters 

75 if request.max_tokens is not None: 

76 # Use max_completion_tokens for reasoning models (o1, o3, o4) 

77 if self.is_reasoning_model(request.model): 

78 openai_request["max_completion_tokens"] = request.max_tokens 

79 else: 

80 openai_request["max_tokens"] = request.max_tokens 

81 

82 # Handle temperature for reasoning models 

83 if self.is_reasoning_model(request.model): 

84 # For reasoning models, default to temperature=1 

85 default_temperature = 1.0 

86 

87 if request.temperature is not None: 

88 # If user provided a different temperature, warn them and use temperature=1 

89 if request.temperature != default_temperature: 

90 warnings.warn( 

91 f"WARNING: Model '{request.model}' is a reasoning model that requires temperature=1. " 

92 f"Your specified temperature={request.temperature} will be overridden to temperature=1.", 

93 UserWarning 

94 ) 

95 print(f"WARNING: Model '{request.model}' is a reasoning model that requires temperature=1. " 

96 f"Your specified temperature={request.temperature} will be overridden to temperature=1.") 

97 

98 # Always use temperature=1 for reasoning models 

99 openai_request["temperature"] = default_temperature 

100 else: 

101 # For non-reasoning models, use the provided temperature 

102 if request.temperature is not None: 

103 openai_request["temperature"] = request.temperature 

104 

105 if request.top_p is not None: 

106 openai_request["top_p"] = request.top_p 

107 if request.stop is not None: 

108 openai_request["stop"] = request.stop 

109 if request.stream: 

110 openai_request["stream"] = request.stream 

111 if request.frequency_penalty is not None: 

112 openai_request["frequency_penalty"] = request.frequency_penalty 

113 if request.presence_penalty is not None: 

114 openai_request["presence_penalty"] = request.presence_penalty 

115 if request.seed is not None: 

116 openai_request["seed"] = request.seed 

117 

118 # Add reasoning parameter for thinking models 

119 if (request.reasoning_effort is not None and 

120 self.is_reasoning_model(request.model)): 

121 openai_request["reasoning"] = {"effort": request.reasoning_effort} 

122 

123 return openai_request 

124 

125 def parse_response( 

126 self, 

127 response_data: dict[str, Any], 

128 original_request: ChatCompletionRequest 

129 ) -> ChatCompletionResponse: 

130 """Parse OpenAI response.""" 

131 # Check for errors first 

132 success = True 

133 error_message = None 

134 

135 if "error" in response_data: 

136 success = False 

137 error_data = response_data["error"] 

138 error_message = error_data.get("message", "Unknown error") 

139 

140 choices = [] 

141 for choice_data in response_data.get("choices", []): 

142 message_data = choice_data.get("message", {}) 

143 message = Message( 

144 role=MessageRole(message_data.get("role", "assistant")), 

145 content=message_data.get("content", "") 

146 ) 

147 choice = Choice( 

148 index=choice_data.get("index", 0), 

149 message=message, 

150 finish_reason=choice_data.get("finish_reason") 

151 ) 

152 choices.append(choice) 

153 

154 # Parse usage 

155 usage = None 

156 if "usage" in response_data: 

157 usage_data = response_data["usage"] 

158 usage = Usage( 

159 prompt_tokens=usage_data.get("prompt_tokens", 0), 

160 completion_tokens=usage_data.get("completion_tokens", 0), 

161 total_tokens=usage_data.get("total_tokens", 0) 

162 ) 

163 

164 return ChatCompletionResponse( 

165 id=response_data.get("id", ""), 

166 model=response_data.get("model", original_request.model), 

167 choices=choices, 

168 usage=usage, 

169 created=response_data.get("created"), 

170 success=success, 

171 error_message=error_message, 

172 provider="openai", 

173 raw_response=response_data 

174 ) 

175 

176 def get_endpoint(self) -> str: 

177 return "/chat/completions" 

178 

179 

180class AnthropicAdapter(ProviderAdapter): 

181 """Adapter for Anthropic API.""" 

182 

183 def format_request(self, request: ChatCompletionRequest) -> dict[str, Any]: 

184 """Convert to Anthropic format.""" 

185 # Anthropic has a different message format 

186 messages: list[dict[str, Any]] = [] 

187 system_message: str | None = None 

188 

189 for msg in request.messages: 

190 if msg.role == MessageRole.SYSTEM: 

191 # Anthropic handles system messages separately 

192 system_message = msg.content 

193 else: 

194 messages.append({ 

195 "role": msg.role.value, 

196 "content": msg.content 

197 }) 

198 

199 # Build request 

200 anthropic_request: dict[str, Any] = { 

201 "model": request.model, 

202 "messages": messages, 

203 "max_tokens": request.max_tokens or 1024, # Required for Anthropic 

204 } 

205 

206 # Add system message if present 

207 if system_message: 

208 anthropic_request["system"] = system_message 

209 

210 # Add optional parameters 

211 if request.temperature is not None: 

212 anthropic_request["temperature"] = request.temperature 

213 if request.top_p is not None: 

214 anthropic_request["top_p"] = request.top_p 

215 if request.stop is not None: 

216 anthropic_request["stop_sequences"] = ( 

217 [request.stop] if isinstance(request.stop, str) else request.stop 

218 ) 

219 if request.stream: 

220 anthropic_request["stream"] = request.stream 

221 if request.top_k is not None: 

222 anthropic_request["top_k"] = request.top_k 

223 if request.seed is not None: 

224 anthropic_request["seed"] = request.seed 

225 

226 return anthropic_request 

227 

228 def parse_response( 

229 self, 

230 response_data: dict[str, Any], 

231 original_request: ChatCompletionRequest 

232 ) -> ChatCompletionResponse: 

233 """Parse Anthropic response.""" 

234 # Check for errors first 

235 success = True 

236 error_message = None 

237 

238 if "error" in response_data: 

239 success = False 

240 error_data = response_data["error"] 

241 error_message = error_data.get("message", "Unknown error") 

242 

243 # Anthropic returns content differently 

244 content_blocks = response_data.get("content", []) 

245 content = "" 

246 if content_blocks: 

247 # Extract text from content blocks 

248 for block in content_blocks: 

249 if block.get("type") == "text": 

250 content += block.get("text", "") 

251 

252 message = Message( 

253 role=MessageRole.ASSISTANT, 

254 content=content 

255 ) 

256 

257 choice = Choice( 

258 index=0, 

259 message=message, 

260 finish_reason=response_data.get("stop_reason") 

261 ) 

262 

263 # Parse usage 

264 usage = None 

265 if "usage" in response_data: 

266 usage_data = response_data["usage"] 

267 usage = Usage( 

268 prompt_tokens=usage_data.get("input_tokens", 0), 

269 completion_tokens=usage_data.get("output_tokens", 0), 

270 total_tokens=usage_data.get("input_tokens", 0) + usage_data.get("output_tokens", 0) 

271 ) 

272 

273 return ChatCompletionResponse( 

274 id=response_data.get("id", ""), 

275 model=response_data.get("model", original_request.model), 

276 choices=[choice], 

277 usage=usage, 

278 created=int(time.time()), # Anthropic doesn't provide created timestamp 

279 success=success, 

280 error_message=error_message, 

281 provider="anthropic", 

282 raw_response=response_data 

283 ) 

284 

285 def get_endpoint(self) -> str: 

286 return "/messages" 

287 

288 

289class OpenRouterAdapter(ProviderAdapter): 

290 """Adapter for OpenRouter API.""" 

291 

292 def format_request(self, request: ChatCompletionRequest) -> dict[str, Any]: 

293 """Convert to OpenRouter format (similar to OpenAI).""" 

294 # OpenRouter uses OpenAI-compatible format 

295 messages: list[dict[str, Any]] = [] 

296 for msg in request.messages: 

297 messages.append({ 

298 "role": msg.role.value, 

299 "content": msg.content 

300 }) 

301 

302 # Build request 

303 openrouter_request: dict[str, Any] = { 

304 "model": request.model, 

305 "messages": messages, 

306 } 

307 

308 # Add optional parameters 

309 if request.max_tokens is not None: 

310 openrouter_request["max_tokens"] = request.max_tokens 

311 if request.temperature is not None: 

312 openrouter_request["temperature"] = request.temperature 

313 if request.top_p is not None: 

314 openrouter_request["top_p"] = request.top_p 

315 if request.stop is not None: 

316 openrouter_request["stop"] = request.stop 

317 if request.stream: 

318 openrouter_request["stream"] = request.stream 

319 if request.frequency_penalty is not None: 

320 openrouter_request["frequency_penalty"] = request.frequency_penalty 

321 if request.presence_penalty is not None: 

322 openrouter_request["presence_penalty"] = request.presence_penalty 

323 if request.top_k is not None: 

324 openrouter_request["top_k"] = request.top_k 

325 if request.seed is not None: 

326 openrouter_request["seed"] = request.seed 

327 

328 # Add reasoning parameter for thinking models 

329 if (request.reasoning_effort is not None and 

330 self.is_reasoning_model(request.model)): 

331 openrouter_request["reasoning"] = {"effort": request.reasoning_effort} 

332 

333 # Add provider routing if specified 

334 if request.providers is not None: 

335 openrouter_request["provider"] = { 

336 "order": request.providers, 

337 "allow_fallbacks": False 

338 } 

339 

340 return openrouter_request 

341 

342 def parse_response( 

343 self, 

344 response_data: dict[str, Any], 

345 original_request: ChatCompletionRequest 

346 ) -> ChatCompletionResponse: 

347 """Parse OpenRouter response (similar to OpenAI).""" 

348 # Check for errors first 

349 success = True 

350 error_message = None 

351 

352 if "error" in response_data: 

353 success = False 

354 error_data = response_data["error"] 

355 error_message = error_data.get("message", "Unknown error") 

356 

357 choices = [] 

358 for choice_data in response_data.get("choices", []): 

359 message_data = choice_data.get("message", {}) 

360 message = Message( 

361 role=MessageRole(message_data.get("role", "assistant")), 

362 content=message_data.get("content", "") 

363 ) 

364 choice = Choice( 

365 index=choice_data.get("index", 0), 

366 message=message, 

367 finish_reason=choice_data.get("finish_reason") 

368 ) 

369 choices.append(choice) 

370 

371 # Parse usage 

372 usage = None 

373 if "usage" in response_data: 

374 usage_data = response_data["usage"] 

375 usage = Usage( 

376 prompt_tokens=usage_data.get("prompt_tokens", 0), 

377 completion_tokens=usage_data.get("completion_tokens", 0), 

378 total_tokens=usage_data.get("total_tokens", 0) 

379 ) 

380 

381 return ChatCompletionResponse( 

382 id=response_data.get("id", ""), 

383 model=response_data.get("model", original_request.model), 

384 choices=choices, 

385 usage=usage, 

386 created=response_data.get("created"), 

387 success=success, 

388 error_message=error_message, 

389 provider="openrouter", 

390 raw_response=response_data 

391 ) 

392 

393 def get_endpoint(self) -> str: 

394 return "/chat/completions" 

395 

396 

397# Provider adapter registry 

398PROVIDER_ADAPTERS = { 

399 Provider.OPENAI: OpenAIAdapter(), 

400 Provider.ANTHROPIC: AnthropicAdapter(), 

401 Provider.OPENROUTER: OpenRouterAdapter(), 

402} 

403 

404 

405def get_adapter(provider: Provider) -> ProviderAdapter: 

406 """Get the appropriate adapter for a provider.""" 

407 return PROVIDER_ADAPTERS[provider]