c58ef78d52
- Updated the response format to include file paths in filename keys - Added docstrings to chat_completion functions for better clarity
96 lines
3.5 KiB
Python
96 lines
3.5 KiB
Python
import re
|
||
import sys
|
||
import json
|
||
|
||
import openai
|
||
|
||
|
||
def _try_remove_markdown_block_flag(content):
|
||
"""
|
||
如果content是一个markdown块，则删除它的头部```xxx和尾部```
|
||
"""
|
||
# 定义正则表达式模式，用于匹配markdown块的头部和尾部
|
||
pattern = r'^\s*```\s*(\w+)\s*\n(.*?)\n\s*```\s*$'
|
||
|
||
# 使用re模块进行匹配
|
||
match = re.search(pattern, content, re.DOTALL | re.MULTILINE)
|
||
|
||
if match:
|
||
# 如果匹配成功，则提取出markdown块的内容并返回
|
||
language = match.group(1)
|
||
markdown_content = match.group(2)
|
||
return markdown_content.strip()
|
||
else:
|
||
# 如果匹配失败，则返回原始内容
|
||
return content
|
||
|
||
|
||
def chat_completion_no_stream(messages, llm_config, error_out: bool=True) -> str:
|
||
"""
|
||
通过ChatCompletion API获取OpenAI聊天机器人的回复。
|
||
|
||
Args:
|
||
messages: 一个列表，包含用户输入的消息。
|
||
llm_config: 一个字典，包含ChatCompletion API的配置信息。
|
||
error_out: 如果为True，遇到异常时输出错误信息并返回None，否则返回None。
|
||
|
||
Returns:
|
||
如果成功获取到聊天机器人的回复，返回一个字符串类型的回复消息。如果连接失败，则返回None。
|
||
|
||
"""
|
||
connection_error = ''
|
||
for _1 in range(3):
|
||
try:
|
||
response = openai.ChatCompletion.create(
|
||
messages=messages,
|
||
**llm_config,
|
||
stream=False
|
||
)
|
||
|
||
response_dict = json.loads(str(response))
|
||
if 'choices' not in response_dict:
|
||
if error_out:
|
||
print("Response Error:", response_dict, file=sys.stderr, flush=True)
|
||
return None
|
||
respose_message = response_dict["choices"][0]["message"]
|
||
# print("=> llm response:", respose_message, end="\n\n")
|
||
return respose_message
|
||
except ConnectionError as err:
|
||
connection_error = err
|
||
continue
|
||
except Exception as err:
|
||
if error_out:
|
||
print("Exception:", err, file=sys.stderr, flush=True)
|
||
return None
|
||
if error_out:
|
||
print("Connect Error:", connection_error, file=sys.stderr, flush=True)
|
||
return None
|
||
|
||
def chat_completion_no_stream_return_json(messages, llm_config, error_out: bool=True):
|
||
"""
|
||
尝试三次从聊天完成API获取结果，并返回JSON对象。如果无法解析JSON，将尝试三次，直到出现错误或达到最大尝试次数。
|
||
|
||
Args:
|
||
messages (List[str]): 用户输入的消息列表。
|
||
llm_config (Dict[str, Any]): 聊天配置字典。
|
||
error_out (bool, optional): 如果为True，则如果出现错误将打印错误消息并返回None。默认为True。
|
||
|
||
Returns:
|
||
Dict[str, Any]: 从聊天完成API获取的JSON对象。如果无法解析JSON或达到最大尝试次数，则返回None。
|
||
"""
|
||
for _1 in range(3):
|
||
response = chat_completion_no_stream(messages, llm_config)
|
||
if response is None:
|
||
return None
|
||
|
||
try:
|
||
# json will format as ```json ... ``` in 1106 model
|
||
response_content = _try_remove_markdown_block_flag(response["content"])
|
||
response_obj = json.loads(response_content)
|
||
return response_obj
|
||
except Exception:
|
||
continue
|
||
if error_out:
|
||
print("Not valid json response:", response["content"], file=sys.stderr, flush=True)
|
||
return None
|