fiatrete
135 lines
4.8 KiB
Python
135 lines
4.8 KiB
Python
import asyncio
|
|
|
|
import openai
|
|
from openai.error import RateLimitError, APIError, Timeout
|
|
|
|
from jarvis import CFG
|
|
from jarvis.gpt.message import Message
|
|
from jarvis.logger import logger
|
|
from typing import Callable
|
|
|
|
openai.api_key = CFG.openai_api_key
|
|
if CFG.openai_url_base is not None:
|
|
openai.api_base = CFG.openai_url_base
|
|
|
|
print_total_cost = CFG.debug_mode
|
|
|
|
|
|
async def acreate_chat_completion_once(
|
|
messages: list, # type: ignore
|
|
model: str | None = None,
|
|
temperature: float = CFG.temperature,
|
|
max_tokens: int | None = None,
|
|
deployment_id=None,
|
|
request_timeout=40,
|
|
) -> str:
|
|
"""
|
|
Create a chat completion and update the cost.
|
|
Args:
|
|
messages (list): The list of messages to send to the API.
|
|
model (str): The model to use for the API call.
|
|
temperature (float): The temperature to use for the API call.
|
|
max_tokens (int): The maximum number of tokens for the API call.
|
|
Returns:
|
|
str: The AI's response.
|
|
"""
|
|
if deployment_id is not None:
|
|
response = await openai.ChatCompletion.acreate(
|
|
deployment_id=deployment_id,
|
|
model=model,
|
|
messages=messages,
|
|
temperature=temperature,
|
|
max_tokens=max_tokens,
|
|
request_timeout=request_timeout
|
|
)
|
|
else:
|
|
response = await openai.ChatCompletion.acreate(
|
|
model=model,
|
|
messages=messages,
|
|
temperature=temperature,
|
|
max_tokens=max_tokens,
|
|
request_timeout=request_timeout
|
|
)
|
|
if CFG.debug_mode:
|
|
logger.debug(f"Response: {response}")
|
|
# prompt_tokens = response.usage.prompt_tokens
|
|
# completion_tokens = response.usage.completion_tokens
|
|
return response
|
|
|
|
|
|
# Overly simple abstraction until we create something better
|
|
# simple retry mechanism when getting a rate error or a bad gateway
|
|
async def acreate_chat_completion(
|
|
messages: list[Message], # type: ignore
|
|
model: str = None,
|
|
temperature: float = CFG.temperature,
|
|
max_tokens: int = None,
|
|
request_timeout: int = 40,
|
|
num_retries=3,
|
|
on_single_request_timeout: Callable = None
|
|
):
|
|
"""Create a chat completion using the OpenAI API
|
|
|
|
Args:
|
|
messages (List[Message]): The messages to send to the chat completion
|
|
model (str, optional): The model to use. Defaults to None.
|
|
temperature (float, optional): The temperature to use. Defaults to 0.9.
|
|
max_tokens (int, optional): The max tokens to use. Defaults to None.
|
|
request_timeout (int, optional): The request_timeout of a single openai request.
|
|
num_retries (int, optional): The max retries.
|
|
on_single_request_timeout (Callable, optional): This function will be called each time a single openai request
|
|
timeout, must be an async function, the last timeout will not emit callback.
|
|
|
|
Returns:
|
|
str: The response from the chat completion
|
|
"""
|
|
if CFG.debug_mode:
|
|
logger.debug(
|
|
f"Creating chat completion with model {model}, temperature {temperature}, max_tokens {max_tokens}"
|
|
)
|
|
|
|
response = None
|
|
|
|
for attempt in range(num_retries):
|
|
backoff = min(2 ** (attempt + 2), 8)
|
|
try:
|
|
if CFG.use_azure:
|
|
response = await acreate_chat_completion_once(
|
|
deployment_id=CFG.get_azure_deployment_id_for_model(model),
|
|
model=model,
|
|
messages=messages,
|
|
temperature=temperature,
|
|
max_tokens=max_tokens,
|
|
request_timeout=request_timeout,
|
|
)
|
|
else:
|
|
response = await acreate_chat_completion_once(
|
|
model=model,
|
|
messages=messages,
|
|
temperature=temperature,
|
|
max_tokens=max_tokens,
|
|
request_timeout=request_timeout,
|
|
)
|
|
break
|
|
except RateLimitError:
|
|
if CFG.debug_mode:
|
|
logger.debug(f"Error: Reached rate limit, passing...")
|
|
except (APIError, Timeout) as e:
|
|
if isinstance(e, Timeout):
|
|
if on_single_request_timeout:
|
|
await on_single_request_timeout(num_retries < num_retries - 1)
|
|
if e.http_status != 502:
|
|
raise
|
|
if attempt == num_retries - 1:
|
|
raise
|
|
if CFG.debug_mode:
|
|
logger.debug(
|
|
f"Error: API Bad gateway. Waiting {backoff} seconds..."
|
|
)
|
|
await asyncio.sleep(backoff)
|
|
if response is None:
|
|
logger.error(f"Failed to get response from GPT after {num_retries} retries")
|
|
raise RuntimeError(f"Failed to get response after {num_retries} retries")
|
|
resp = response.choices[0].message["content"]
|
|
return resp
|