File size: 2,340 Bytes
39c930a d4ec913 39c930a d4ec913 39c930a d4ec913 39c930a d4ec913 39c930a d4ec913 39c930a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
"""Backend for OpenAI API."""
import json
import logging
import time
from .utils import FunctionSpec, OutputType, opt_messages_to_list, backoff_create
from funcy import notnone, once, select_values
import openai
logger = logging.getLogger("aide")
_client: openai.OpenAI = None # type: ignore
OPENAI_TIMEOUT_EXCEPTIONS = (
openai.RateLimitError,
openai.APIConnectionError,
openai.APITimeoutError,
openai.InternalServerError,
)
@once
def _setup_openai_client():
global _client
_client = openai.OpenAI(max_retries=0)
def query(
system_message: str | None,
user_message: str | None,
func_spec: FunctionSpec | None = None,
**model_kwargs,
) -> tuple[OutputType, float, int, int, dict]:
_setup_openai_client()
filtered_kwargs: dict = select_values(notnone, model_kwargs) # type: ignore
messages = opt_messages_to_list(system_message, user_message)
if func_spec is not None:
filtered_kwargs["tools"] = [func_spec.as_openai_tool_dict]
# force the model the use the function
filtered_kwargs["tool_choice"] = func_spec.openai_tool_choice_dict
t0 = time.time()
completion = backoff_create(
_client.chat.completions.create,
OPENAI_TIMEOUT_EXCEPTIONS,
messages=messages,
**filtered_kwargs,
)
req_time = time.time() - t0
choice = completion.choices[0]
if func_spec is None:
output = choice.message.content
else:
assert (
choice.message.tool_calls
), f"function_call is empty, it is not a function call: {choice.message}"
assert (
choice.message.tool_calls[0].function.name == func_spec.name
), "Function name mismatch"
try:
output = json.loads(choice.message.tool_calls[0].function.arguments)
except json.JSONDecodeError as e:
logger.error(
f"Error decoding the function arguments: {choice.message.tool_calls[0].function.arguments}"
)
raise e
in_tokens = completion.usage.prompt_tokens
out_tokens = completion.usage.completion_tokens
info = {
"system_fingerprint": completion.system_fingerprint,
"model": completion.model,
"created": completion.created,
}
return output, req_time, in_tokens, out_tokens, info
|