Add timeout/retry logic and refactor parameter building in BaseOpenAILLMService

- Add timeout (default 5.0s) and retry_on_timeout parameters to constructor
- Implement timeout/retry logic in get_chat_completions using asyncio.wait_for
- Extract build_chat_completion_params() as public method for subclass customization
This commit is contained in:
Mark Backman
2025-08-07 11:14:15 -04:00
parent 54a4d8a9f8
commit 5c86f8e687

View File

@@ -6,6 +6,7 @@
"""Base OpenAI LLM service implementation.""" """Base OpenAI LLM service implementation."""
import asyncio
import base64 import base64
import json import json
from typing import Any, Dict, List, Mapping, Optional from typing import Any, Dict, List, Mapping, Optional
@@ -14,6 +15,7 @@ import httpx
from loguru import logger from loguru import logger
from openai import ( from openai import (
NOT_GIVEN, NOT_GIVEN,
APITimeoutError,
AsyncOpenAI, AsyncOpenAI,
AsyncStream, AsyncStream,
DefaultAsyncHttpxClient, DefaultAsyncHttpxClient,
@@ -91,6 +93,8 @@ class BaseOpenAILLMService(LLMService):
project=None, project=None,
default_headers: Optional[Mapping[str, str]] = None, default_headers: Optional[Mapping[str, str]] = None,
params: Optional[InputParams] = None, params: Optional[InputParams] = None,
timeout: Optional[float] = 5.0,
retry_on_timeout: Optional[bool] = False,
**kwargs, **kwargs,
): ):
"""Initialize the BaseOpenAILLMService. """Initialize the BaseOpenAILLMService.
@@ -103,6 +107,8 @@ class BaseOpenAILLMService(LLMService):
project: OpenAI project ID. project: OpenAI project ID.
default_headers: Additional HTTP headers to include in requests. default_headers: Additional HTTP headers to include in requests.
params: Input parameters for model configuration and behavior. params: Input parameters for model configuration and behavior.
timeout: Request timeout in seconds. Defaults to 5.0 seconds.
retry_on_timeout: Whether to retry the request once if it times out.
**kwargs: Additional arguments passed to the parent LLMService. **kwargs: Additional arguments passed to the parent LLMService.
""" """
super().__init__(**kwargs) super().__init__(**kwargs)
@@ -119,6 +125,8 @@ class BaseOpenAILLMService(LLMService):
"max_completion_tokens": params.max_completion_tokens, "max_completion_tokens": params.max_completion_tokens,
"extra": params.extra if isinstance(params.extra, dict) else {}, "extra": params.extra if isinstance(params.extra, dict) else {},
} }
self._timeout = timeout
self._retry_on_timeout = retry_on_timeout
self.set_model_name(model) self.set_model_name(model)
self._client = self.create_client( self._client = self.create_client(
api_key=api_key, api_key=api_key,
@@ -175,7 +183,7 @@ class BaseOpenAILLMService(LLMService):
async def get_chat_completions( async def get_chat_completions(
self, context: OpenAILLMContext, messages: List[ChatCompletionMessageParam] self, context: OpenAILLMContext, messages: List[ChatCompletionMessageParam]
) -> AsyncStream[ChatCompletionChunk]: ) -> AsyncStream[ChatCompletionChunk]:
"""Get streaming chat completions from OpenAI API. """Get streaming chat completions from OpenAI API with optional timeout and retry.
Args: Args:
context: The LLM context containing tools and configuration. context: The LLM context containing tools and configuration.
@@ -184,6 +192,36 @@ class BaseOpenAILLMService(LLMService):
Returns: Returns:
Async stream of chat completion chunks. Async stream of chat completion chunks.
""" """
params = self.build_chat_completion_params(context, messages)
if self._retry_on_timeout:
try:
chunks = await asyncio.wait_for(
self._client.chat.completions.create(**params), timeout=self._timeout
)
return chunks
except (APITimeoutError, asyncio.TimeoutError):
# Retry, this time without a timeout so we get a response
chunks = await self._client.chat.completions.create(**params)
return chunks
else:
chunks = await self._client.chat.completions.create(**params)
return chunks
def build_chat_completion_params(
self, context: OpenAILLMContext, messages: List[ChatCompletionMessageParam]
) -> dict:
"""Build parameters for chat completion request.
Subclasses can override this to customize parameters for different providers.
Args:
context: The LLM context containing tools and configuration.
messages: List of chat completion messages to send.
Returns:
Dictionary of parameters for the chat completion request.
"""
params = { params = {
"model": self.model_name, "model": self.model_name,
"stream": True, "stream": True,
@@ -201,9 +239,7 @@ class BaseOpenAILLMService(LLMService):
} }
params.update(self._settings["extra"]) params.update(self._settings["extra"])
return params
chunks = await self._client.chat.completions.create(**params)
return chunks
async def _stream_chat_completions( async def _stream_chat_completions(
self, context: OpenAILLMContext self, context: OpenAILLMContext