68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
import json
|
|
import openai
|
|
import os
|
|
import re
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
openai.api_key = os.getenv("OPENAI_API_KEY")
|
|
|
|
|
|
MAX_TOKENS = 4097
|
|
TOP_P = 0.9
|
|
FREQUENCY_PENALTY = 0.5
|
|
|
|
TRY_LIMIT = 1
|
|
|
|
try_count = 0
|
|
def process_response(input_string, quotation_check_field):
|
|
if '{' in input_string:
|
|
try:
|
|
# Find the index of the first occurrence of '{'
|
|
index = input_string.index('{')
|
|
# Extract everything after the first '{' (inclusive)
|
|
result = input_string[index:]
|
|
if re.search(r"'" + quotation_check_field + "':\s*'(.*?)'", result):
|
|
parsed_string = result.replace("\"", "\\\"")
|
|
pattern = r"(?<!\w)'|'(?!\w)"
|
|
parsed_string = re.sub(pattern, '"', parsed_string)
|
|
parsed_string = parsed_string.replace("\\\"", "'")
|
|
parsed_string = parsed_string.replace("\n\n", " ")
|
|
|
|
json_obj = json.loads(parsed_string)
|
|
return json_obj
|
|
else:
|
|
json_obj = json.loads(result)
|
|
return json_obj
|
|
except json.JSONDecodeError:
|
|
print("Invalid JSON string!")
|
|
else:
|
|
return input_string
|
|
|
|
def check_fields(obj, fields):
|
|
return all(field in obj for field in fields)
|
|
|
|
def make_openai_call(model, messages, token_count, fields_to_check, temperature):
|
|
global try_count
|
|
result = openai.ChatCompletion.create(
|
|
model=model,
|
|
max_tokens=int(MAX_TOKENS - token_count - 300),
|
|
temperature=float(temperature),
|
|
top_p=float(TOP_P),
|
|
frequency_penalty=float(FREQUENCY_PENALTY),
|
|
messages=messages
|
|
)
|
|
processed_response = process_response(result["choices"][0]["message"]["content"], fields_to_check[0])
|
|
if check_fields(processed_response, fields_to_check) is False and try_count < TRY_LIMIT:
|
|
try_count = try_count + 1
|
|
return make_openai_call(model, messages, token_count, fields_to_check, temperature)
|
|
elif try_count >= TRY_LIMIT:
|
|
try_count = 0
|
|
return result["choices"][0]["message"]["content"]
|
|
else:
|
|
try_count = 0
|
|
return processed_response
|
|
|
|
|