Implements LLMProvider + BatchProvider for OpenAI.
| 53 | |
| 54 | |
| 55 | class OpenAIProvider: |
| 56 | """Implements LLMProvider + BatchProvider for OpenAI.""" |
| 57 | |
| 58 | def __init__( |
| 59 | self, |
| 60 | model: str, |
| 61 | *, |
| 62 | reasoning_effort: str | None = None, |
| 63 | stall_timeout: int = STALL_TIMEOUT, |
| 64 | ) -> None: |
| 65 | self._model = model |
| 66 | if model.startswith("openai/"): |
| 67 | self._backend = "openai" |
| 68 | self._api_model = model.removeprefix("openai/") |
| 69 | elif model.startswith("azure/"): |
| 70 | self._backend = "azure" |
| 71 | self._api_model = model.removeprefix("azure/") |
| 72 | else: |
| 73 | raise ValueError(f"unsupported OpenAI-compatible model prefix: {model!r}") |
| 74 | self._reasoning_effort = reasoning_effort |
| 75 | self._stall_timeout = stall_timeout |
| 76 | if self._backend == "azure": |
| 77 | api_key = os.environ.get("AZURE_OPENAI_API_KEY") |
| 78 | if not api_key: |
| 79 | raise ValueError("azure/ models require AZURE_OPENAI_API_KEY") |
| 80 | self.client = OpenAI( |
| 81 | api_key=api_key, |
| 82 | base_url=_azure_base_url(), |
| 83 | timeout=LLM_TIMEOUT_SECONDS, |
| 84 | ) |
| 85 | else: |
| 86 | self.client = OpenAI(timeout=LLM_TIMEOUT_SECONDS) |
| 87 | |
| 88 | def call(self, user_content: str) -> tuple[str, TokenUsage]: |
| 89 | kwargs: dict = {} |
| 90 | if self._reasoning_effort: |
| 91 | kwargs["reasoning"] = {"effort": self._reasoning_effort} |
| 92 | response = self.client.responses.create( |
| 93 | model=self._api_model, |
| 94 | input=_openai_input(user_content), |
| 95 | text={"format": {"type": "json_object"}}, |
| 96 | **kwargs, |
| 97 | ) |
| 98 | usage = TokenUsage() |
| 99 | if response.usage: |
| 100 | usage.input_tokens = response.usage.input_tokens |
| 101 | usage.output_tokens = response.usage.output_tokens |
| 102 | usage.reasoning_tokens = ( |
| 103 | response.usage.output_tokens_details.reasoning_tokens |
| 104 | ) |
| 105 | return response.output_text, usage |
| 106 | |
| 107 | @property |
| 108 | def retryable_exceptions(self) -> tuple[type[Exception], ...]: |
| 109 | return ( |
| 110 | openai.RateLimitError, |
| 111 | openai.APITimeoutError, |
| 112 | openai.APIConnectionError, |
no outgoing calls