workflows/libs/llm_api/openai.py

# flake8: noqa: E402
import json
import os
import re
import sys
from typing import Dict, List

import openai

sys.path.append(os.path.join(os.path.dirname(__file__), ".."))

from ide_services.services import log_warn


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块的内容并返回
        _ = match.group(1)  # language
        markdown_content = match.group(2)
        return markdown_content.strip()
    else:
        # 如果匹配失败，则返回原始内容
        return content


def chat_completion_stream(
    messages: List[Dict],  # [{"role": "user", "content": "hello"}]
    llm_config: Dict,  # {"model": "...", ...}
    error_out: bool = True,
    stream_out=False,
) -> str:
    """
    通过ChatCompletion API获取OpenAI聊天机器人的回复。

    Args:
        messages: 一个列表，包含用户输入的消息。
        llm_config: 一个字典，包含ChatCompletion API的配置信息。
        error_out: 如果为True，遇到异常时输出错误信息并返回None，否则返回None。

    Returns:
        如果成功获取到聊天机器人的回复，返回一个字符串类型的回复消息。如果连接失败，则返回None。

    """
    for try_times in range(3):
        try:
            client = openai.OpenAI(
                api_key=os.environ.get("OPENAI_API_KEY", None),
                base_url=os.environ.get("OPENAI_API_BASE", None),
            )

            llm_config["stream"] = True
            llm_config["timeout"] = 8
            response = client.chat.completions.create(messages=messages, **llm_config)

            response_result = {"content": None, "function_name": None, "parameters": ""}
            for chunk in response:  # pylint: disable=E1133
                chunk = chunk.dict()
                delta = chunk["choices"][0]["delta"]
                if "tool_calls" in delta and delta["tool_calls"]:
                    tool_call = delta["tool_calls"][0]["function"]
                    if tool_call.get("name", None):
                        response_result["function_name"] = tool_call["name"]
                    if tool_call.get("arguments", None):
                        response_result["parameters"] += tool_call["arguments"]
                if delta.get("content", None):
                    if stream_out:
                        print(delta["content"], end="", flush=True)
                    if response_result["content"]:
                        response_result["content"] += delta["content"]
                    else:
                        response_result["content"] = delta["content"]
            return response_result
        except (openai.APIConnectionError, openai.APITimeoutError) as err:
            log_warn(f"Exception: {err.__class__.__name__}: {err}")
            if try_times >= 2:
                return {"content": None, "function_name": None, "parameters": "", "error": err}
            continue
        except openai.APIError as err:
            if error_out:
                print("Exception:", err, file=sys.stderr, flush=True)
            return {"content": None, "function_name": None, "parameters": "", "error": err}
        except Exception as err:
            if error_out:
                print("Exception:", err, file=sys.stderr, flush=True)
            return {"content": None, "function_name": None, "parameters": "", "error": err}


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_stream(messages, llm_config)
        if not response["content"]:
            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 json.JSONDecodeError:
            log_warn(f"JSONDecodeError: {response['content']}")
            continue
        except Exception as err:
            if error_out:
                print("Exception: ", err, file=sys.stderr, flush=True)
            return None
    if error_out:
        print("Not valid json response:", response["content"], file=sys.stderr, flush=True)
    return None