| 1 | from __future__ import annotations |
| 2 | |
| 3 | import asyncio |
| 4 | import json |
| 5 | import logging |
| 6 | from dataclasses import dataclass, field |
| 7 | from typing import Any |
| 8 | from uuid import uuid4 |
| 9 | |
| 10 | from openai import APIConnectionError, APITimeoutError, AsyncOpenAI |
| 11 | |
| 12 | from .config import llm_api_key, llm_base_url, llm_model |
| 13 | from .models import ToolCall |
| 14 | |
| 15 | |
| 16 | LLM_MAX_ATTEMPTS = 3 |
| 17 | LLM_RETRY_BACKOFF_SECONDS = (1.0, 4.0) |
| 18 | LLM_REQUEST_TIMEOUT_SECONDS = 300.0 |
| 19 | |
| 20 | |
| 21 | def _is_retryable_llm_error(exc: BaseException) -> bool: |
| 22 | status = getattr(exc, "status_code", None) |
| 23 | if status is not None: |
| 24 | try: |
| 25 | status = int(status) |
| 26 | except (TypeError, ValueError): |
| 27 | return False |
| 28 | return status == 429 or status >= 500 |
| 29 | return isinstance(exc, (APIConnectionError, APITimeoutError)) |
| 30 | |
| 31 | |
| 32 | class LLMResponseShapeError(RuntimeError): |
| 33 | pass |
| 34 | |
| 35 | |
| 36 | @dataclass(slots=True) |
| 37 | class AssistantMessage: |
| 38 | text: str = "" |
| 39 | tool_calls: list[ToolCall] = field(default_factory=list) |
| 40 | raw_message: dict[str, Any] = field(default_factory=dict) |
| 41 | |
| 42 | |
| 43 | class OpenAICompatibleLLM: |
| 44 | def __init__(self, model: str | None = None, base_url: str | None = None, api_key: str | None = None) -> None: |
| 45 | self.model = model or llm_model() |
| 46 | self.base_url = base_url or llm_base_url() |
| 47 | self.api_key = api_key or llm_api_key() |
| 48 | if not self.api_key: |
| 49 | raise RuntimeError("VIMAX_LLM_API_KEY is required for the agent LLM client") |
| 50 | self.client = AsyncOpenAI(api_key=self.api_key, base_url=self.base_url, timeout=LLM_REQUEST_TIMEOUT_SECONDS) |
| 51 | |
| 52 | async def complete(self, messages: list[dict[str, Any]], tools: list[dict[str, Any]]) -> AssistantMessage: |
| 53 | shape_attempts = [ |
| 54 | {"tools": tools or None, "tool_choice": "auto" if tools else None}, |
| 55 | {"tools": tools or None, "tool_choice": "auto" if tools else None}, |
| 56 | ] |
| 57 | if tools: |
| 58 | shape_attempts.append({"tools": None, "tool_choice": None}) |
| 59 | |
| 60 | last_shape_error: Exception | None = None |
| 61 | for attempt in shape_attempts: |
| 62 | try: |
| 63 | response = await self._create_completion_with_retries(messages, attempt["tools"], attempt["tool_choice"]) |
| 64 | return _assistant_message_from_response(response) |
| 65 | except LLMResponseShapeError as exc: |
| 66 | last_shape_error = exc |
| 67 | continue |
| 68 | assert last_shape_error is not None |
| 69 | raise last_shape_error |
| 70 | |
| 71 | async def _create_completion_with_retries(self, messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None, tool_choice: str | None) -> Any: |
| 72 | for attempt in range(LLM_MAX_ATTEMPTS): |
| 73 | try: |
| 74 | return await self._create_completion(messages, tools, tool_choice) |
| 75 | except Exception as exc: |
| 76 | if isinstance(exc, LLMResponseShapeError) or attempt == LLM_MAX_ATTEMPTS - 1 or not _is_retryable_llm_error(exc): |
| 77 | raise |
| 78 | delay = LLM_RETRY_BACKOFF_SECONDS[min(attempt, len(LLM_RETRY_BACKOFF_SECONDS) - 1)] |
| 79 | logging.warning("LLM call failed (%s); retrying in %.1fs (attempt %d/%d)", exc, delay, attempt + 1, LLM_MAX_ATTEMPTS) |
| 80 | await asyncio.sleep(delay) |
| 81 | |
| 82 | async def _create_completion(self, messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None, tool_choice: str | None) -> Any: |
| 83 | kwargs: dict[str, Any] = { |
| 84 | "model": self.model, |
| 85 | "messages": messages, |
| 86 | "stream": False, |
| 87 | } |
| 88 | if tools: |
| 89 | kwargs["tools"] = tools |
| 90 | if tool_choice: |
| 91 | kwargs["tool_choice"] = tool_choice |
| 92 | return await self.client.chat.completions.create(**kwargs) |
| 93 | |
| 94 | |
| 95 | def _assistant_message_from_response(response: Any) -> AssistantMessage: |
| 96 | message = _extract_message(response) |
| 97 | text = _message_value(message, "content") or "" |
| 98 | calls: list[ToolCall] = [] |
| 99 | for call in _message_value(message, "tool_calls") or []: |
| 100 | function = _message_value(call, "function") or {} |
| 101 | try: |
| 102 | arguments = json.loads(_message_value(function, "arguments") or "{}") |
| 103 | except json.JSONDecodeError: |
| 104 | arguments = {} |
| 105 | calls.append(ToolCall(id=_message_value(call, "id") or f"tool-{uuid4().hex[:12]}", name=_message_value(function, "name"), arguments=arguments)) |
| 106 | return AssistantMessage(text=text, tool_calls=calls, raw_message=_dump_message(message)) |
| 107 | |
| 108 | |
| 109 | def _extract_message(response: Any) -> Any: |
| 110 | if isinstance(response, str): |
| 111 | try: |
| 112 | response = json.loads(response) |
| 113 | except json.JSONDecodeError as exc: |
| 114 | raise LLMResponseShapeError(f"LLM provider returned a string instead of a chat completion object: {response[:300]}") from exc |
| 115 | choices = _message_value(response, "choices") |
| 116 | if not choices: |
| 117 | raise LLMResponseShapeError(f"LLM provider response missing choices: {str(response)[:500]}") |
| 118 | first_choice = choices[0] |
| 119 | message = _message_value(first_choice, "message") |
| 120 | if message is None: |
| 121 | raise LLMResponseShapeError(f"LLM provider response missing choice.message: {str(response)[:500]}") |
| 122 | return message |
| 123 | |
| 124 | |
| 125 | def _message_value(obj: Any, key: str) -> Any: |
| 126 | if isinstance(obj, dict): |
| 127 | return obj.get(key) |
| 128 | return getattr(obj, key, None) |
| 129 | |
| 130 | |
| 131 | def _dump_message(message: Any) -> dict[str, Any]: |
| 132 | if isinstance(message, dict): |
| 133 | return message |
| 134 | if hasattr(message, "model_dump"): |
| 135 | return message.model_dump() |
| 136 | return {"content": str(message)} |
| 137 |