aient 1.0.29__py3-none-any.whl

aient/core/utils.py

@@ -0,0 +1,655 @@
+import re
+import io
+import os
+import json
+import httpx
+import base64
+import asyncio
+import urllib.parse
+from time import time
+from PIL import Image
+from fastapi import HTTPException
+from urllib.parse import urlparse
+from collections import defaultdict
+
+from .log_config import logger
+
+def get_model_dict(provider):
+    model_dict = {}
+    for model in provider['model']:
+        if type(model) == str:
+            model_dict[model] = model
+        if isinstance(model, dict):
+            model_dict.update({new: old for old, new in model.items()})
+    return model_dict
+
+class BaseAPI:
+    def __init__(
+        self,
+        api_url: str = "https://api.openai.com/v1/chat/completions",
+    ):
+        if api_url == "":
+            api_url = "https://api.openai.com/v1/chat/completions"
+        self.source_api_url: str = api_url
+        from urllib.parse import urlparse, urlunparse
+        parsed_url = urlparse(self.source_api_url)
+        # print("parsed_url", parsed_url)
+        if parsed_url.scheme == "":
+            raise Exception("Error: API_URL is not set")
+        if parsed_url.path != '/':
+            before_v1 = parsed_url.path.split("chat/completions")[0]
+            if not before_v1.endswith("/"):
+                before_v1 = before_v1 + "/"
+        else:
+            before_v1 = ""
+        self.base_url: str = urlunparse(parsed_url[:2] + ("",) + ("",) * 3)
+        self.v1_url: str = urlunparse(parsed_url[:2]+ (before_v1,) + ("",) * 3)
+        self.v1_models: str = urlunparse(parsed_url[:2] + (before_v1 + "models",) + ("",) * 3)
+        self.chat_url: str = urlunparse(parsed_url[:2] + (before_v1 + "chat/completions",) + ("",) * 3)
+        self.image_url: str = urlunparse(parsed_url[:2] + (before_v1 + "images/generations",) + ("",) * 3)
+        self.audio_transcriptions: str = urlunparse(parsed_url[:2] + (before_v1 + "audio/transcriptions",) + ("",) * 3)
+        self.moderations: str = urlunparse(parsed_url[:2] + (before_v1 + "moderations",) + ("",) * 3)
+        self.embeddings: str = urlunparse(parsed_url[:2] + (before_v1 + "embeddings",) + ("",) * 3)
+        self.audio_speech: str = urlunparse(parsed_url[:2] + (before_v1 + "audio/speech",) + ("",) * 3)
+
+        if parsed_url.hostname == "generativelanguage.googleapis.com":
+            self.base_url = api_url
+            self.v1_url = api_url
+            self.chat_url = api_url
+            self.embeddings = api_url
+
+def get_engine(provider, endpoint=None, original_model=""):
+    parsed_url = urlparse(provider['base_url'])
+    # print("parsed_url", parsed_url)
+    engine = None
+    stream = None
+    if parsed_url.path.endswith("/v1beta") or parsed_url.path.endswith("/v1") or parsed_url.netloc == 'generativelanguage.googleapis.com':
+        engine = "gemini"
+    elif parsed_url.netloc.rstrip('/').endswith('aiplatform.googleapis.com'):
+        engine = "vertex"
+    elif parsed_url.netloc.rstrip('/').endswith('openai.azure.com') or parsed_url.netloc.rstrip('/').endswith('services.ai.azure.com'):
+        engine = "azure"
+    elif parsed_url.netloc == 'api.cloudflare.com':
+        engine = "cloudflare"
+    elif parsed_url.netloc == 'api.anthropic.com' or parsed_url.path.endswith("v1/messages"):
+        engine = "claude"
+    elif parsed_url.netloc == 'api.cohere.com':
+        engine = "cohere"
+        stream = True
+    else:
+        engine = "gpt"
+
+    original_model = original_model.lower()
+    if original_model \
+    and "claude" not in original_model \
+    and "gpt" not in original_model \
+    and "deepseek" not in original_model \
+    and "o1" not in original_model \
+    and "o3" not in original_model \
+    and "gemini" not in original_model \
+    and "learnlm" not in original_model \
+    and "grok" not in original_model \
+    and parsed_url.netloc != 'api.cloudflare.com' \
+    and parsed_url.netloc != 'api.cohere.com':
+        engine = "openrouter"
+
+    if "claude" in original_model and engine == "vertex":
+        engine = "vertex-claude"
+
+    if "gemini" in original_model and engine == "vertex":
+        engine = "vertex-gemini"
+
+    if provider.get("engine"):
+        engine = provider["engine"]
+
+    if endpoint == "/v1/images/generations" or "stable-diffusion" in original_model:
+        engine = "dalle"
+        stream = False
+
+    if endpoint == "/v1/audio/transcriptions":
+        engine = "whisper"
+        stream = False
+
+    if endpoint == "/v1/moderations":
+        engine = "moderation"
+        stream = False
+
+    if endpoint == "/v1/embeddings":
+        engine = "embedding"
+
+    if endpoint == "/v1/audio/speech":
+        engine = "tts"
+        stream = False
+
+    return engine, stream
+
+def get_proxy(proxy, client_config = {}):
+    if proxy:
+        # 解析代理URL
+        parsed = urlparse(proxy)
+        scheme = parsed.scheme.rstrip('h')
+
+        if scheme == 'socks5':
+            try:
+                from httpx_socks import AsyncProxyTransport
+                proxy = proxy.replace('socks5h://', 'socks5://')
+                transport = AsyncProxyTransport.from_url(proxy)
+                client_config["transport"] = transport
+                # print("proxy", proxy)
+            except ImportError:
+                logger.error("httpx-socks package is required for SOCKS proxy support")
+                raise ImportError("Please install httpx-socks package for SOCKS proxy support: pip install httpx-socks")
+        else:
+            client_config["proxies"] = {
+                "http://": proxy,
+                "https://": proxy
+            }
+    return client_config
+
+def update_initial_model(provider):
+    try:
+        engine, stream_mode = get_engine(provider, endpoint=None, original_model="")
+        # print("engine", engine, provider)
+        api_url = provider['base_url']
+        api = provider['api']
+        proxy = safe_get(provider, "preferences", "proxy", default=None)
+        client_config = get_proxy(proxy)
+        if engine == "gemini":
+            url = "https://generativelanguage.googleapis.com/v1beta/models"
+            params = {"key": api}
+            with httpx.Client(**client_config) as client:
+                response = client.get(url, params=params)
+
+            original_models = response.json()
+            if original_models.get("error"):
+                raise Exception({"error": original_models.get("error"), "endpoint": url, "api": api})
+
+            models = {"data": []}
+            for model in original_models["models"]:
+                models["data"].append({
+                    "id": model["name"].split("models/")[-1],
+                })
+        else:
+            endpoint = BaseAPI(api_url=api_url)
+            endpoint_models_url = endpoint.v1_models
+            if isinstance(api, list):
+                api = api[0]
+            headers = {"Authorization": f"Bearer {api}"}
+            response = httpx.get(
+                endpoint_models_url,
+                headers=headers,
+                **client_config
+            )
+            models = response.json()
+            if models.get("error"):
+                logger.error({"error": models.get("error"), "endpoint": endpoint_models_url, "api": api})
+                return []
+
+        # print(models)
+        models_list = models["data"]
+        models_id = [model["id"] for model in models_list]
+        set_models = set()
+        for model_item in models_id:
+            set_models.add(model_item)
+        models_id = list(set_models)
+        # print(models_id)
+        return models_id
+    except Exception as e:
+        # print("error:", e)
+        import traceback
+        traceback.print_exc()
+        return []
+
+def safe_get(data, *keys, default=None):
+    for key in keys:
+        try:
+            data = data[key] if isinstance(data, (dict, list)) else data.get(key)
+        except (KeyError, IndexError, AttributeError, TypeError):
+            return default
+    if not data:
+        return default
+    return data
+
+def parse_rate_limit(limit_string):
+    # 定义时间单位到秒的映射
+    time_units = {
+        's': 1, 'sec': 1, 'second': 1,
+        'm': 60, 'min': 60, 'minute': 60,
+        'h': 3600, 'hr': 3600, 'hour': 3600,
+        'd': 86400, 'day': 86400,
+        'mo': 2592000, 'month': 2592000,
+        'y': 31536000, 'year': 31536000
+    }
+
+    # 处理多个限制条件
+    limits = []
+    for limit in limit_string.split(','):
+        limit = limit.strip()
+        # 使用正则表达式匹配数字和单位
+        match = re.match(r'^(\d+)/(\w+)$', limit)
+        if not match:
+            raise ValueError(f"Invalid rate limit format: {limit}")
+
+        count, unit = match.groups()
+        count = int(count)
+
+        # 转换单位到秒
+        if unit not in time_units:
+            raise ValueError(f"Unknown time unit: {unit}")
+
+        seconds = time_units[unit]
+        limits.append((count, seconds))
+
+    return limits
+
+class ThreadSafeCircularList:
+    def __init__(self, items = [], rate_limit={"default": "999999/min"}, schedule_algorithm="round_robin"):
+        if schedule_algorithm == "random":
+            import random
+            self.items = random.sample(items, len(items))
+            self.schedule_algorithm = "random"
+        elif schedule_algorithm == "round_robin":
+            self.items = items
+            self.schedule_algorithm = "round_robin"
+        elif schedule_algorithm == "fixed_priority":
+            self.items = items
+            self.schedule_algorithm = "fixed_priority"
+        else:
+            self.items = items
+            logger.warning(f"Unknown schedule algorithm: {schedule_algorithm}, use (round_robin, random, fixed_priority) instead")
+            self.schedule_algorithm = "round_robin"
+        self.index = 0
+        self.lock = asyncio.Lock()
+        # 修改为二级字典，第一级是item，第二级是model
+        self.requests = defaultdict(lambda: defaultdict(list))
+        self.cooling_until = defaultdict(float)
+        self.rate_limits = {}
+        if isinstance(rate_limit, dict):
+            for rate_limit_model, rate_limit_value in rate_limit.items():
+                self.rate_limits[rate_limit_model] = parse_rate_limit(rate_limit_value)
+        elif isinstance(rate_limit, str):
+            self.rate_limits["default"] = parse_rate_limit(rate_limit)
+        else:
+            logger.error(f"Error ThreadSafeCircularList: Unknown rate_limit type: {type(rate_limit)}, rate_limit: {rate_limit}")
+
+    async def set_cooling(self, item: str, cooling_time: int = 60):
+        """设置某个 item 进入冷却状态
+
+        Args:
+            item: 需要冷却的 item
+            cooling_time: 冷却时间(秒)，默认60秒
+        """
+        if item == None:
+            return
+        now = time()
+        async with self.lock:
+            self.cooling_until[item] = now + cooling_time
+            # 清空该 item 的请求记录
+            # self.requests[item] = []
+            logger.warning(f"API key {item} 已进入冷却状态，冷却时间 {cooling_time} 秒")
+
+    async def is_rate_limited(self, item, model: str = None) -> bool:
+        now = time()
+        # 检查是否在冷却中
+        if now < self.cooling_until[item]:
+            return True
+
+        # 获取适用的速率限制
+
+        if model:
+            model_key = model
+        else:
+            model_key = "default"
+
+        rate_limit = None
+        # 先尝试精确匹配
+        if model and model in self.rate_limits:
+            rate_limit = self.rate_limits[model]
+        else:
+            # 如果没有精确匹配，尝试模糊匹配
+            for limit_model in self.rate_limits:
+                if limit_model != "default" and model and limit_model in model:
+                    rate_limit = self.rate_limits[limit_model]
+                    break
+
+        # 如果都没匹配到，使用默认值
+        if rate_limit is None:
+            rate_limit = self.rate_limits.get("default", [(999999, 60)])  # 默认限制
+
+        # 检查所有速率限制条件
+        for limit_count, limit_period in rate_limit:
+            # 使用特定模型的请求记录进行计算
+            recent_requests = sum(1 for req in self.requests[item][model_key] if req > now - limit_period)
+            if recent_requests >= limit_count:
+                logger.warning(f"API key {item} 对模型 {model_key} 已达到速率限制 ({limit_count}/{limit_period}秒)")
+                return True
+
+        # 清理太旧的请求记录
+        max_period = max(period for _, period in rate_limit)
+        self.requests[item][model_key] = [req for req in self.requests[item][model_key] if req > now - max_period]
+
+        # 记录新的请求
+        self.requests[item][model_key].append(now)
+        return False
+
+    async def next(self, model: str = None):
+        async with self.lock:
+            if self.schedule_algorithm == "fixed_priority":
+                self.index = 0
+            start_index = self.index
+            while True:
+                item = self.items[self.index]
+                self.index = (self.index + 1) % len(self.items)
+
+                if not await self.is_rate_limited(item, model):
+                    return item
+
+                # 如果已经检查了所有的 API key 都被限制
+                if self.index == start_index:
+                    logger.warning(f"All API keys are rate limited!")
+                    raise HTTPException(status_code=429, detail="Too many requests")
+
+    async def after_next_current(self):
+        # 返回当前取出的 API，因为已经调用了 next，所以当前API应该是上一个
+        if len(self.items) == 0:
+            return None
+        async with self.lock:
+            item = self.items[(self.index - 1) % len(self.items)]
+            return item
+
+    def get_items_count(self) -> int:
+        """返回列表中的项目数量
+
+        Returns:
+            int: items列表的长度
+        """
+        return len(self.items)
+
+def circular_list_encoder(obj):
+    if isinstance(obj, ThreadSafeCircularList):
+        return obj.to_dict()
+    raise TypeError(f'Object of type {obj.__class__.__name__} is not JSON serializable')
+
+provider_api_circular_list = defaultdict(ThreadSafeCircularList)
+
+# 【GCP-Vertex AI 目前有這些區域可用】 https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude?hl=zh_cn
+# https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations?hl=zh-cn#available-regions
+
+# c3.5s
+# us-east5
+# europe-west1
+
+# c3s
+# us-east5
+# us-central1
+# asia-southeast1
+
+# c3o
+# us-east5
+
+# c3h
+# us-east5
+# us-central1
+# europe-west1
+# europe-west4
+
+
+c35s = ThreadSafeCircularList(["us-east5", "europe-west1"])
+c3s = ThreadSafeCircularList(["us-east5", "us-central1", "asia-southeast1"])
+c3o = ThreadSafeCircularList(["us-east5"])
+c3h = ThreadSafeCircularList(["us-east5", "us-central1", "europe-west1", "europe-west4"])
+gemini1 = ThreadSafeCircularList(["us-central1", "us-east4", "us-west1", "us-west4", "europe-west1", "europe-west2"])
+gemini2 = ThreadSafeCircularList(["us-central1"])
+
+
+
+# end_of_line = "\n\r\n"
+# end_of_line = "\r\n"
+# end_of_line = "\n\r"
+end_of_line = "\n\n"
+# end_of_line = "\r"
+# end_of_line = "\n"
+
+import random
+import string
+async def generate_sse_response(timestamp, model, content=None, tools_id=None, function_call_name=None, function_call_content=None, role=None, total_tokens=0, prompt_tokens=0, completion_tokens=0, reasoning_content=None):
+    random.seed(timestamp)
+    random_str = ''.join(random.choices(string.ascii_letters + string.digits, k=29))
+
+    delta_content = {"role": "assistant", "content": content} if content else {}
+    if reasoning_content:
+        delta_content = {"role": "assistant", "content": "", "reasoning_content": reasoning_content}
+
+    sample_data = {
+        "id": f"chatcmpl-{random_str}",
+        "object": "chat.completion.chunk",
+        "created": timestamp,
+        "model": model,
+        "choices": [
+            {
+                "index": 0,
+                "delta": delta_content,
+                "logprobs": None,
+                "finish_reason": None if content else "stop"
+            }
+        ],
+        "usage": None,
+        "system_fingerprint": "fp_d576307f90",
+    }
+    if function_call_content:
+        sample_data["choices"][0]["delta"] = {"tool_calls":[{"index":0,"function":{"arguments": function_call_content}}]}
+    if tools_id and function_call_name:
+        sample_data["choices"][0]["delta"] = {"tool_calls":[{"index":0,"id": tools_id,"type":"function","function":{"name": function_call_name, "arguments":""}}]}
+        # sample_data["choices"][0]["delta"] = {"tool_calls":[{"index":0,"function":{"id": tools_id, "name": function_call_name}}]}
+    if role:
+        sample_data["choices"][0]["delta"] = {"role": role, "content": ""}
+    if total_tokens:
+        total_tokens = prompt_tokens + completion_tokens
+        sample_data["usage"] = {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": total_tokens}
+        sample_data["choices"] = []
+    json_data = json.dumps(sample_data, ensure_ascii=False)
+
+    # 构建SSE响应
+    sse_response = f"data: {json_data}" + end_of_line
+
+    return sse_response
+
+async def generate_no_stream_response(timestamp, model, content=None, tools_id=None, function_call_name=None, function_call_content=None, role=None, total_tokens=0, prompt_tokens=0, completion_tokens=0):
+    random.seed(timestamp)
+    random_str = ''.join(random.choices(string.ascii_letters + string.digits, k=29))
+    sample_data = {
+        "id": f"chatcmpl-{random_str}",
+        "object": "chat.completion",
+        "created": timestamp,
+        "model": model,
+        "choices": [
+            {
+                "index": 0,
+                "message": {
+                    "role": role,
+                    "content": content,
+                    "refusal": None
+                },
+                "logprobs": None,
+                "finish_reason": "stop"
+            }
+        ],
+        "usage": None,
+        "system_fingerprint": "fp_a7d06e42a7"
+    }
+
+    if function_call_name:
+        if not tools_id:
+            tools_id = f"call_{random_str}"
+        sample_data = {
+            "id": f"chatcmpl-{random_str}",
+            "object": "chat.completion",
+            "created": timestamp,
+            "model": model,
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {
+                    "role": "assistant",
+                    "content": None,
+                    "tool_calls": [
+                        {
+                            "id": tools_id,
+                            "type": "function",
+                            "function": {
+                                "name": function_call_name,
+                                "arguments": json.dumps(function_call_content, ensure_ascii=False)
+                            }
+                        }
+                    ],
+                    "refusal": None
+                    },
+                    "logprobs": None,
+                    "finish_reason": "tool_calls"
+                }
+            ],
+            "usage": None,
+            "service_tier": "default",
+            "system_fingerprint": "fp_4691090a87"
+        }
+
+    if total_tokens:
+        total_tokens = prompt_tokens + completion_tokens
+        sample_data["usage"] = {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": total_tokens}
+
+    json_data = json.dumps(sample_data, ensure_ascii=False)
+
+    return json_data
+
+def get_image_format(file_content):
+    try:
+        img = Image.open(io.BytesIO(file_content))
+        return img.format.lower()
+    except:
+        return None
+
+def encode_image(image_path):
+    with open(image_path, "rb") as image_file:
+        file_content = image_file.read()
+        img_format = get_image_format(file_content)
+        if not img_format:
+            raise ValueError("无法识别的图片格式")
+        base64_encoded = base64.b64encode(file_content).decode('utf-8')
+
+        if img_format == 'png':
+            return f"data:image/png;base64,{base64_encoded}"
+        elif img_format in ['jpg', 'jpeg']:
+            return f"data:image/jpeg;base64,{base64_encoded}"
+        else:
+            raise ValueError(f"不支持的图片格式: {img_format}")
+
+async def get_doc_from_url(url):
+    filename = urllib.parse.unquote(url.split("/")[-1])
+    transport = httpx.AsyncHTTPTransport(
+        http2=True,
+        verify=False,
+        retries=1
+    )
+    async with httpx.AsyncClient(transport=transport) as client:
+        try:
+            response = await client.get(
+                url,
+                timeout=30.0
+            )
+            with open(filename, 'wb') as f:
+                f.write(response.content)
+
+        except httpx.RequestError as e:
+            print(f"An error occurred while requesting {e.request.url!r}.")
+
+    return filename
+
+async def get_encode_image(image_url):
+    filename = await get_doc_from_url(image_url)
+    image_path = os.getcwd() + "/" + filename
+    base64_image = encode_image(image_path)
+    os.remove(image_path)
+    return base64_image
+
+# from PIL import Image
+# import io
+# def validate_image(image_data, image_type):
+#     try:
+#         decoded_image = base64.b64decode(image_data)
+#         image = Image.open(io.BytesIO(decoded_image))
+
+#         # 检查图片格式是否与声明的类型匹配
+#         # print("image.format", image.format)
+#         if image_type == "image/png" and image.format != "PNG":
+#             raise ValueError("Image is not a valid PNG")
+#         elif image_type == "image/jpeg" and image.format not in ["JPEG", "JPG"]:
+#             raise ValueError("Image is not a valid JPEG")
+
+#         # 如果没有异常,则图片有效
+#         return True
+#     except Exception as e:
+#         print(f"Image validation failed: {str(e)}")
+#         return False
+
+async def get_image_message(base64_image, engine = None):
+    if base64_image.startswith("http"):
+        base64_image = await get_encode_image(base64_image)
+    colon_index = base64_image.index(":")
+    semicolon_index = base64_image.index(";")
+    image_type = base64_image[colon_index + 1:semicolon_index]
+
+    if image_type == "image/webp":
+        # 将webp转换为png
+
+        # 解码base64获取图片数据
+        image_data = base64.b64decode(base64_image.split(",")[1])
+
+        # 使用PIL打开webp图片
+        image = Image.open(io.BytesIO(image_data))
+
+        # 转换为PNG格式
+        png_buffer = io.BytesIO()
+        image.save(png_buffer, format="PNG")
+        png_base64 = base64.b64encode(png_buffer.getvalue()).decode('utf-8')
+
+        # 返回PNG格式的base64
+        base64_image = f"data:image/png;base64,{png_base64}"
+        image_type = "image/png"
+
+    if "gpt" == engine or "openrouter" == engine or "azure" == engine:
+        return {
+            "type": "image_url",
+            "image_url": {
+                "url": base64_image,
+            }
+        }
+    if "claude" == engine or "vertex-claude" == engine:
+        # if not validate_image(base64_image.split(",")[1], image_type):
+        #     raise ValueError(f"Invalid image format. Expected {image_type}")
+        return {
+            "type": "image",
+            "source": {
+                "type": "base64",
+                "media_type": image_type,
+                "data": base64_image.split(",")[1],
+            }
+        }
+    if "gemini" == engine or "vertex-gemini" == engine:
+        return {
+            "inlineData": {
+                "mimeType": image_type,
+                "data": base64_image.split(",")[1],
+            }
+        }
+    raise ValueError("Unknown engine")
+
+async def get_text_message(message, engine = None):
+    if "gpt" == engine or "claude" == engine or "openrouter" == engine or "vertex-claude" == engine or "azure" == engine:
+        return {"type": "text", "text": message}
+    if "gemini" == engine or "vertex-gemini" == engine:
+        return {"text": message}
+    if engine == "cloudflare":
+        return message
+    if engine == "cohere":
+        return message
+    raise ValueError("Unknown engine")

aient/models/__init__.py

@@ -0,0 +1,9 @@
+from .chatgpt import *
+from .claude import *
+from .gemini import *
+from .vertex import *
+from .groq import *
+from .audio import *
+from .duckduckgo import *
+
+# __all__ = ["chatgpt", "claude", "claude3", "gemini", "groq"]