chore: 修改包名

pkg/provider/api/chat_completion.py

@@ -0,0 +1,232 @@
+import json
+import logging
+
+import openai
+from openai.types.chat import chat_completion_message
+
+from .model import RequestBase
+from .. import funcmgr
+from ...plugin import host
+from ...utils import context
+
+
+class ChatCompletionRequest(RequestBase):
+    """调用ChatCompletion接口的请求类。
+    
+    此类保证每一次返回的角色为assistant的信息的finish_reason一定为stop。
+    若有函数调用响应，本类的返回瀑布是：函数调用请求->函数调用结果->...->assistant的信息->stop。
+    """
+
+    model: str
+    messages: list[dict[str, str]]
+    kwargs: dict
+
+    stopped: bool = False
+
+    pending_func_call: chat_completion_message.FunctionCall = None
+
+    pending_msg: str
+
+    def flush_pending_msg(self):
+        self.append_message(
+            role="assistant",
+            content=self.pending_msg
+        )
+        self.pending_msg = ""
+
+    def append_message(self, role: str, content: str, name: str=None, function_call: dict=None):
+        msg = {
+            "role": role,
+            "content": content
+        }
+
+        if name is not None:
+            msg['name'] = name
+
+        if function_call is not None:
+            msg['function_call'] = function_call
+
+        self.messages.append(msg)
+
+    def __init__(
+        self,
+        client: openai.Client,
+        model: str,
+        messages: list[dict[str, str]],
+        **kwargs
+    ):
+        self.client = client
+        self.model = model
+        self.messages = messages.copy()
+
+        self.kwargs = kwargs
+
+        self.req_func = self.client.chat.completions.create
+
+        self.pending_func_call = None
+
+        self.stopped = False
+
+        self.pending_msg = ""
+
+    def __iter__(self):
+        return self
+
+    def __next__(self) -> dict:
+        if self.stopped:
+            raise StopIteration()
+
+        if self.pending_func_call is None:  # 没有待处理的函数调用请求
+
+            args = {
+                "model": self.model,
+                "messages": self.messages,
+            }
+
+            funcs = funcmgr.get_func_schema_list()
+
+            if len(funcs) > 0:
+                args['functions'] = funcs
+
+            # 拼接kwargs
+            args = {**args, **self.kwargs}
+            
+            from openai.types.chat import chat_completion
+
+            resp: chat_completion.ChatCompletion = self._req(**args)
+
+            choice0 = resp.choices[0]
+
+            # 如果不是函数调用，且finish_reason为stop，则停止迭代
+            if choice0.finish_reason == 'stop':  #  and choice0["finish_reason"] == "stop"
+                self.stopped = True
+            
+            if hasattr(choice0.message, 'function_call') and choice0.message.function_call is not None:
+                self.pending_func_call = choice0.message.function_call
+
+                self.append_message(
+                    role="assistant",
+                    content=choice0.message.content,
+                    function_call=choice0.message.function_call
+                )
+
+                return {
+                    "id": resp.id,
+                    "choices": [
+                        {
+                            "index": choice0.index,
+                            "message": {
+                                "role": "assistant",
+                                "type": "function_call",
+                                "content": choice0.message.content,
+                                "function_call": {
+                                    "name": choice0.message.function_call.name,
+                                    "arguments": choice0.message.function_call.arguments
+                                }
+                            },
+                            "finish_reason": "function_call"
+                        }
+                    ],
+                    "usage": {
+                        "prompt_tokens": resp.usage.prompt_tokens,
+                        "completion_tokens": resp.usage.completion_tokens,
+                        "total_tokens": resp.usage.total_tokens
+                    }
+                }
+            else:
+
+                # self.pending_msg += choice0['message']['content']
+                # 普通回复一定处于最后方，故不用再追加进内部messages
+
+                return {
+                    "id": resp.id,
+                    "choices": [
+                        {
+                            "index": choice0.index,
+                            "message": {
+                                "role": "assistant",
+                                "type": "text",
+                                "content": choice0.message.content
+                            },
+                            "finish_reason": choice0.finish_reason
+                        }
+                    ],
+                    "usage": {
+                        "prompt_tokens": resp.usage.prompt_tokens,
+                        "completion_tokens": resp.usage.completion_tokens,
+                        "total_tokens": resp.usage.total_tokens
+                    }
+                }
+        else:  # 处理函数调用请求
+
+            cp_pending_func_call = self.pending_func_call.copy()
+            
+            self.pending_func_call = None
+
+            func_name = cp_pending_func_call.name
+            arguments = {}
+
+            try:
+
+                try:
+                    arguments = json.loads(cp_pending_func_call.arguments)
+                # 若不是json格式的异常处理
+                except json.decoder.JSONDecodeError:
+                    # 获取函数的参数列表
+                    func_schema = funcmgr.get_func_schema(func_name)
+
+                    arguments = {
+                        func_schema['parameters']['required'][0]: cp_pending_func_call.arguments
+                    }
+                
+                logging.info("执行函数调用: name={}, arguments={}".format(func_name, arguments))
+
+                # 执行函数调用
+                ret = ""
+                try:
+                    ret = funcmgr.execute_function(func_name, arguments)
+
+                    logging.info("函数执行完成。")
+                except Exception as e:
+                    ret = "error: execute function failed: {}".format(str(e))
+                    logging.error("函数执行失败: {}".format(str(e)))
+
+                # 上报数据
+                plugin_info = host.get_plugin_info_for_audit(func_name.split('-')[0])
+                audit_func_name = func_name.split('-')[1]
+                audit_func_desc = funcmgr.get_func_schema(func_name)['description']
+                context.get_center_v2_api().usage.post_function_record(
+                    plugin=plugin_info,
+                    function_name=audit_func_name,
+                    function_description=audit_func_desc,
+                )
+
+                self.append_message(
+                    role="function",
+                    content=json.dumps(ret, ensure_ascii=False),
+                    name=func_name
+                )
+
+                return {
+                    "id": -1,
+                    "choices": [
+                        {
+                            "index": -1,
+                            "message": {
+                                "role": "function",
+                                "type": "function_return",
+                                "function_name": func_name,
+                                "content": json.dumps(ret, ensure_ascii=False)
+                            },
+                            "finish_reason": "function_return"
+                        }
+                    ],
+                    "usage": {
+                        "prompt_tokens": 0,
+                        "completion_tokens": 0,
+                        "total_tokens": 0
+                    }
+                }
+
+            except funcmgr.ContentFunctionNotFoundError:
+                raise Exception("没有找到函数: {}".format(func_name))

pkg/provider/api/completion.py

@@ -0,0 +1,100 @@
+import openai
+from openai.types import completion, completion_choice
+
+from . import model
+
+
+class CompletionRequest(model.RequestBase):
+    """调用Completion接口的请求类。
+    
+    调用方可以一直next completion直到finish_reason为stop。
+    """
+
+    model: str
+    prompt: str
+    kwargs: dict
+
+    stopped: bool = False
+
+    def __init__(
+        self,
+        client: openai.Client,
+        model: str,
+        messages: list[dict[str, str]],
+        **kwargs
+    ):
+        self.client = client
+        self.model = model
+        self.prompt = ""
+        
+        for message in messages:
+            self.prompt += message["role"] + ": " + message["content"] + "\n"
+        
+        self.prompt += "assistant: "
+
+        self.kwargs = kwargs
+
+        self.req_func = self.client.completions.create
+
+    def __iter__(self):
+        return self
+
+    def __next__(self) -> dict:
+        """调用Completion接口，返回生成的文本
+        
+        {
+            "id": "id",
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {
+                        "role": "assistant",
+                        "type": "text",
+                        "content": "message"
+                    },
+                    "finish_reason": "reason"
+                }
+            ],
+            "usage": {
+                "prompt_tokens": 10,
+                "completion_tokens": 20,
+                "total_tokens": 30
+            }
+        }
+        """
+
+        if self.stopped:
+            raise StopIteration()
+
+        resp: completion.Completion = self._req(
+            model=self.model,
+            prompt=self.prompt,
+            **self.kwargs
+        )
+
+        if resp.choices[0].finish_reason == "stop":
+            self.stopped = True
+
+        choice0: completion_choice.CompletionChoice = resp.choices[0]
+
+        self.prompt += choice0.text
+
+        return {
+            "id": resp.id,
+            "choices": [
+                {
+                    "index": choice0.index,
+                    "message": {
+                        "role": "assistant",
+                        "type": "text",
+                        "content": choice0.text
+                    },
+                    "finish_reason": choice0.finish_reason
+                }
+            ],
+            "usage": {
+                "prompt_tokens": resp.usage.prompt_tokens,
+                "completion_tokens": resp.usage.completion_tokens,
+                "total_tokens": resp.usage.total_tokens
+            }
+        }

pkg/provider/api/model.py

@@ -0,0 +1,40 @@
+# 定义不同接口请求的模型
+import logging
+
+import openai
+
+from ...utils import context
+
+
+class RequestBase:
+
+    client: openai.Client
+
+    req_func: callable
+
+    def __init__(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def _next_key(self):
+        switched, name = context.get_openai_manager().key_mgr.auto_switch()
+        logging.debug("切换api-key: switched={}, name={}".format(switched, name))
+        self.client.api_key = context.get_openai_manager().key_mgr.get_using_key()
+
+    def _req(self, **kwargs):
+        """处理代理问题"""
+        logging.debug("请求接口参数: %s", str(kwargs))
+        config = context.get_config_manager().data
+
+        ret = self.req_func(**kwargs)
+        logging.debug("接口请求返回：%s", str(ret))
+
+        if config['switch_strategy'] == 'active':
+            self._next_key()
+
+        return ret
+
+    def __iter__(self):
+        raise self
+
+    def __next__(self):
+        raise NotImplementedError