unionllm 0.1.0__py3-none-any.whl

cnlitellm/api.py

@@ -0,0 +1,40 @@
+import logging
+from typing import Any, List
+from .providers import base_provider
+from .providers import zhipu, moonshot, minimax, qwen, tiangong, baichuan, wenxin, xunfei, dify, fastgpt, coze
+from .exceptions import ProviderError
+
+logger = logging.getLogger(__name__)
+
+class CNLiteLLM:
+    def __init__(self, provider: str, **kwargs):
+        self.provider = provider.lower()
+        if self.provider == "zhipuai":
+            self.provider_instance = zhipu.ZhipuAIProvider(**kwargs)
+        elif self.provider == "moonshot":
+            self.provider_instance = moonshot.MoonshotAIProvider(**kwargs)
+        elif self.provider == "minimax":
+            self.provider_instance = minimax.MinimaxAIProvider(**kwargs)
+        elif self.provider == "qwen":
+            self.provider_instance = qwen.QwenAIProvider(**kwargs)
+        elif self.provider == "tiangong":
+            self.provider_instance = tiangong.TianGongAIProvider(**kwargs)
+        elif self.provider == "baichuan":
+            self.provider_instance = baichuan.BaiChuanAIProvider(**kwargs)
+        elif self.provider == "wenxin":
+            self.provider_instance = wenxin.WenXinAIProvider(**kwargs)
+        elif self.provider == "xunfei":
+            self.provider_instance = xunfei.XunfeiAIProvider(**kwargs)
+        elif self.provider == "dify":
+            self.provider_instance = dify.DifyAIProvider(**kwargs)
+        elif self.provider == "fastgpt":
+            self.provider_instance = fastgpt.FastGPTProvider(**kwargs)
+        elif self.provider == "coze":
+            self.provider_instance = coze.CozeAIProvider(**kwargs)
+        else:
+            raise ProviderError(f"Provider '{self.provider}' is not supported.")
+
+    def completion(self, model: str, messages: List[str], **kwargs) -> Any:
+        if not self.provider_instance:
+            raise ProviderError(f"Provider '{self.provider}' is not initialized.")
+        return self.provider_instance.completion(model, messages, **kwargs)

cnlitellm/exceptions.py

@@ -0,0 +1,11 @@
+class CNLiteLLMError(Exception):
+    """Base exception class for CNLiteLLM."""
+    pass
+
+class ProviderError(CNLiteLLMError):
+    """Exception raised when there is an issue with the model provider."""
+    pass
+
+class APICallError(CNLiteLLMError):
+    """Exception raised when an API call fails."""
+    pass

cnlitellm/models.py

@@ -0,0 +1,28 @@
+from typing import List, Optional
+
+class ResponseModel:
+    def __init__(self, prompt_tokens: List[str], completion_tokens: List[str], total_tokens: int, total_attempts: int, raw_response: Optional[dict] = None, usage: Optional[dict] = None):
+        self.prompt_tokens = prompt_tokens
+        self.completion_tokens = completion_tokens
+        self.total_tokens = total_tokens
+        self.total_attempts = total_attempts
+        self.raw_response = raw_response
+        self.usage = usage
+
+    def to_dict(self) -> dict:
+        return {
+            "prompt_tokens": self.prompt_tokens,
+            "completion_tokens": self.completion_tokens,
+            "total_tokens": self.total_tokens,
+            "total_attempts": self.total_attempts,
+            "raw_response": self.raw_response,
+            "usage": self.usage
+        }
+
+    def get_completions(self) -> list[str]:
+        completions = []
+        if self.raw_response and "choices" in self.raw_response:
+            for choice in self.raw_response["choices"]:
+                if "message" in choice and "content" in choice["message"]:
+                    completions.append(choice["message"]["content"])
+        return completions

cnlitellm/providers/baichuan.py

@@ -0,0 +1,128 @@
+import json
+from .base_provider import BaseProvider
+from cnlitellm.utils import ModelResponse, Message, Choices, Usage, Delta, StreamingChoices
+import requests
+import logging
+
+class BaiChuanOpenAIError(Exception):
+    def __init__(
+        self,
+        status_code,
+        message,
+    ):
+        self.status_code = status_code
+        self.message = message
+        super().__init__(self.message)
+
+class BaiChuanAIProvider(BaseProvider):
+    def __init__(self, **model_kwargs):
+        self.api_key = model_kwargs.get("api_key")
+        self.endpoint_url = "https://api.baichuan-ai.com/v1/chat/completions"
+
+    def pre_processing(self, **kwargs):
+        supported_params = [
+            "model", "messages", "max_tokens", "temperature", "top_p", "n",
+            "logprobs", "stream", "stop", "presence_penalty", "frequency_penalty",
+            "best_of", "logit_bias"
+        ]
+        for key in list(kwargs.keys()):
+            if key not in supported_params:
+                kwargs.pop(key)
+        return kwargs
+
+    def post_stream_processing_wrapper(self, model, messages, **new_kwargs):
+        payload = json.dumps({"model": model, "messages": messages, **new_kwargs})
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+        response = requests.post(self.endpoint_url, headers=headers, data=payload)
+        for line in response.iter_lines():
+            if line:
+                new_line = line.decode("utf-8").replace("data: ", "")
+                if new_line == "[DONE]":
+                    break
+                data = json.loads(new_line)
+                chunk_choices = []
+                for choice in data["choices"]:
+                    chunk_delta = Delta()
+                    delta = choice.get("delta")
+                    if delta:
+                        if "role" in choice['delta']:
+                            chunk_delta.role = choice['delta']["role"]
+                        if "content" in choice['delta']:
+                            chunk_delta.content = choice['delta']["content"]
+                        chunk_choices.append(StreamingChoices(index=choice['index'], delta=chunk_delta))
+
+                if "usage" in data:
+                    chunk_usage = Usage()
+                    if "prompt_tokens" in data["usage"]:
+                        chunk_usage.prompt_tokens = data["usage"]["prompt_tokens"]
+                    if "completion_tokens" in data["usage"]:
+                        chunk_usage.completion_tokens = data["usage"]["completion_tokens"]
+                    if "total_tokens" in data["usage"]:
+                        chunk_usage.total_tokens = data["usage"]["total_tokens"]
+                
+                chunk_response = ModelResponse(
+                    id=data["id"],
+                    choices=chunk_choices,
+                    created=data["created"],
+                    model=model,
+                    usage=chunk_usage if "usage" in data else None,
+                    stream=True
+                )
+                yield chunk_response
+
+    def create_model_response_wrapper(self, result, model):
+        response_dict = result.json()
+        choices = []
+        for choice in response_dict["choices"]:
+            message = Message(
+                content=choice["message"]["content"], role=choice["message"]["role"]
+            )
+            choices.append(
+                Choices(
+                    message=message,
+                    index=choice["index"],
+                    finish_reason=choice["finish_reason"],
+                )
+            )
+        usage = Usage(
+            prompt_tokens=response_dict["usage"]["prompt_tokens"],
+            completion_tokens=response_dict["usage"]["completion_tokens"],
+            total_tokens=response_dict["usage"]["total_tokens"],
+        )
+        response = ModelResponse(
+            id=response_dict["id"],
+            choices=choices,
+            created=response_dict["created"],
+            model=model,
+            usage=usage,
+        )
+        return response
+
+
+    def completion(self, model: str, messages: list, **kwargs):
+        try:
+            if model is None or messages is None:
+                raise BaiChuanOpenAIError(
+                    status_code=422, message=f"Missing model or messages"
+                )
+            new_kwargs = self.pre_processing(**kwargs)
+            stream = kwargs.get("stream", False)
+
+            if stream:
+                return self.post_stream_processing_wrapper(model, messages, **new_kwargs)
+            else:
+                payload = json.dumps({"model": model, "messages": messages, **new_kwargs})
+                headers = {
+                    "Authorization": f"Bearer {self.api_key}",
+                    "Content-Type": "application/json",
+                }
+                result = requests.post(self.endpoint_url, headers=headers, data=payload)
+                return self.create_model_response_wrapper(result, model=model)
+        except Exception as e:
+            if hasattr(e, "status_code"):
+                raise BaiChuanOpenAIError(status_code=e.status_code, message=str(e))
+            else:
+                raise BaiChuanOpenAIError(status_code=500, message=str(e))

cnlitellm/providers/base_provider.py

@@ -0,0 +1,44 @@
+from abc import ABC, abstractmethod
+from ..models import ResponseModel
+from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, List
+from cnlitellm.utils import Message, Choices, Usage, ModelResponse
+
+if TYPE_CHECKING:
+    from dataclasses import dataclass
+else:
+    from pydantic.dataclasses import dataclass
+
+@dataclass
+class BaseProvider(ABC):
+    cls: str
+    args: Optional[Dict[str, Any]] = None
+    key: Optional[str] = None
+    group: Optional[str] = None
+
+    @abstractmethod
+    def completion(self, model: str, messages: list) -> ResponseModel:
+        pass
+
+    def create_model_response(self, openai_response: Any, model: str) -> ResponseModel:
+        print("openai_response: ", openai_response)
+        choices = []
+
+        for choice in openai_response.choices:
+            print("choice.message.content: ", choice.message.content)
+            message = Message(content=choice.message.content, role=choice.message.role)
+            print("message: ", message)
+            choices.append(Choices(message=message, index=choice.index, finish_reason=choice.finish_reason))
+
+        usage = Usage(
+            prompt_tokens=openai_response.usage.prompt_tokens,
+            completion_tokens=openai_response.usage.completion_tokens,
+            total_tokens=openai_response.usage.total_tokens
+        )
+        response =  ModelResponse(
+            id=openai_response.id,
+            choices=choices,
+            created=openai_response.created,
+            model=model,
+            usage=usage
+        )
+        pass

cnlitellm/providers/coze.py

@@ -0,0 +1,230 @@
+from .base_provider import BaseProvider
+from cnlitellm.utils import ResponseModelInterface
+from cnlitellm.utils import ModelResponse, Message, Choices, Usage, Context, generate_unique_uid, Delta, StreamingChoices
+from openai import OpenAI
+import logging, json, time, requests
+
+
+class CozeAIError(Exception):
+    def __init__(
+        self,
+        status_code,
+        message,
+    ):
+        self.status_code = status_code
+        self.message = message
+        super().__init__(self.message)
+
+
+class CozeAIProvider(BaseProvider):
+    def __init__(self, **model_kwargs):
+        self.api_key = model_kwargs.get("api_key")
+        self.base_url = "https://api.coze.com/open_api/v2"
+        self.endpoint_url = self.base_url + "/chat"
+
+    def pre_processing(self, **kwargs):
+        # 处理参数兼容性问题，不支持的参数全部舍弃
+        supported_params = [
+            "model", "messages", "max_tokens", "temperature", "top_p", "n",
+            "logprobs", "stream", "stop", "presence_penalty", "frequency_penalty",
+            "best_of", "logit_bias"
+        ]
+        for key in list(kwargs.keys()):
+            if key not in supported_params:
+                kwargs.pop(key)
+        return kwargs
+
+    def to_formatted_prompt(self, messages):
+        message = messages[-1]
+        # if last message role is not user, return error
+        if message["role"] != "user":
+            raise CozeAIError(
+                status_code=422, message=f"Last message role should be user"
+            )
+        query = message["content"]
+        return messages, query
+
+    def post_stream_processing_wrapper(self, model, messages, **new_kwargs):
+        # 预处理消息内容并提取用户问题
+        messages, query = self.to_formatted_prompt(messages)
+
+        # 接收user_id作为用户唯一标识传入参数
+        if 'user_id' not in new_kwargs:
+            user_id = generate_unique_uid()
+        else:
+            user_id = new_kwargs['user_id']
+        payload = json.dumps({"query": query,"user": user_id,"bot_id": model,"chat_history":[], "stream": True})
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+        response = requests.post(self.endpoint_url, headers=headers, data=payload)
+
+        for line in response.iter_lines():
+            if line:
+                if line.startswith(b"data:"):
+                    new_line = line.decode("utf-8").replace("data:", "")
+                    if new_line == "[DONE]":
+                        continue
+
+                    data = json.loads(new_line)
+                    print("data is:",data)
+                    chunk_line = {}
+                    chunk_context = []
+                    chunk_choices = []
+                    index = 0
+                    if "message" in data:
+                        message = data["message"]
+                        if message['role'] == 'assistant' and message['type'] == 'knowledge':
+                            # 解析知识型内容
+                            content = message['content']
+                            conversation_id = data.get('conversation_id', 0)
+                            # 假设knowledge类型内容是由"---\nrecall slice X:\n"分隔的
+                            slices = content.split('---\n')
+                            index = 0
+                            for slice in slices:
+                                if slice.strip().startswith('recall slice'):
+                                    # 去掉前缀找到JSON部分
+                                    json_part = slice.strip().split('\n', 1)[-1].strip()
+                                    # 如果不是以\"}结尾，则强制添加
+                                    if not json_part.endswith('\"}'):
+                                        json_part += '\"}'
+                                    try:
+                                        # 尝试解析JSON数据
+                                        recall_data = json.loads(json_part)
+                                        chunk_context.append({
+                                            "id": index,
+                                            "content": str(recall_data)    
+                                        })
+                                        index += 1
+                                    except json.JSONDecodeError:
+                                        raise CozeAIError(status_code=500, message=f"Error decoding JSON from slice: {json_part}")
+
+                        elif message['role'] == 'assistant' and message['type'] == 'answer':
+                            chunk_choices = []
+                            chunk_delta = Delta()
+                            if "role" in message:
+                                chunk_delta.role = message["role"]
+                            if "content" in message:
+                                chunk_delta.content = message["content"]
+                            chunk_choices.append(StreamingChoices(index=str(index), delta=chunk_delta))
+                        
+                        chunk_usage = Usage()
+                        chunk_response = ModelResponse(
+                            id=conversation_id,
+                            choices=chunk_choices,
+                            context=chunk_context,
+                            created=int(time.time()),
+                            model=model,
+                            usage=chunk_usage if chunk_usage else None,
+                            stream=True
+                        )
+                        index += 1
+                        yield chunk_response
+
+    def create_model_response_wrapper(self, result, model):
+        choices = []
+        context = []
+
+        result_dict = result.json()
+        messages = result_dict.get('messages', [])
+        results = []
+        cotext_pre = []
+
+        print("messages is:",messages)
+
+        for message in messages:
+            # 根据消息类型处理数据
+            if message['role'] == 'assistant' and message['type'] == 'knowledge':
+                # 解析知识型内容
+                content = message['content']
+                print("content is:",content)
+
+                # 假设knowledge类型内容是由"---\nrecall slice X:\n"分隔的
+                slices = content.split('---\n')
+                for slice in slices:
+                    if slice.strip().startswith('recall slice'):
+                        # 去掉前缀找到JSON部分
+                        json_part = slice.strip().split('\n', 1)[-1].strip()
+                        # 如果不是以\"}结尾，则强制添加
+                        if not json_part.endswith('\"}'):
+                            json_part += '\"}'
+
+                        try:
+                            # 尝试解析JSON数据
+                            recall_data = json.loads(json_part)
+                            cotext_pre.append(str(recall_data))
+                        except json.JSONDecodeError:
+                            raise CozeAIError(status_code=500, message=f"Error decoding JSON from slice: {json_part}")
+
+            elif message['role'] == 'assistant' and message['type'] == 'answer':
+                # 直接添加回复内容
+                results.append(message['content'])
+
+        message = Message(
+            content=results[0],
+            role="assistant"
+        )
+        choices.append(
+            Choices(
+                message=message,
+                index=0,
+                finish_reason="stop",
+            )
+        )
+
+        usage = Usage(
+            prompt_tokens=0,
+            completion_tokens=0,
+            total_tokens=0,
+        )
+
+        # Assume response_dict contains a field 'responseData' with quoteList to extract context
+        if cotext_pre:
+            i = 0
+            for context_item in cotext_pre:
+                i += 1
+                context.append(Context(id=i, content=context_item))
+
+        response = ModelResponse(
+            id=result_dict.get('conversation_id'),
+            choices=choices,
+            context=context,
+            created=int(time.time()),
+            model='',
+            usage=usage,
+        )
+        return response
+
+    def completion(self, model: str, messages: list, **kwargs):
+        try:
+            if model is None or messages is None:
+                raise CozeAIError(
+                    status_code=422, message="Missing model or messages"
+                )
+            new_kwargs = self.pre_processing(**kwargs)
+            stream = kwargs.get("stream", False)
+
+            if stream:
+                return self.post_stream_processing_wrapper(model, messages, **new_kwargs)
+            else:
+                messages, query = self.to_formatted_prompt(messages)
+
+                # 接受user_id作为用户唯一身份标识传入参数
+                if 'user_id' not in kwargs:
+                    user_id = generate_unique_uid()
+                else:
+                    user_id = kwargs['user_id']
+                payload = json.dumps({"query": query,"user": user_id,"bot_id": model,"chat_history":[],"steam":False})
+
+                headers = {
+                    "Authorization": f"Bearer {self.api_key}",
+                    "Content-Type": "application/json",
+                }
+                result = requests.post(self.endpoint_url, headers=headers, data=payload)
+                return self.create_model_response_wrapper(result, model=model)
+        except Exception as e:
+            if hasattr(e, "status_code"):
+                raise CozeAIError(status_code=e.status_code, message=str(e))
+            else:
+                raise CozeAIError(status_code=500, message=str(e))

cnlitellm/providers/dify.py

@@ -0,0 +1,192 @@
+import json
+from .base_provider import BaseProvider
+from cnlitellm.utils import ModelResponse, Message, Choices, Usage, Context, Delta, StreamingChoices
+import requests
+import logging, json, time
+
+class DifyOpenAIError(Exception):
+    def __init__(
+        self,
+        status_code,
+        message,
+    ):
+        self.status_code = status_code
+        self.message = message
+        super().__init__(self.message)
+
+class DifyAIProvider(BaseProvider):
+    def __init__(self, **model_kwargs):
+        self.api_key = model_kwargs.get("api_key")
+        self.endpoint_url = "https://api.dify.ai/v1/chat-messages"
+
+    def pre_processing(self, **kwargs):
+        supported_params = [
+            "model", "messages", "max_tokens", "temperature", "top_p", "n",
+            "logprobs", "stream", "stop", "presence_penalty", "frequency_penalty",
+            "best_of", "logit_bias"
+        ]
+        for key in list(kwargs.keys()):
+            if key not in supported_params:
+                kwargs.pop(key)
+        return kwargs
+    
+    def to_formatted_prompt(self, messages):
+        message = messages[-1]
+        # if last message role is not user, return error
+        if message["role"] != "user":
+            raise DifyOpenAIError(
+                status_code=422, message=f"Last message role should be user"
+            )
+        query = message["content"]
+        return messages, query
+
+
+    def post_stream_processing_wrapper(self, model, messages, **new_kwargs):
+        # 预处理对话内容并返回最近的用户问题
+        messages, query = self.to_formatted_prompt(messages)
+        mode = "streaming"
+        payload = json.dumps({"query": query, "response_mode": mode, "user": "abc-123","conversation_id": "","inputs":{}})
+
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+        response = requests.post(self.endpoint_url, headers=headers, data=payload)
+
+        chunk_usage = Usage()
+        index = 0
+        for line in response.iter_lines():
+            if line:
+                chunk_choices = []
+                chunk_context = []
+                chunk_delta = Delta()
+                # judge if the new_line begins with "data:"
+                if line.startswith(b"data:"):
+                    new_line = line.decode("utf-8").replace("data: ", "")
+                    if new_line == "[DONE]":
+                        continue
+                    
+                    
+                    data = json.loads(new_line)
+                    print("data is:",data)
+                    event = data.get("event")
+                    msg_id = data.get("message_id")
+                    if event == "message":
+                        if 'answer' in data:
+                            chunk_message = data["answer"]
+                            if chunk_message:
+                                chunk_delta.role = "assistant"
+                                chunk_delta.content = chunk_message
+                                chunk_choices.append(StreamingChoices(index=str(index), delta=chunk_delta))
+
+                    if event == "message_end" and "metadata" in data:
+                        metadata = data["metadata"]
+                        if "usage" in metadata:
+                            usage_info = metadata["usage"]
+                            if "prompt_tokens" in usage_info:
+                                chunk_usage.prompt_tokens = usage_info["prompt_tokens"]
+                            if "completion_tokens" in usage_info:
+                                chunk_usage.completion_tokens = usage_info["completion_tokens"]
+                            if "total_tokens" in usage_info:
+                                chunk_usage.total_tokens = usage_info["total_tokens"]
+
+                        if 'retriever_resources' in metadata:
+                            for resource in metadata['retriever_resources']:
+                                chunk_context.append({
+                                    "id": resource["position"],
+                                    "content": resource["content"],
+                                    "score": resource["score"],    
+                                })
+                    chunk_response = ModelResponse(
+                        id=msg_id,
+                        choices=chunk_choices,
+                        context=chunk_context,
+                        created=int(time.time()),
+                        model=model,
+                        usage=chunk_usage,
+                        stream=True
+                    )
+                    index += 1
+                    yield chunk_response
+
+    def create_model_response_wrapper(self, result, model):
+
+        print("result of dify:", result.json())
+
+        response_dict = result.json()
+        choices = []
+        context = []
+
+        message = Message(
+            content=response_dict['answer'], role="assistant"
+        )
+        choices.append(
+            Choices(
+                message=message,
+                index=0,
+                finish_reason="stop",
+            )
+        )
+
+        usage = Usage(
+            prompt_tokens=response_dict['metadata']["usage"]["prompt_tokens"],
+            completion_tokens=response_dict['metadata']["usage"]["completion_tokens"],
+            total_tokens=response_dict['metadata']["usage"]["total_tokens"],
+        )
+
+        if 'retriever_resources' in response_dict['metadata']:
+            retrieved_resources = response_dict['metadata']["retriever_resources"]
+            if retrieved_resources is not None:
+                for resource in retrieved_resources:
+                    context.append(
+                        Context(
+                            id=resource["position"],
+                            content=resource["content"],
+                            score=resource["score"],
+                        )
+                    )
+        else:
+            context = []
+
+        response = ModelResponse(
+            id=response_dict["id"],
+            choices=choices,
+            context=context,
+            created=int(time.time()),            
+            model=model,
+            usage=usage,
+        )
+        return response
+
+
+    def completion(self, model: str, messages: list, **kwargs):
+        try:
+            if model is None or messages is None:
+                raise DifyOpenAIError(
+                    status_code=422, message=f"Missing model or messages"
+                )
+            new_kwargs = self.pre_processing(**kwargs)
+            stream = kwargs.get("stream", False)
+
+            print("stream is:", stream)
+
+            if stream:
+                return self.post_stream_processing_wrapper(model, messages, **new_kwargs)
+            else:
+                messages, query = self.to_formatted_prompt(messages)
+                mode = "blocking"
+
+                print("mode is:", mode)
+
+                payload = json.dumps({"query": query, "response_mode": mode, "user": "abc-123","conversation_id": "","inputs":{}})
+                headers = {
+                    "Authorization": f"Bearer {self.api_key}",
+                    "Content-Type": "application/json",
+                }
+                result = requests.post(self.endpoint_url, headers=headers, data=payload)
+                return self.create_model_response_wrapper(result, model=model)
+        except Exception as e:
+            if hasattr(e, "status_code"):
+                raise DifyOpenAIError(status_code=e.status_code, message=str(e))
+            else:
+                raise DifyOpenAIError(status_code=500, message=str(e))