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:
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user