aient 1.2.49__tar.gz → 1.2.50__tar.gz

{aient-1.2.49 → aient-1.2.50}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aient
-Version: 1.2.49
+Version: 1.2.50
 Summary: Aient: The Awakening of Agent.
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown

{aient-1.2.49 → aient-1.2.50}/aient/core/request.py

@@ -1,13 +1,24 @@
 import re
 import json
 import copy
+import hmac
+import time
 import httpx
 import base64
 import asyncio
+import hashlib
+import datetime
+import urllib.parse
 from io import IOBase
 from typing import Tuple
+from datetime import timezone
 from urllib.parse import urlparse
 
+
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.asymmetric import padding
+from cryptography.hazmat.primitives.serialization import load_pem_private_key
+
 from .models import RequestModel, Message
 from .utils import (
     c3s,
@@ -16,7 +27,7 @@ from .utils import (
     c35s,
     c4,
     gemini1,
-    gemini2,
+    gemini_preview,
     gemini2_5_pro_exp,
     BaseAPI,
     safe_get,
@@ -57,7 +68,7 @@ async def get_gemini_payload(request, engine, provider, api_key=None):
 
     try:
         request_messages = [Message(role="user", content=request.prompt)]
-    except:
+    except Exception:
         request_messages = copy.deepcopy(request.messages)
     for msg in request_messages:
         if msg.role == "assistant":
@@ -300,11 +311,6 @@ async def get_gemini_payload(request, engine, provider, api_key=None):
 
     return url, headers, payload
 
-import time
-from cryptography.hazmat.primitives import hashes
-from cryptography.hazmat.primitives.asymmetric import padding
-from cryptography.hazmat.primitives.serialization import load_pem_private_key
-
 def create_jwt(client_email, private_key):
     # JWT Header
     header = json.dumps({
@@ -375,12 +381,14 @@ async def get_vertex_gemini_payload(request, engine, provider, api_key=None):
 
     # https://cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/2-0-flash?hl=zh-cn
     pro_models = ["gemini-2.5"]
-    if any(pro_model in original_model for pro_model in pro_models):
+    if "gemini-2.5-flash-image-preview" in original_model:
+        location = gemini_preview
+    elif any(pro_model in original_model for pro_model in pro_models):
         location = gemini2_5_pro_exp
     else:
         location = gemini1
 
-    if "google-vertex-ai" in provider.get("base_url", ""):
+    if "google-vertex-ai" in provider.get("base_url", "") or "gemini-2.5-flash-image-preview" in original_model:
         url = provider.get("base_url").rstrip('/') + "/v1/projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/{MODEL_ID}:{stream}".format(
             LOCATION=await location.next(),
             PROJECT_ID=project_id,
@@ -754,18 +762,12 @@ async def get_vertex_claude_payload(request, engine, provider, api_key=None):
                         "type": "any"
                     }
 
-    if provider.get("tools") == False:
+    if provider.get("tools") is False:
         payload.pop("tools", None)
         payload.pop("tool_choice", None)
 
     return url, headers, payload
 
-import hashlib
-import hmac
-import datetime
-import urllib.parse
-from datetime import timezone
-
 def sign(key, msg):
     return hmac.new(key, msg.encode('utf-8'), hashlib.sha256).digest()
 
@@ -837,7 +839,7 @@ async def get_aws_payload(request, engine, provider, api_key=None):
     url = f"{base_url}/model/{original_model}/invoke-with-response-stream"
 
     messages = []
-    system_prompt = None
+    # system_prompt = None
     tool_id = None
     for msg in request.messages:
         tool_call_id = None
@@ -887,8 +889,8 @@ async def get_aws_payload(request, engine, provider, api_key=None):
             }]})
         elif msg.role != "system":
             messages.append({"role": msg.role, "content": content})
-        elif msg.role == "system":
-            system_prompt = content
+        # elif msg.role == "system":
+        #     system_prompt = content
 
     conversation_len = len(messages) - 1
     message_index = 0
@@ -971,7 +973,7 @@ async def get_aws_payload(request, engine, provider, api_key=None):
                         "type": "any"
                     }
 
-    if provider.get("tools") == False:
+    if provider.get("tools") is False:
         payload.pop("tools", None)
         payload.pop("tool_choice", None)
 
@@ -1079,7 +1081,7 @@ async def get_gpt_payload(request, engine, provider, api_key=None):
             else:
                 payload[field] = value
 
-    if provider.get("tools") == False or "o1-mini" in original_model or "chatgpt-4o-latest" in original_model or "grok" in original_model:
+    if provider.get("tools") is False or "o1-mini" in original_model or "chatgpt-4o-latest" in original_model or "grok" in original_model:
         payload.pop("tools", None)
         payload.pop("tool_choice", None)
 
@@ -1241,7 +1243,7 @@ async def get_azure_payload(request, engine, provider, api_key=None):
             else:
                 payload[field] = value
 
-    if provider.get("tools") == False or "o1" in original_model or "chatgpt-4o-latest" in original_model or "grok" in original_model:
+    if provider.get("tools") is False or "o1" in original_model or "chatgpt-4o-latest" in original_model or "grok" in original_model:
         payload.pop("tools", None)
         payload.pop("tool_choice", None)
 
@@ -1336,7 +1338,7 @@ async def get_azure_databricks_payload(request, engine, provider, api_key=None):
             else:
                 payload[field] = value
 
-    if provider.get("tools") == False or "o1" in original_model or "chatgpt-4o-latest" in original_model or "grok" in original_model:
+    if provider.get("tools") is False or "o1" in original_model or "chatgpt-4o-latest" in original_model or "grok" in original_model:
         payload.pop("tools", None)
         payload.pop("tool_choice", None)
 
@@ -1626,7 +1628,7 @@ async def gpt2claude_tools_json(json_dict):
     for old_key, new_key in keys_to_change.items():
         if old_key in json_dict:
             if new_key:
-                if json_dict[old_key] == None:
+                if json_dict[old_key] is None:
                     json_dict[old_key] = {
                         "type": "object",
                         "properties": {}
@@ -1788,7 +1790,7 @@ async def get_claude_payload(request, engine, provider, api_key=None):
                         "type": "any"
                     }
 
-    if provider.get("tools") == False:
+    if provider.get("tools") is False:
         payload.pop("tools", None)
         payload.pop("tool_choice", None)
 
@@ -1868,8 +1870,6 @@ async def get_upload_certificate(client: httpx.AsyncClient, api_key: str, model:
         print(f"获取凭证时发生未知错误: {e}")
         return None
 
-from mimetypes import guess_type
-
 async def upload_file_to_oss(client: httpx.AsyncClient, certificate: dict, file: Tuple[str, IOBase, str]) -> str:
     """第二步：使用凭证将文件内容上传到OSS"""
     upload_host = certificate.get("upload_host")

{aient-1.2.49 → aient-1.2.50}/aient/core/utils.py

@@ -4,12 +4,16 @@ import ast
 import json
 import httpx
 import base64
+import random
+import string
 import asyncio
+import traceback
 from time import time
 from PIL import Image
 from fastapi import HTTPException
-from urllib.parse import urlparse
 from collections import defaultdict
+from httpx_socks import AsyncProxyTransport
+from urllib.parse import urlparse, urlunparse
 
 from .log_config import logger
 
@@ -33,7 +37,6 @@ class BaseAPI:
         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 == "":
@@ -152,7 +155,6 @@ def get_engine(provider, endpoint=None, original_model=""):
 
     return engine, stream
 
-from httpx_socks import AsyncProxyTransport
 def get_proxy(proxy, client_config = {}):
     if proxy:
         # 解析代理URL
@@ -223,9 +225,7 @@ async def update_initial_model(provider):
         models_id = list(set_models)
         # print(models_id)
         return models_id
-    except Exception as e:
-        # print("error:", e)
-        import traceback
+    except Exception:
         traceback.print_exc()
         return []
 
@@ -278,27 +278,31 @@ def parse_rate_limit(limit_string):
     return limits
 
 class ThreadSafeCircularList:
-    def __init__(self, items = [], rate_limit={"default": "999999/min"}, schedule_algorithm="round_robin"):
+    def __init__(self, items = [], rate_limit={"default": "999999/min"}, schedule_algorithm="round_robin", provider_name=None):
+        self.provider_name = provider_name
+        self.original_items = list(items)
+        self.schedule_algorithm = schedule_algorithm
+
         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"
+        elif schedule_algorithm == "smart_round_robin":
+            self.items = items
         else:
             self.items = items
-            logger.warning(f"Unknown schedule algorithm: {schedule_algorithm}, use (round_robin, random, fixed_priority) instead")
+            logger.warning(f"Unknown schedule algorithm: {schedule_algorithm}, use (round_robin, random, fixed_priority, smart_round_robin) 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 = {}
+        self.reordering_task = None
+
         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)
@@ -307,6 +311,38 @@ class ThreadSafeCircularList:
         else:
             logger.error(f"Error ThreadSafeCircularList: Unknown rate_limit type: {type(rate_limit)}, rate_limit: {rate_limit}")
 
+        if self.schedule_algorithm == "smart_round_robin":
+            logger.info(f"Initializing '{self.provider_name}' with 'smart_round_robin' algorithm.")
+            self._trigger_reorder()
+
+    async def reset_items(self, new_items: list):
+        """Safely replaces the current list of items with a new one."""
+        async with self.lock:
+            if self.items != new_items:
+                self.items = new_items
+                self.index = 0
+                logger.info(f"Provider '{self.provider_name}' API key list has been reset and reordered.")
+
+    def _trigger_reorder(self):
+        """Asynchronously triggers the reordering task if not already running."""
+        if self.provider_name and (self.reordering_task is None or self.reordering_task.done()):
+            logger.info(f"Triggering reorder for provider '{self.provider_name}'...")
+            try:
+                loop = asyncio.get_running_loop()
+                self.reordering_task = loop.create_task(self._reorder_keys())
+            except RuntimeError:
+                logger.warning(f"No running event loop to trigger reorder for '{self.provider_name}'.")
+
+    async def _reorder_keys(self):
+        """Performs the actual reordering logic."""
+        from utils import get_sorted_api_keys
+        try:
+            sorted_keys = await get_sorted_api_keys(self.provider_name, self.original_items, group_size=100)
+            if sorted_keys:
+                await self.reset_items(sorted_keys)
+        except Exception as e:
+            logger.error(f"Error during key reordering for provider '{self.provider_name}': {e}")
+
     async def set_cooling(self, item: str, cooling_time: int = 60):
         """设置某个 item 进入冷却状态
 
@@ -314,7 +350,7 @@ class ThreadSafeCircularList:
             item: 需要冷却的 item
             cooling_time: 冷却时间(秒)，默认60秒
         """
-        if item == None:
+        if item is None:
             return
         now = time()
         async with self.lock:
@@ -374,6 +410,11 @@ class ThreadSafeCircularList:
         async with self.lock:
             if self.schedule_algorithm == "fixed_priority":
                 self.index = 0
+
+            # 检查是否即将完成一个循环，并据此触发重排序
+            if self.schedule_algorithm == "smart_round_robin" and len(self.items) > 0 and self.index == len(self.items) - 1:
+                self._trigger_reorder()
+
             start_index = self.index
             while True:
                 item = self.items[self.index]
@@ -384,7 +425,7 @@ class ThreadSafeCircularList:
 
                 # 如果已经检查了所有的 API key 都被限制
                 if self.index == start_index:
-                    logger.warning(f"All API keys are rate limited!")
+                    logger.warning("All API keys are rate limited!")
                     raise HTTPException(status_code=429, detail="Too many requests")
 
     async def is_tpr_exceeded(self, model: str = None, tokens: int = 0) -> bool:
@@ -488,7 +529,7 @@ c3o = ThreadSafeCircularList(["us-east5"])
 c4 = ThreadSafeCircularList(["us-east5", "us-central1", "europe-west4", "asia-southeast1"])
 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"])
+gemini_preview = ThreadSafeCircularList(["global"])
 # gemini2_5_pro_exp = ThreadSafeCircularList(["global"])
 gemini2_5_pro_exp = ThreadSafeCircularList([
   "us-central1",
@@ -515,8 +556,6 @@ 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, stop=None):
     random.seed(timestamp)
     random_str = ''.join(random.choices(string.ascii_letters + string.digits, k=29))
@@ -654,7 +693,7 @@ def get_image_format(file_content: bytes):
     try:
         img = Image.open(io.BytesIO(file_content))
         return img.format.lower()
-    except:
+    except Exception:
         return None
 
 def encode_image(file_content: bytes):

{aient-1.2.49 → aient-1.2.50}/aient.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aient
-Version: 1.2.49
+Version: 1.2.50
 Summary: Aient: The Awakening of Agent.
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown

{aient-1.2.49 → aient-1.2.50}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "aient"
-version = "1.2.49"
+version = "1.2.50"
 description = "Aient: The Awakening of Agent."
 readme = "README.md"
 requires-python = ">=3.11"