# Adapted from ReEvo: https://github.com/ai4co/reevo/blob/main/utils/llm_client/openai.py
# Licensed under the MIT License (see THIRD-PARTY-LICENSES.txt)
from dataclasses import dataclass
import logging
from llamda.llm_client.base import BaseClient, BaseLLMClientConfig
try:
from openai import OpenAI
except ImportError:
OpenAI = "openai" # type: ignore[misc,assignment]
logger = logging.getLogger(__name__)
[docs]
@dataclass
class OpenAIClientConfig(BaseLLMClientConfig):
base_url: str | None = None
api_key: str | None = None
[docs]
class OpenAIClient(BaseClient):
ClientClass = OpenAI
def __init__(
self,
config: OpenAIClientConfig,
) -> None:
super().__init__(config=config)
if isinstance(self.ClientClass, str):
logger.fatal(f"Package `{self.ClientClass}` is required")
exit(-1)
self.client = self.ClientClass(api_key=config.api_key, base_url=config.base_url)
def _chat_completion_api(
self, messages: list[dict], temperature: float, n: int = 1
) -> list:
response = self.client.chat.completions.create(
model=self.model,
messages=messages,
temperature=temperature,
n=n,
stream=False,
)
return response.choices