unionllm 0.2.0__py2.py3-none-any.whl

unionllm/__init__.py

@@ -0,0 +1 @@
+from .main import UnionLLM, unionchat

unionllm/exceptions.py

@@ -0,0 +1,146 @@
+class UnionLLMError(Exception):
+    """Base exception class for UnionLLM."""
+    pass
+
+class ProviderError(UnionLLMError):
+    """Exception raised when there is an issue with the model provider."""
+    pass
+
+class APICallError(UnionLLMError):
+    """Exception raised when an API call fails."""
+    pass
+
+from openai import (
+    AuthenticationError,
+    BadRequestError,
+    RateLimitError,
+    APIStatusError,
+    OpenAIError,
+    APIError, 
+    APITimeoutError, 
+    APIConnectionError, 
+)
+import httpx
+
+class AuthenticationError(AuthenticationError):  # type: ignore
+    def __init__(self, message, llm_provider, model, response: httpx.Response):
+        self.status_code = 401
+        self.message = message
+        self.llm_provider = llm_provider
+        self.model = model
+        super().__init__(
+            self.message,
+            response=response,
+            body=None
+        )  # Call the base class constructor with the parameters it needs
+
+class BadRequestError(BadRequestError):  # type: ignore
+    def __init__(self, message, model, llm_provider, response: httpx.Response):
+        self.status_code = 400
+        self.message = message
+        self.model = model
+        self.llm_provider = llm_provider
+        super().__init__(
+            self.message,
+            response=response,
+            body=None
+        )  # Call the base class constructor with the parameters it needs
+
+class Timeout(APITimeoutError):  # type: ignore
+    def __init__(self, message, model, llm_provider, request: httpx.Request):
+        self.status_code = 408
+        self.message = message
+        self.model = model
+        self.llm_provider = llm_provider
+        super().__init__(
+            request=request
+        )  # Call the base class constructor with the parameters it needs
+
+class RateLimitError(RateLimitError):  # type: ignore
+    def __init__(self, message, llm_provider, model, response: httpx.Response):
+        self.status_code = 429
+        self.message = message
+        self.llm_provider = llm_provider
+        self.modle = model
+        super().__init__(
+            self.message,
+            response=response,
+            body=None
+        )  # Call the base class constructor with the parameters it needs
+
+# sub class of rate limit error - meant to give more granularity for error handling context window exceeded errors
+class ContextWindowExceededError(BadRequestError):  # type: ignore
+    def __init__(self, message, model, llm_provider, response: httpx.Response):
+        self.status_code = 400
+        self.message = message
+        self.model = model
+        self.llm_provider = llm_provider
+        super().__init__(
+            message=self.message, 
+            model=self.model, # type: ignore
+            llm_provider=self.llm_provider, # type: ignore
+            response=response
+        )  # Call the base class constructor with the parameters it needs
+
+class ServiceUnavailableError(APIStatusError):  # type: ignore
+    def __init__(self, message, llm_provider, model):
+        self.status_code = 503
+        self.message = message
+        self.llm_provider = llm_provider
+        self.model = model
+        super().__init__(
+            self.message
+        )  # Call the base class constructor with the parameters it needs
+
+
+# raise this when the API returns an invalid response object - https://github.com/openai/openai-python/blob/1be14ee34a0f8e42d3f9aa5451aa4cb161f1781f/openai/api_requestor.py#L401
+class APIError(APIError): # type: ignore 
+    def __init__(self, status_code, message, llm_provider, model, request: httpx.Request):
+        self.status_code = status_code 
+        self.message = message
+        self.llm_provider = llm_provider
+        self.model = model
+        super().__init__(
+            self.message,
+            request=request # type: ignore
+        )
+
+# raised if an invalid request (not get, delete, put, post) is made
+class APIConnectionError(APIConnectionError):  # type: ignore 
+    def __init__(self, message, llm_provider, model):
+        self.message = message
+        self.llm_provider = llm_provider
+        self.model = model
+        super().__init__(
+            self.message
+        )
+
+class OpenAIError(OpenAIError):  # type: ignore
+    def __init__(self, original_exception):
+        self.status_code = original_exception.http_status
+        super().__init__(
+            http_body=original_exception.http_body,
+            http_status=original_exception.http_status,
+            json_body=original_exception.json_body,
+            headers=original_exception.headers,
+            code=original_exception.code,
+        )
+        self.llm_provider = "openai"
+
+class BudgetExceededError(Exception):
+    def __init__(self, current_cost, max_budget):
+        self.current_cost = current_cost
+        self.max_budget = max_budget
+        message = f"Budget has been exceeded! Current cost: {current_cost}, Max budget: {max_budget}"
+        super().__init__(message)
+
+## DEPRECATED ## 
+class InvalidRequestError(BadRequestError):  # type: ignore
+    def __init__(self, message, model, llm_provider):
+        self.status_code = 400
+        self.message = message
+        self.model = model
+        self.llm_provider = llm_provider
+        super().__init__(
+            self.message, f"{self.model}"
+        )  # Call the base class constructor with the parameters it needs

unionllm/main.py

@@ -0,0 +1,142 @@
+import logging
+import asyncio
+import os
+from functools import partial
+
+from typing import Any, List, Optional
+from .providers import zhipu, moonshot, xai, minimax, qwen, tiangong, baichuan, wenxin, xunfei, xunfei_http, dify, fastgpt, coze, litellm, lingyi, stepfun, doubao, deepseek, gemini
+from .exceptions import ProviderError
+# from litellm import completion as litellm_completion
+
+logger = logging.getLogger(__name__)
+
+class UnionLLM:
+    def __init__(self, provider: Optional[str] = None, **kwargs):
+        self.provider = provider.lower() if provider else None
+        self.litellm_call_type = None
+
+        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":
+            # 根据model判断是否需要使用dashscope的api, 否则使用openai的兼容api
+            model = kwargs.get("model")
+            # 预留特殊模型通过litellm调用
+            if model is not None and model not in ["qwen-special-model"]:
+                self.provider_instance = qwen.QwenAIProvider(**kwargs)
+            else:
+                # 使用DASHSCOPE_API_KEY作为OPENAI_API_KEY
+                os.environ["OPENAI_API_KEY"] = os.environ.get("DASHSCOPE_API_KEY")
+                self.provider_instance = litellm.LiteLLMProvider(**kwargs)
+                self.litellm_call_type = 3
+        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 == "xunfei_http":
+            self.provider_instance = xunfei_http.XunfeiHTTPProvider(**kwargs)
+        # elif self.provider == "xai":
+        #     self.provider_instance = xai.XAIHTTPProvider(**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)
+        elif self.provider == "lingyi":
+            self.provider_instance = lingyi.LingyiAIProvider(**kwargs)
+        elif self.provider == "stepfun":
+            self.provider_instance = stepfun.StepfunAIProvider(**kwargs)
+        elif self.provider == "doubao":
+            self.provider_instance = doubao.DouBaoAIProvider(**kwargs)
+        elif self.provider == "deepseek":
+            self.provider_instance = deepseek.DeepSeekAIProvider(**kwargs)
+        # elif self.provider == "gemini":
+        #     self.provider_instance = gemini.GeminiAIProvider(**kwargs)
+        elif self.provider == "gemini" and kwargs.get("multimodal") and kwargs['multimodal'] == True:
+            self.provider_instance = gemini.GeminiAIProvider(**kwargs)
+        elif self.provider == "gemini" and kwargs.get("fileinput") and kwargs['fileinput'] == True:
+            self.provider_instance = gemini.GeminiAIProvider(**kwargs)
+        elif self.provider == "gemini" and kwargs.get("videoinput") and kwargs['videoinput'] == True:
+            self.provider_instance = gemini.GeminiAIProvider(**kwargs)
+
+        elif self.provider:
+            if_litellm_support, support_type = self.check_litellm_providers(provider=self.provider)
+            if if_litellm_support:
+                self.provider_instance = litellm.LiteLLMProvider(**kwargs)
+                if support_type == 1:
+                    self.litellm_call_type = 1
+                elif support_type == 2:
+                    self.litellm_call_type = 2
+                elif support_type == 3:
+                    self.litellm_call_type = 3
+            else:
+                raise ProviderError(f"Provider '{self.provider}' is not supported.")
+        else:
+            self.provider_instance = litellm.LiteLLMProvider(**kwargs)
+
+    def completion(self, model: str, messages: List[str], **kwargs) -> Any:
+        if not self.provider_instance:
+            raise ProviderError(f"Provider '{self.provider}' is not initialized.")
+        if self.litellm_call_type:
+            if self.litellm_call_type == 1:
+                # Jugde whether the model starts with self.provider, if not, add it
+                if not model.startswith(self.provider+"/"):
+                    model = f"{self.provider}/{model}"     
+                return self.provider_instance.completion(model, messages, **kwargs)
+            elif self.litellm_call_type == 2:
+                return self.provider_instance.completion(model, messages, **kwargs)
+            elif self.litellm_call_type == 3:
+                # append OpenAI as provider name
+                if self.provider == 'xai':
+                    kwargs['api_base'] = "https://api.x.ai/v1"
+                if self.provider == 'qwen':
+                    kwargs['api_base'] = "https://dashscope.aliyuncs.com/compatible-mode/v1"
+                model = f"openai/{model}"
+                return self.provider_instance.completion(model, messages, **kwargs)
+        else:
+            return self.provider_instance.completion(model, messages, **kwargs)
+        
+    async def acompletion(self, model: str, messages: List[str], **kwargs) -> Any:
+        loop = asyncio.get_event_loop()
+        if not self.provider_instance:
+            raise ProviderError(f"Provider '{self.provider}' is not initialized.")
+        
+        if self.litellm_call_type:
+            if self.litellm_call_type == 1:
+                if not model.startswith(self.provider + "/"):
+                    model = f"{self.provider}/{model}"
+                func = partial(self.provider_instance.completion, model, messages, **kwargs)
+                return await loop.run_in_executor(None, func)
+            elif self.litellm_call_type == 2:
+                func = partial(self.provider_instance.completion, model, messages, **kwargs)
+                return await loop.run_in_executor(None, func)
+        else:
+            func = partial(self.provider_instance.completion, model, messages, **kwargs)
+            return await loop.run_in_executor(None, func)
+
+    def check_litellm_providers(self, provider: str) -> bool:
+        # Judge whether the provider is supported by LiteLLM, and if provider name should be added to the model name
+        if provider in ['azure', 'azure_ai', 'anthropic', 'deepseek', 'sagemaker', 'bedrock', 'vertex_ai', 'vertex_ai_beta', 'palm', 'gemini', 'mistral', 'cloudflare', 'huggingface', 'replicate', 'together_ai', 'openrouter', 'baseten', 'nlp_cloud', 'petals', 'ollama', 'perplexity', 'groq', 'anyscale', 'watsonx', 'voyage', 'xinference']:
+            # provider name should be added to the model name
+            return True, 1
+        elif provider in ['openai', 'cohere', 'ai21', 'deepinfra', 'ai21', 'alpha_alpha']:
+            # provider name should not be added to the model name
+            return True, 2
+        elif provider in ['xai', 'qwen']:
+            # append OpenAI as procall as OpenAI campatible API
+            return True, 3
+        else:
+            return False, 0
+        
+# This is the new function to simplify calling
+def unionchat(model: str, messages: List[dict], **kwargs) -> Any:
+    client = UnionLLM(**kwargs)
+    return client.completion(model=model, messages=messages, **kwargs)

unionllm/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

unionllm/providers/baichuan.py

@@ -0,0 +1,143 @@
+import json
+from .base_provider import BaseProvider
+from unionllm.utils import ModelResponse, Message, Choices, Usage, Delta, StreamingChoices
+import requests
+import logging, os
+
+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):        
+        # Get BAICHUAN_API_KEY from environment variables
+        _env_api_key = os.environ.get("BAICHUAN_API_KEY")
+        self.api_key = model_kwargs.get("api_key") if model_kwargs.get("api_key") else _env_api_key
+        if not self.api_key:
+            raise BaiChuanOpenAIError(
+                status_code=422, message=f"Missing 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", "tools", "tool_choice"
+        ]
+        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"
+                )
+
+            message_check_result = self.check_prompt("baichuan", model, messages)    
+            if message_check_result['pass_check']:
+                messages = message_check_result['messages']
+            else:
+                raise BaiChuanOpenAIError(
+                    status_code=422, message=message_check_result['reason']
+                )
+                
+            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))