| 1 | from dataclasses import dataclass, field |
| 2 | from enum import Enum |
| 3 | import json |
| 4 | import logging |
| 5 | import os |
| 6 | import threading |
| 7 | from typing import ( |
| 8 | Any, |
| 9 | Awaitable, |
| 10 | Callable, |
| 11 | List, |
| 12 | Optional, |
| 13 | Iterator, |
| 14 | AsyncIterator, |
| 15 | Tuple, |
| 16 | TypedDict, |
| 17 | ) |
| 18 | |
| 19 | from litellm import embedding |
| 20 | import litellm |
| 21 | import openai |
| 22 | |
| 23 | from helpers import dotenv |
| 24 | from helpers import settings, images |
| 25 | from helpers.dotenv import load_dotenv |
| 26 | from helpers.providers import ModelType as ProviderModelType, get_provider_config |
| 27 | from helpers.rate_limiter import RateLimiter |
| 28 | from helpers.tokens import approximate_tokens |
| 29 | from helpers.extension import extensible # extensible: allows plugins to intercept get_api_key() |
| 30 | from helpers.litellm_transport import LiteLLMTransport, ResponsesTransport |
| 31 | from helpers.llm_result import LLMResult |
| 32 | |
| 33 | from langchain_core.language_models.chat_models import SimpleChatModel |
| 34 | from langchain_core.outputs.chat_generation import ChatGenerationChunk |
| 35 | from langchain_core.callbacks.manager import ( |
| 36 | CallbackManagerForLLMRun, |
| 37 | AsyncCallbackManagerForLLMRun, |
| 38 | ) |
| 39 | from langchain_core.messages import ( |
| 40 | BaseMessage, |
| 41 | AIMessageChunk, |
| 42 | HumanMessage, |
| 43 | SystemMessage, |
| 44 | ) |
| 45 | from langchain.embeddings.base import Embeddings |
| 46 | from sentence_transformers import SentenceTransformer |
| 47 | from pydantic import ConfigDict |
| 48 | |
| 49 | |
| 50 | DEFAULT_LITELLM_GLOBAL_KWARGS: dict[str, Any] = { |
| 51 | "drop_params": True, |
| 52 | } |
| 53 | |
| 54 | # LiteLLM documents drop_params as both a module-level switch and per-call kwarg. |
| 55 | # Other entries in litellm_global_kwargs, such as timeout or additional_drop_params, |
| 56 | # are kept as per-call kwargs instead of becoming arbitrary module attributes. |
| 57 | LITELLM_MODULE_GLOBAL_KEYS = frozenset({"drop_params"}) |
| 58 | |
| 59 | |
| 60 | def _normalize_litellm_kwargs(values: dict[str, Any]) -> dict[str, Any]: |
| 61 | # Normalize .env/UI-style scalar strings into native types for LiteLLM. |
| 62 | result: dict[str, Any] = {} |
| 63 | for k, v in values.items(): |
| 64 | if isinstance(v, str): |
| 65 | stripped = v.strip() |
| 66 | lowered = stripped.lower() |
| 67 | if lowered == "true": |
| 68 | result[k] = True |
| 69 | elif lowered == "false": |
| 70 | result[k] = False |
| 71 | elif lowered in ("none", "null"): |
| 72 | result[k] = None |
| 73 | else: |
| 74 | try: |
| 75 | result[k] = int(stripped) |
| 76 | except ValueError: |
| 77 | try: |
| 78 | result[k] = float(stripped) |
| 79 | except ValueError: |
| 80 | result[k] = v |
| 81 | else: |
| 82 | result[k] = v |
| 83 | return result |
| 84 | |
| 85 | |
| 86 | def get_litellm_global_kwargs() -> dict[str, Any]: |
| 87 | kwargs = _normalize_litellm_kwargs(DEFAULT_LITELLM_GLOBAL_KWARGS) |
| 88 | try: |
| 89 | configured = settings.get_settings().get("litellm_global_kwargs", {}) # type: ignore[union-attr] |
| 90 | except Exception: |
| 91 | configured = {} |
| 92 | if isinstance(configured, dict): |
| 93 | kwargs.update(_normalize_litellm_kwargs(configured)) |
| 94 | return kwargs |
| 95 | |
| 96 | |
| 97 | # keep provider logging quiet in normal operation |
| 98 | def turn_off_logging(): |
| 99 | os.environ["LITELLM_LOG"] = "ERROR" # only errors |
| 100 | litellm.suppress_debug_info = True |
| 101 | # Silence **all** LiteLLM sub-loggers (utils, cost_calculator…) |
| 102 | for name in logging.Logger.manager.loggerDict: |
| 103 | if name.lower().startswith("litellm"): |
| 104 | logging.getLogger(name).setLevel(logging.ERROR) |
| 105 | |
| 106 | |
| 107 | def set_litellm_params(): |
| 108 | global_kwargs = get_litellm_global_kwargs() |
| 109 | for key, value in global_kwargs.items(): |
| 110 | if key not in LITELLM_MODULE_GLOBAL_KEYS: |
| 111 | continue |
| 112 | setattr(litellm, key, value) |
| 113 | return global_kwargs |
| 114 | |
| 115 | |
| 116 | def configure_litellm(): |
| 117 | turn_off_logging() |
| 118 | set_litellm_params() |
| 119 | |
| 120 | |
| 121 | def _merge_litellm_call_kwargs(*overrides: dict[str, Any] | None) -> dict[str, Any]: |
| 122 | kwargs = get_litellm_global_kwargs() |
| 123 | for override in overrides: |
| 124 | if isinstance(override, dict): |
| 125 | kwargs.update(override) |
| 126 | return kwargs |
| 127 | |
| 128 | |
| 129 | # init |
| 130 | load_dotenv() |
| 131 | configure_litellm() |
| 132 | |
| 133 | |
| 134 | class ModelType(Enum): |
| 135 | CHAT = "Chat" |
| 136 | EMBEDDING = "Embedding" |
| 137 | |
| 138 | |
| 139 | @dataclass |
| 140 | class ModelConfig: |
| 141 | type: ModelType |
| 142 | provider: str |
| 143 | name: str |
| 144 | api_key: str = "" |
| 145 | api_base: str = "" |
| 146 | ctx_length: int = 0 |
| 147 | limit_requests: int = 0 |
| 148 | limit_input: int = 0 |
| 149 | limit_output: int = 0 |
| 150 | vision: bool = False |
| 151 | kwargs: dict = field(default_factory=dict) |
| 152 | |
| 153 | def build_kwargs(self): |
| 154 | kwargs = self.kwargs.copy() or {} |
| 155 | if self.api_key and "api_key" not in kwargs: |
| 156 | kwargs["api_key"] = self.api_key |
| 157 | if self.api_base and "api_base" not in kwargs: |
| 158 | kwargs["api_base"] = self.api_base |
| 159 | return kwargs |
| 160 | |
| 161 | |
| 162 | class ChatChunk(TypedDict): |
| 163 | """Simplified response chunk for chat models.""" |
| 164 | response_delta: str |
| 165 | reasoning_delta: str |
| 166 | |
| 167 | |
| 168 | class ChatGenerationResult: |
| 169 | """Chat generation result object""" |
| 170 | def __init__(self, chunk: ChatChunk|None = None): |
| 171 | self.reasoning = "" |
| 172 | self.response = "" |
| 173 | self.thinking = False |
| 174 | self.thinking_tag = "" |
| 175 | self.unprocessed = "" |
| 176 | self.native_reasoning = False |
| 177 | self.thinking_pairs = [("<think>", "</think>"), ("<reasoning>", "</reasoning>")] |
| 178 | if chunk: |
| 179 | self.add_chunk(chunk) |
| 180 | |
| 181 | def add_chunk(self, chunk: ChatChunk) -> ChatChunk: |
| 182 | if chunk["reasoning_delta"]: |
| 183 | self.native_reasoning = True |
| 184 | |
| 185 | # if native reasoning detection works, there's no need to worry about thinking tags |
| 186 | if self.native_reasoning: |
| 187 | processed_chunk = ChatChunk(response_delta=chunk["response_delta"], reasoning_delta=chunk["reasoning_delta"]) |
| 188 | else: |
| 189 | # if the model outputs thinking tags, we ned to parse them manually as reasoning |
| 190 | processed_chunk = self._process_thinking_chunk(chunk) |
| 191 | |
| 192 | self.reasoning += processed_chunk.get("reasoning_delta", "") |
| 193 | self.response += processed_chunk.get("response_delta", "") |
| 194 | |
| 195 | return processed_chunk |
| 196 | |
| 197 | def _process_thinking_chunk(self, chunk: ChatChunk) -> ChatChunk: |
| 198 | response_delta = self.unprocessed + chunk["response_delta"] |
| 199 | self.unprocessed = "" |
| 200 | return self._process_thinking_tags(response_delta, chunk["reasoning_delta"]) |
| 201 | |
| 202 | def _process_thinking_tags(self, response: str, reasoning: str) -> ChatChunk: |
| 203 | if self.thinking: |
| 204 | close_pos = response.find(self.thinking_tag) |
| 205 | if close_pos != -1: |
| 206 | reasoning += response[:close_pos] |
| 207 | response = response[close_pos + len(self.thinking_tag):] |
| 208 | self.thinking = False |
| 209 | self.thinking_tag = "" |
| 210 | else: |
| 211 | if self._is_partial_closing_tag(response): |
| 212 | self.unprocessed = response |
| 213 | response = "" |
| 214 | else: |
| 215 | reasoning += response |
| 216 | response = "" |
| 217 | else: |
| 218 | for opening_tag, closing_tag in self.thinking_pairs: |
| 219 | if response.startswith(opening_tag): |
| 220 | response = response[len(opening_tag):] |
| 221 | self.thinking = True |
| 222 | self.thinking_tag = closing_tag |
| 223 | |
| 224 | close_pos = response.find(closing_tag) |
| 225 | if close_pos != -1: |
| 226 | reasoning += response[:close_pos] |
| 227 | response = response[close_pos + len(closing_tag):] |
| 228 | self.thinking = False |
| 229 | self.thinking_tag = "" |
| 230 | else: |
| 231 | if self._is_partial_closing_tag(response): |
| 232 | self.unprocessed = response |
| 233 | response = "" |
| 234 | else: |
| 235 | reasoning += response |
| 236 | response = "" |
| 237 | break |
| 238 | elif len(response) < len(opening_tag) and self._is_partial_opening_tag(response, opening_tag): |
| 239 | self.unprocessed = response |
| 240 | response = "" |
| 241 | break |
| 242 | |
| 243 | return ChatChunk(response_delta=response, reasoning_delta=reasoning) |
| 244 | |
| 245 | def _is_partial_opening_tag(self, text: str, opening_tag: str) -> bool: |
| 246 | for i in range(1, len(opening_tag)): |
| 247 | if text == opening_tag[:i]: |
| 248 | return True |
| 249 | return False |
| 250 | |
| 251 | def _is_partial_closing_tag(self, text: str) -> bool: |
| 252 | if not self.thinking_tag or not text: |
| 253 | return False |
| 254 | max_check = min(len(text), len(self.thinking_tag) - 1) |
| 255 | for i in range(1, max_check + 1): |
| 256 | if text.endswith(self.thinking_tag[:i]): |
| 257 | return True |
| 258 | return False |
| 259 | |
| 260 | def output(self) -> ChatChunk: |
| 261 | response = self.response |
| 262 | reasoning = self.reasoning |
| 263 | if self.unprocessed: |
| 264 | if reasoning and not response: |
| 265 | reasoning += self.unprocessed |
| 266 | else: |
| 267 | response += self.unprocessed |
| 268 | return ChatChunk(response_delta=response, reasoning_delta=reasoning) |
| 269 | |
| 270 | |
| 271 | rate_limiters: dict[str, RateLimiter] = {} |
| 272 | api_keys_round_robin: dict[str, int] = {} |
| 273 | |
| 274 | |
| 275 | @extensible |
| 276 | def get_api_key(service: str) -> str: |
| 277 | # get api key for the service |
| 278 | key = ( |
| 279 | dotenv.get_dotenv_value(f"API_KEY_{service.upper()}") |
| 280 | or dotenv.get_dotenv_value(f"{service.upper()}_API_KEY") |
| 281 | or dotenv.get_dotenv_value(f"{service.upper()}_API_TOKEN") |
| 282 | or "None" |
| 283 | ) |
| 284 | # if the key contains a comma, use round-robin |
| 285 | if "," in key: |
| 286 | api_keys = [k.strip() for k in key.split(",") if k.strip()] |
| 287 | api_keys_round_robin[service] = api_keys_round_robin.get(service, -1) + 1 |
| 288 | key = api_keys[api_keys_round_robin[service] % len(api_keys)] |
| 289 | return key |
| 290 | |
| 291 | |
| 292 | def get_rate_limiter( |
| 293 | provider: str, name: str, requests: int, input: int, output: int |
| 294 | ) -> RateLimiter: |
| 295 | key = f"{provider}\\{name}" |
| 296 | rate_limiters[key] = limiter = rate_limiters.get(key, RateLimiter(seconds=60)) |
| 297 | limiter.limits["requests"] = requests or 0 |
| 298 | limiter.limits["input"] = input or 0 |
| 299 | limiter.limits["output"] = output or 0 |
| 300 | return limiter |
| 301 | |
| 302 | |
| 303 | def _is_transient_litellm_error(exc: Exception) -> bool: |
| 304 | """Uses status_code when available, else falls back to exception types""" |
| 305 | # Prefer explicit status codes if present |
| 306 | status_code = getattr(exc, "status_code", None) |
| 307 | if isinstance(status_code, int): |
| 308 | if status_code in (408, 429, 500, 502, 503, 504): |
| 309 | return True |
| 310 | # Treat other 5xx as retriable |
| 311 | if status_code >= 500: |
| 312 | return True |
| 313 | return False |
| 314 | |
| 315 | # Fallback to exception classes mapped by LiteLLM/OpenAI |
| 316 | transient_types = ( |
| 317 | getattr(openai, "APITimeoutError", Exception), |
| 318 | getattr(openai, "APIConnectionError", Exception), |
| 319 | getattr(openai, "RateLimitError", Exception), |
| 320 | getattr(openai, "APIError", Exception), |
| 321 | getattr(openai, "InternalServerError", Exception), |
| 322 | # Some providers map overloads to ServiceUnavailable-like errors |
| 323 | getattr(openai, "APIStatusError", Exception), |
| 324 | ) |
| 325 | return isinstance(exc, transient_types) |
| 326 | |
| 327 | |
| 328 | async def apply_rate_limiter( |
| 329 | model_config: ModelConfig | None, |
| 330 | input_text: str, |
| 331 | rate_limiter_callback: ( |
| 332 | Callable[[str, str, int, int], Awaitable[bool]] | None |
| 333 | ) = None, |
| 334 | ): |
| 335 | if not model_config: |
| 336 | return |
| 337 | limiter = get_rate_limiter( |
| 338 | model_config.provider, |
| 339 | model_config.name, |
| 340 | model_config.limit_requests, |
| 341 | model_config.limit_input, |
| 342 | model_config.limit_output, |
| 343 | ) |
| 344 | limiter.add(input=approximate_tokens(input_text)) |
| 345 | limiter.add(requests=1) |
| 346 | await limiter.wait(rate_limiter_callback) |
| 347 | return limiter |
| 348 | |
| 349 | |
| 350 | def apply_rate_limiter_sync( |
| 351 | model_config: ModelConfig | None, |
| 352 | input_text: str, |
| 353 | rate_limiter_callback: ( |
| 354 | Callable[[str, str, int, int], Awaitable[bool]] | None |
| 355 | ) = None, |
| 356 | ): |
| 357 | if not model_config: |
| 358 | return |
| 359 | import asyncio, nest_asyncio |
| 360 | |
| 361 | nest_asyncio.apply() |
| 362 | return asyncio.run( |
| 363 | apply_rate_limiter(model_config, input_text, rate_limiter_callback) |
| 364 | ) |
| 365 | |
| 366 | |
| 367 | class LiteLLMChatWrapper(SimpleChatModel): |
| 368 | model_name: str |
| 369 | provider: str |
| 370 | kwargs: dict = {} |
| 371 | |
| 372 | model_config = ConfigDict( |
| 373 | arbitrary_types_allowed=True, |
| 374 | extra="allow", |
| 375 | validate_assignment=False, |
| 376 | ) |
| 377 | |
| 378 | def __init__( |
| 379 | self, |
| 380 | model: str, |
| 381 | provider: str, |
| 382 | model_config: Optional[ModelConfig] = None, |
| 383 | **kwargs: Any, |
| 384 | ): |
| 385 | model_value = f"{provider}/{model}" |
| 386 | super().__init__(model_name=model_value, provider=provider, kwargs=kwargs) # type: ignore |
| 387 | # Set A0 model config as instance attribute after parent init |
| 388 | self.a0_model_conf = model_config |
| 389 | |
| 390 | @property |
| 391 | def _llm_type(self) -> str: |
| 392 | return "litellm-chat" |
| 393 | |
| 394 | def _convert_messages(self, messages: List[BaseMessage], explicit_caching: bool = False) -> List[dict]: |
| 395 | result = [] |
| 396 | # Map LangChain message types to LiteLLM roles |
| 397 | role_mapping = { |
| 398 | "human": "user", |
| 399 | "ai": "assistant", |
| 400 | "system": "system", |
| 401 | "tool": "tool", |
| 402 | } |
| 403 | for m in messages: |
| 404 | role = role_mapping.get(m.type, m.type) |
| 405 | message_dict = {"role": role, "content": images.prepare_content(m.content)} |
| 406 | |
| 407 | # Handle tool calls for AI messages |
| 408 | tool_calls = getattr(m, "tool_calls", None) |
| 409 | if tool_calls: |
| 410 | # Convert LangChain tool calls to LiteLLM format |
| 411 | new_tool_calls = [] |
| 412 | for tool_call in tool_calls: |
| 413 | # Ensure arguments is a JSON string |
| 414 | args = tool_call["args"] |
| 415 | if isinstance(args, dict): |
| 416 | import json |
| 417 | |
| 418 | args_str = json.dumps(args) |
| 419 | else: |
| 420 | args_str = str(args) |
| 421 | |
| 422 | new_tool_calls.append( |
| 423 | { |
| 424 | "id": tool_call.get("id", ""), |
| 425 | "type": "function", |
| 426 | "function": { |
| 427 | "name": tool_call["name"], |
| 428 | "arguments": args_str, |
| 429 | }, |
| 430 | } |
| 431 | ) |
| 432 | message_dict["tool_calls"] = new_tool_calls |
| 433 | |
| 434 | # Handle tool call ID for ToolMessage |
| 435 | tool_call_id = getattr(m, "tool_call_id", None) |
| 436 | if tool_call_id: |
| 437 | message_dict["tool_call_id"] = tool_call_id |
| 438 | |
| 439 | # fix messages with empty content, this breaks some LLMs |
| 440 | content = message_dict.get("content") |
| 441 | has_content = bool(content) if not isinstance(content, list) else len(content) > 0 |
| 442 | if not has_content: |
| 443 | message_dict["content"] = "empty" |
| 444 | |
| 445 | result.append(message_dict) |
| 446 | |
| 447 | return result |
| 448 | |
| 449 | def _call( |
| 450 | self, |
| 451 | messages: List[BaseMessage], |
| 452 | stop: Optional[List[str]] = None, |
| 453 | run_manager: Optional[CallbackManagerForLLMRun] = None, |
| 454 | **kwargs: Any, |
| 455 | ) -> str: |
| 456 | configure_litellm() |
| 457 | msgs = self._convert_messages(messages) |
| 458 | |
| 459 | # Apply rate limiting if configured |
| 460 | apply_rate_limiter_sync(self.a0_model_conf, str(msgs)) |
| 461 | |
| 462 | call_kwargs = _merge_litellm_call_kwargs(self.kwargs, kwargs) |
| 463 | transport = LiteLLMTransport( |
| 464 | model=self.model_name, |
| 465 | messages=msgs, |
| 466 | kwargs=call_kwargs, |
| 467 | stop=stop, |
| 468 | ) |
| 469 | parsed = transport.complete() |
| 470 | output = ChatGenerationResult(parsed).output() |
| 471 | return output["response_delta"] |
| 472 | |
| 473 | def _stream( |
| 474 | self, |
| 475 | messages: List[BaseMessage], |
| 476 | stop: Optional[List[str]] = None, |
| 477 | run_manager: Optional[CallbackManagerForLLMRun] = None, |
| 478 | **kwargs: Any, |
| 479 | ) -> Iterator[ChatGenerationChunk]: |
| 480 | configure_litellm() |
| 481 | msgs = self._convert_messages(messages) |
| 482 | |
| 483 | # Apply rate limiting if configured |
| 484 | apply_rate_limiter_sync(self.a0_model_conf, str(msgs)) |
| 485 | |
| 486 | result = ChatGenerationResult() |
| 487 | call_kwargs = _merge_litellm_call_kwargs(self.kwargs, kwargs) |
| 488 | transport = LiteLLMTransport( |
| 489 | model=self.model_name, |
| 490 | messages=msgs, |
| 491 | kwargs=call_kwargs, |
| 492 | stop=stop, |
| 493 | ) |
| 494 | for parsed in transport.stream(): |
| 495 | output = result.add_chunk(parsed) |
| 496 | if output["response_delta"]: |
| 497 | yield ChatGenerationChunk( |
| 498 | message=AIMessageChunk(content=output["response_delta"]) |
| 499 | ) |
| 500 | |
| 501 | async def _astream( |
| 502 | self, |
| 503 | messages: List[BaseMessage], |
| 504 | stop: Optional[List[str]] = None, |
| 505 | run_manager: Optional[AsyncCallbackManagerForLLMRun] = None, |
| 506 | **kwargs: Any, |
| 507 | ) -> AsyncIterator[ChatGenerationChunk]: |
| 508 | configure_litellm() |
| 509 | msgs = self._convert_messages(messages) |
| 510 | |
| 511 | # Apply rate limiting if configured |
| 512 | await apply_rate_limiter(self.a0_model_conf, str(msgs)) |
| 513 | |
| 514 | result = ChatGenerationResult() |
| 515 | call_kwargs = _merge_litellm_call_kwargs(self.kwargs, kwargs) |
| 516 | transport = LiteLLMTransport( |
| 517 | model=self.model_name, |
| 518 | messages=msgs, |
| 519 | kwargs=call_kwargs, |
| 520 | stop=stop, |
| 521 | ) |
| 522 | async for parsed in transport.astream(): |
| 523 | output = result.add_chunk(parsed) |
| 524 | if output["response_delta"]: |
| 525 | yield ChatGenerationChunk( |
| 526 | message=AIMessageChunk(content=output["response_delta"]) |
| 527 | ) |
| 528 | |
| 529 | @extensible |
| 530 | async def unified_call( |
| 531 | self, |
| 532 | system_message="", |
| 533 | user_message="", |
| 534 | messages: List[BaseMessage] | None = None, |
| 535 | response_callback: Callable[[str, str], Awaitable[str | None]] | None = None, |
| 536 | reasoning_callback: Callable[[str, str], Awaitable[None]] | None = None, |
| 537 | tokens_callback: Callable[[str, int], Awaitable[None]] | None = None, |
| 538 | rate_limiter_callback: ( |
| 539 | Callable[[str, str, int, int], Awaitable[bool]] | None |
| 540 | ) = None, |
| 541 | explicit_caching: bool = False, |
| 542 | **kwargs: Any, |
| 543 | ) -> Tuple[str, str]: |
| 544 | |
| 545 | configure_litellm() |
| 546 | |
| 547 | if not messages: |
| 548 | messages = [] |
| 549 | # construct messages |
| 550 | if system_message: |
| 551 | messages.insert(0, SystemMessage(content=system_message)) |
| 552 | if user_message: |
| 553 | messages.append(HumanMessage(content=user_message)) |
| 554 | |
| 555 | # convert to litellm format |
| 556 | msgs_conv = self._convert_messages(messages, explicit_caching=explicit_caching) |
| 557 | |
| 558 | # Apply rate limiting if configured |
| 559 | limiter = await apply_rate_limiter( |
| 560 | self.a0_model_conf, str(msgs_conv), rate_limiter_callback |
| 561 | ) |
| 562 | |
| 563 | # Prepare call kwargs and retry config (strip A0-only params before calling LiteLLM) |
| 564 | call_kwargs: dict[str, Any] = _merge_litellm_call_kwargs( |
| 565 | self.kwargs, kwargs |
| 566 | ) |
| 567 | if explicit_caching: |
| 568 | call_kwargs["a0_explicit_prompt_caching"] = True |
| 569 | max_retries: int = int(call_kwargs.pop("a0_retry_attempts", 2)) |
| 570 | retry_delay_s: float = float(call_kwargs.pop("a0_retry_delay_seconds", 1.5)) |
| 571 | stream = reasoning_callback is not None or response_callback is not None or tokens_callback is not None |
| 572 | transport = LiteLLMTransport( |
| 573 | model=self.model_name, |
| 574 | messages=msgs_conv, |
| 575 | kwargs=call_kwargs, |
| 576 | ) |
| 577 | |
| 578 | # results |
| 579 | result = ChatGenerationResult() |
| 580 | |
| 581 | attempt = 0 |
| 582 | while True: |
| 583 | got_any_chunk = False |
| 584 | try: |
| 585 | if stream: |
| 586 | stop_response: str | None = None |
| 587 | async for parsed in transport.astream(): |
| 588 | got_any_chunk = True |
| 589 | output = result.add_chunk(parsed) |
| 590 | |
| 591 | # collect reasoning delta and call callbacks |
| 592 | if output["reasoning_delta"]: |
| 593 | if reasoning_callback: |
| 594 | await reasoning_callback(output["reasoning_delta"], result.reasoning) |
| 595 | if tokens_callback: |
| 596 | await tokens_callback( |
| 597 | output["reasoning_delta"], |
| 598 | approximate_tokens(output["reasoning_delta"]), |
| 599 | ) |
| 600 | # Add output tokens to rate limiter if configured |
| 601 | if limiter: |
| 602 | limiter.add(output=approximate_tokens(output["reasoning_delta"])) |
| 603 | # collect response delta and call callbacks |
| 604 | if output["response_delta"]: |
| 605 | if response_callback: |
| 606 | stop_response = await response_callback( |
| 607 | output["response_delta"], result.response |
| 608 | ) |
| 609 | if tokens_callback: |
| 610 | await tokens_callback( |
| 611 | output["response_delta"], |
| 612 | approximate_tokens(output["response_delta"]), |
| 613 | ) |
| 614 | # Add output tokens to rate limiter if configured |
| 615 | if limiter: |
| 616 | limiter.add(output=approximate_tokens(output["response_delta"])) |
| 617 | if stop_response is not None: |
| 618 | result.response = stop_response |
| 619 | break |
| 620 | |
| 621 | # non-stream response |
| 622 | else: |
| 623 | parsed = await transport.acomplete() |
| 624 | output = result.add_chunk(parsed) |
| 625 | if limiter: |
| 626 | if output["response_delta"]: |
| 627 | limiter.add(output=approximate_tokens(output["response_delta"])) |
| 628 | if output["reasoning_delta"]: |
| 629 | limiter.add(output=approximate_tokens(output["reasoning_delta"])) |
| 630 | |
| 631 | # Successful completion of stream |
| 632 | return result.response, result.reasoning |
| 633 | |
| 634 | except Exception as e: |
| 635 | import asyncio |
| 636 | |
| 637 | # Retry only if no chunks received and error is transient |
| 638 | if got_any_chunk or not _is_transient_litellm_error(e) or attempt >= max_retries: |
| 639 | raise |
| 640 | attempt += 1 |
| 641 | await asyncio.sleep(retry_delay_s) |
| 642 | |
| 643 | @extensible |
| 644 | async def unified_turn( |
| 645 | self, |
| 646 | system_message="", |
| 647 | user_message="", |
| 648 | messages: List[BaseMessage] | None = None, |
| 649 | response_callback: Callable[[str, str], Awaitable[str | None]] | None = None, |
| 650 | reasoning_callback: Callable[[str, str], Awaitable[None]] | None = None, |
| 651 | tokens_callback: Callable[[str, int], Awaitable[None]] | None = None, |
| 652 | rate_limiter_callback: ( |
| 653 | Callable[[str, str, int, int], Awaitable[bool]] | None |
| 654 | ) = None, |
| 655 | explicit_caching: bool = False, |
| 656 | **kwargs: Any, |
| 657 | ) -> LLMResult: |
| 658 | """Canonical internal LLM turn with Responses metadata. |
| 659 | |
| 660 | Public plugin-facing callers should keep using ``unified_call``. Core |
| 661 | orchestration uses this method when it needs response ids, native output |
| 662 | items, and state/capability metadata. |
| 663 | """ |
| 664 | |
| 665 | configure_litellm() |
| 666 | |
| 667 | if not messages: |
| 668 | messages = [] |
| 669 | if system_message: |
| 670 | messages.insert(0, SystemMessage(content=system_message)) |
| 671 | if user_message: |
| 672 | messages.append(HumanMessage(content=user_message)) |
| 673 | |
| 674 | msgs_conv = self._convert_messages(messages, explicit_caching=explicit_caching) |
| 675 | |
| 676 | limiter = await apply_rate_limiter( |
| 677 | self.a0_model_conf, str(msgs_conv), rate_limiter_callback |
| 678 | ) |
| 679 | |
| 680 | call_kwargs: dict[str, Any] = _merge_litellm_call_kwargs( |
| 681 | self.kwargs, kwargs |
| 682 | ) |
| 683 | if explicit_caching: |
| 684 | call_kwargs["a0_explicit_prompt_caching"] = True |
| 685 | max_retries: int = int(call_kwargs.pop("a0_retry_attempts", 2)) |
| 686 | retry_delay_s: float = float(call_kwargs.pop("a0_retry_delay_seconds", 1.5)) |
| 687 | stream = ( |
| 688 | reasoning_callback is not None |
| 689 | or response_callback is not None |
| 690 | or tokens_callback is not None |
| 691 | ) |
| 692 | transport = LiteLLMTransport( |
| 693 | model=self.model_name, |
| 694 | messages=msgs_conv, |
| 695 | kwargs=call_kwargs, |
| 696 | ) |
| 697 | |
| 698 | result = ChatGenerationResult() |
| 699 | |
| 700 | attempt = 0 |
| 701 | while True: |
| 702 | got_any_chunk = False |
| 703 | try: |
| 704 | if stream: |
| 705 | stop_response: str | None = None |
| 706 | async for parsed in transport.astream(): |
| 707 | got_any_chunk = True |
| 708 | output = result.add_chunk(parsed) |
| 709 | |
| 710 | if output["reasoning_delta"]: |
| 711 | if reasoning_callback: |
| 712 | await reasoning_callback( |
| 713 | output["reasoning_delta"], result.reasoning |
| 714 | ) |
| 715 | if tokens_callback: |
| 716 | await tokens_callback( |
| 717 | output["reasoning_delta"], |
| 718 | approximate_tokens(output["reasoning_delta"]), |
| 719 | ) |
| 720 | if limiter: |
| 721 | limiter.add( |
| 722 | output=approximate_tokens( |
| 723 | output["reasoning_delta"] |
| 724 | ) |
| 725 | ) |
| 726 | |
| 727 | if output["response_delta"]: |
| 728 | if response_callback: |
| 729 | stop_response = await response_callback( |
| 730 | output["response_delta"], result.response |
| 731 | ) |
| 732 | if tokens_callback: |
| 733 | await tokens_callback( |
| 734 | output["response_delta"], |
| 735 | approximate_tokens(output["response_delta"]), |
| 736 | ) |
| 737 | if limiter: |
| 738 | limiter.add( |
| 739 | output=approximate_tokens( |
| 740 | output["response_delta"] |
| 741 | ) |
| 742 | ) |
| 743 | if ( |
| 744 | stop_response is not None |
| 745 | and not transport.policy.using_responses |
| 746 | ): |
| 747 | result.response = stop_response |
| 748 | break |
| 749 | else: |
| 750 | parsed = await transport.acomplete() |
| 751 | output = result.add_chunk(parsed) |
| 752 | if limiter: |
| 753 | if output["response_delta"]: |
| 754 | limiter.add( |
| 755 | output=approximate_tokens(output["response_delta"]) |
| 756 | ) |
| 757 | if output["reasoning_delta"]: |
| 758 | limiter.add( |
| 759 | output=approximate_tokens(output["reasoning_delta"]) |
| 760 | ) |
| 761 | |
| 762 | llm_result = transport.last_result or LLMResult.from_chat( |
| 763 | response=result.output()["response_delta"], |
| 764 | reasoning=result.output()["reasoning_delta"], |
| 765 | input_items=ResponsesTransport.input_from_messages(msgs_conv), |
| 766 | provider_model_key=self.model_name, |
| 767 | capability=transport._capability_metadata(), |
| 768 | ) |
| 769 | if result.output()["response_delta"] and not llm_result.function_calls: |
| 770 | llm_result.response = result.output()["response_delta"] |
| 771 | if result.output()["reasoning_delta"]: |
| 772 | llm_result.reasoning = result.output()["reasoning_delta"] |
| 773 | return llm_result |
| 774 | |
| 775 | except Exception as e: |
| 776 | import asyncio |
| 777 | |
| 778 | if ( |
| 779 | got_any_chunk |
| 780 | or not _is_transient_litellm_error(e) |
| 781 | or attempt >= max_retries |
| 782 | ): |
| 783 | raise |
| 784 | attempt += 1 |
| 785 | await asyncio.sleep(retry_delay_s) |
| 786 | |
| 787 | |
| 788 | class LiteLLMEmbeddingWrapper(Embeddings): |
| 789 | model_name: str |
| 790 | kwargs: dict = {} |
| 791 | a0_model_conf: Optional[ModelConfig] = None |
| 792 | |
| 793 | def __init__( |
| 794 | self, |
| 795 | model: str, |
| 796 | provider: str, |
| 797 | model_config: Optional[ModelConfig] = None, |
| 798 | **kwargs: Any, |
| 799 | ): |
| 800 | self.model_name = f"{provider}/{model}" |
| 801 | self.kwargs = kwargs |
| 802 | self.a0_model_conf = model_config |
| 803 | |
| 804 | def embed_documents(self, texts: List[str]) -> List[List[float]]: |
| 805 | configure_litellm() |
| 806 | # Apply rate limiting if configured |
| 807 | apply_rate_limiter_sync(self.a0_model_conf, " ".join(texts)) |
| 808 | |
| 809 | resp = embedding( |
| 810 | model=self.model_name, |
| 811 | input=texts, |
| 812 | **_merge_litellm_call_kwargs(self.kwargs), |
| 813 | ) |
| 814 | return [ |
| 815 | item.get("embedding") if isinstance(item, dict) else item.embedding # type: ignore |
| 816 | for item in resp.data # type: ignore |
| 817 | ] |
| 818 | |
| 819 | def embed_query(self, text: str) -> List[float]: |
| 820 | configure_litellm() |
| 821 | # Apply rate limiting if configured |
| 822 | apply_rate_limiter_sync(self.a0_model_conf, text) |
| 823 | |
| 824 | resp = embedding( |
| 825 | model=self.model_name, |
| 826 | input=[text], |
| 827 | **_merge_litellm_call_kwargs(self.kwargs), |
| 828 | ) |
| 829 | item = resp.data[0] # type: ignore |
| 830 | return item.get("embedding") if isinstance(item, dict) else item.embedding # type: ignore |
| 831 | |
| 832 | |
| 833 | _LOCAL_EMBEDDING_MODELS: dict[tuple[str, str], SentenceTransformer] = {} |
| 834 | _LOCAL_EMBEDDING_MODELS_LOCK = threading.Lock() |
| 835 | |
| 836 | |
| 837 | def _get_local_embedding_model( |
| 838 | model: str, kwargs: dict[str, Any] |
| 839 | ) -> SentenceTransformer: |
| 840 | key = (model, json.dumps(kwargs, sort_keys=True, default=repr)) |
| 841 | with _LOCAL_EMBEDDING_MODELS_LOCK: |
| 842 | cached = _LOCAL_EMBEDDING_MODELS.get(key) |
| 843 | if cached is None: |
| 844 | cached = SentenceTransformer(model, **kwargs) |
| 845 | _LOCAL_EMBEDDING_MODELS.clear() |
| 846 | _LOCAL_EMBEDDING_MODELS[key] = cached |
| 847 | return cached |
| 848 | |
| 849 | |
| 850 | class LocalSentenceTransformerWrapper(Embeddings): |
| 851 | """Local wrapper for sentence-transformers models to avoid HuggingFace API calls""" |
| 852 | |
| 853 | def __init__( |
| 854 | self, |
| 855 | provider: str, |
| 856 | model: str, |
| 857 | model_config: Optional[ModelConfig] = None, |
| 858 | **kwargs: Any, |
| 859 | ): |
| 860 | # Clean common user-input mistakes |
| 861 | model = model.strip().strip('"').strip("'") |
| 862 | |
| 863 | # Remove the "sentence-transformers/" prefix if present |
| 864 | if model.startswith("sentence-transformers/"): |
| 865 | model = model[len("sentence-transformers/") :] |
| 866 | |
| 867 | # Filter kwargs for SentenceTransformer only (no LiteLLM params like 'stream_timeout') |
| 868 | st_allowed_keys = { |
| 869 | "device", |
| 870 | "cache_folder", |
| 871 | "use_auth_token", |
| 872 | "revision", |
| 873 | "trust_remote_code", |
| 874 | "model_kwargs", |
| 875 | } |
| 876 | st_kwargs = {k: v for k, v in (kwargs or {}).items() if k in st_allowed_keys} |
| 877 | |
| 878 | self.model = _get_local_embedding_model(model, st_kwargs) |
| 879 | self.model_name = model |
| 880 | self.a0_model_conf = model_config |
| 881 | |
| 882 | def embed_documents(self, texts: List[str]) -> List[List[float]]: |
| 883 | # Apply rate limiting if configured |
| 884 | apply_rate_limiter_sync(self.a0_model_conf, " ".join(texts)) |
| 885 | |
| 886 | embeddings = self.model.encode(texts, convert_to_tensor=False) # type: ignore |
| 887 | return embeddings.tolist() if hasattr(embeddings, "tolist") else embeddings # type: ignore |
| 888 | |
| 889 | def embed_query(self, text: str) -> List[float]: |
| 890 | # Apply rate limiting if configured |
| 891 | apply_rate_limiter_sync(self.a0_model_conf, text) |
| 892 | |
| 893 | embedding = self.model.encode([text], convert_to_tensor=False) # type: ignore |
| 894 | result = ( |
| 895 | embedding[0].tolist() if hasattr(embedding[0], "tolist") else embedding[0] |
| 896 | ) |
| 897 | return result # type: ignore |
| 898 | |
| 899 | |
| 900 | def _get_litellm_chat( |
| 901 | cls: type = LiteLLMChatWrapper, |
| 902 | model_name: str = "", |
| 903 | provider_name: str = "", |
| 904 | model_config: Optional[ModelConfig] = None, |
| 905 | **kwargs: Any, |
| 906 | ): |
| 907 | # use api key from kwargs or env |
| 908 | api_key = kwargs.pop("api_key", None) or get_api_key(provider_name) |
| 909 | |
| 910 | # Only pass API key if key is not a placeholder |
| 911 | if api_key and api_key not in ("None", "NA"): |
| 912 | kwargs["api_key"] = api_key |
| 913 | |
| 914 | provider_name, model_name, kwargs = _adjust_call_args( |
| 915 | provider_name, model_name, kwargs |
| 916 | ) |
| 917 | return cls( |
| 918 | provider=provider_name, model=model_name, model_config=model_config, **kwargs |
| 919 | ) |
| 920 | |
| 921 | |
| 922 | def _get_litellm_embedding( |
| 923 | model_name: str, |
| 924 | provider_name: str, |
| 925 | model_config: Optional[ModelConfig] = None, |
| 926 | **kwargs: Any, |
| 927 | ): |
| 928 | # Check if this is a local sentence-transformers model |
| 929 | if provider_name == "huggingface" and model_name.startswith( |
| 930 | "sentence-transformers/" |
| 931 | ): |
| 932 | # Use local sentence-transformers instead of LiteLLM for local models |
| 933 | provider_name, model_name, kwargs = _adjust_call_args( |
| 934 | provider_name, model_name, kwargs |
| 935 | ) |
| 936 | return LocalSentenceTransformerWrapper( |
| 937 | provider=provider_name, |
| 938 | model=model_name, |
| 939 | model_config=model_config, |
| 940 | **kwargs, |
| 941 | ) |
| 942 | |
| 943 | # use api key from kwargs or env |
| 944 | api_key = kwargs.pop("api_key", None) or get_api_key(provider_name) |
| 945 | |
| 946 | # Only pass API key if key is not a placeholder |
| 947 | if api_key and api_key not in ("None", "NA"): |
| 948 | kwargs["api_key"] = api_key |
| 949 | |
| 950 | provider_name, model_name, kwargs = _adjust_call_args( |
| 951 | provider_name, model_name, kwargs |
| 952 | ) |
| 953 | return LiteLLMEmbeddingWrapper( |
| 954 | model=model_name, provider=provider_name, model_config=model_config, **kwargs |
| 955 | ) |
| 956 | |
| 957 | |
| 958 | |
| 959 | def _adjust_call_args(provider_name: str, model_name: str, kwargs: dict): |
| 960 | |
| 961 | # remap other to openai for litellm |
| 962 | if provider_name == "other": |
| 963 | provider_name = "openai" |
| 964 | |
| 965 | return provider_name, model_name, kwargs |
| 966 | |
| 967 | |
| 968 | def _merge_provider_defaults( |
| 969 | provider_type: ProviderModelType, original_provider: str, kwargs: dict |
| 970 | ) -> tuple[str, dict]: |
| 971 | provider_name = original_provider # default: unchanged |
| 972 | cfg = get_provider_config(provider_type, original_provider) |
| 973 | if cfg: |
| 974 | provider_name = cfg.get("litellm_provider", original_provider).lower() |
| 975 | |
| 976 | # Extra arguments nested under `kwargs` for readability |
| 977 | extra_kwargs = cfg.get("kwargs") if isinstance(cfg, dict) else None # type: ignore[arg-type] |
| 978 | if isinstance(extra_kwargs, dict): |
| 979 | for k, v in extra_kwargs.items(): |
| 980 | kwargs.setdefault(k, v) |
| 981 | |
| 982 | # Inject API key based on the *original* provider id if still missing |
| 983 | if "api_key" not in kwargs: |
| 984 | key = get_api_key(original_provider) |
| 985 | if key and key not in ("None", "NA"): |
| 986 | kwargs["api_key"] = key |
| 987 | |
| 988 | return provider_name, kwargs |
| 989 | |
| 990 | |
| 991 | def get_chat_model( |
| 992 | provider: str, name: str, model_config: Optional[ModelConfig] = None, **kwargs: Any |
| 993 | ) -> LiteLLMChatWrapper: |
| 994 | orig = provider.lower() |
| 995 | provider_name, kwargs = _merge_provider_defaults("chat", orig, kwargs) |
| 996 | return _get_litellm_chat( |
| 997 | LiteLLMChatWrapper, name, provider_name, model_config, **kwargs |
| 998 | ) |
| 999 | |
| 1000 | def get_embedding_model( |
| 1001 | provider: str, name: str, model_config: Optional[ModelConfig] = None, **kwargs: Any |
| 1002 | ) -> LiteLLMEmbeddingWrapper | LocalSentenceTransformerWrapper: |
| 1003 | orig = provider.lower() |
| 1004 | provider_name, kwargs = _merge_provider_defaults("embedding", orig, kwargs) |
| 1005 | return _get_litellm_embedding(name, provider_name, model_config, **kwargs) |