aient 1.1.76__py3-none-any.whl → 1.1.78__py3-none-any.whl

aient/core/request.py

@@ -3,9 +3,9 @@ import json
 import httpx
 import base64
 import asyncio
-import urllib.parse
 from io import IOBase
 from typing import Tuple
+from urllib.parse import urlparse
 
 from .models import RequestModel, Message
 from .utils import (
@@ -41,7 +41,7 @@ async def get_gemini_payload(request, engine, provider, api_key=None):
     else:
         gemini_stream = "generateContent"
     url = provider['base_url']
-    parsed_url = urllib.parse.urlparse(url)
+    parsed_url = urlparse(url)
     if "/v1beta" in parsed_url.path:
         api_version = "v1beta"
     else:
@@ -1952,21 +1952,40 @@ async def get_embedding_payload(request, engine, provider, api_key=None):
     headers = {
         "Content-Type": "application/json",
     }
-    if api_key:
-        headers['Authorization'] = f"Bearer {api_key}"
-    url = provider['base_url']
-    url = BaseAPI(url).embeddings
 
-    payload = {
-        "input": request.input,
-        "model": original_model,
-    }
+    url = provider['base_url']
+    parsed_url = urlparse(url)
+    if "embedding-00" in original_model and "127.0.0.1" not in url and \
+    (parsed_url.path.endswith("/v1beta") or \
+    parsed_url.path.endswith("/v1") or \
+    (parsed_url.netloc == 'generativelanguage.googleapis.com' and "openai/chat/completions" not in parsed_url.path)):
+        if api_key:
+            headers['x-goog-api-key'] = f"{api_key}"
+        parsed_url = urllib.parse.urlparse(url)
+        url = f"{parsed_url.scheme}://{parsed_url.netloc}{parsed_url.path.split('/models')[0].rstrip('/')}/models/{original_model}:embedContent"
+        payload = {
+            "content": {
+                "parts": [
+                    {
+                        "text": request.input
+                    }
+                ]
+            }
+        }
+    else:
+        if api_key:
+            headers['Authorization'] = f"Bearer {api_key}"
+        url = BaseAPI(url).embeddings
+        payload = {
+            "input": request.input,
+            "model": original_model,
+        }
 
-    if request.encoding_format:
-        if url.startswith("https://api.jina.ai"):
-            payload["embedding_type"] = request.encoding_format
-        else:
-            payload["encoding_format"] = request.encoding_format
+        if request.encoding_format:
+            if url.startswith("https://api.jina.ai"):
+                payload["embedding_type"] = request.encoding_format
+            else:
+                payload["encoding_format"] = request.encoding_format
 
     return url, headers, payload
 

aient/core/response.py

@@ -626,6 +626,28 @@ async def fetch_response(client, url, headers, payload, engine, model, timeout=2
         response_json = await asyncio.to_thread(json.loads, response_bytes)
         content = safe_get(response_json, "output", "choices", 0, "message", "content", 0, default=None)
         yield content
+
+    elif "embedContent" in url:
+        response_bytes = await response.aread()
+        response_json = await asyncio.to_thread(json.loads, response_bytes)
+        content = safe_get(response_json, "embedding", "values", default=[])
+        response_embedContent = {
+            "object": "list",
+            "data": [
+                {
+                    "object": "embedding",
+                    "embedding":content,
+                    "index": 0
+                }
+            ],
+            "model": model,
+            "usage": {
+                "prompt_tokens": 0,
+                "total_tokens": 0
+            }
+        }
+
+        yield response_embedContent
     else:
         response_bytes = await response.aread()
         response_json = await asyncio.to_thread(json.loads, response_bytes)

aient/core/utils.py

@@ -518,7 +518,7 @@ async def generate_sse_response(timestamp, model, content=None, tools_id=None, f
         sample_data["choices"][0]["delta"] = {}
         sample_data["choices"][0]["finish_reason"] = stop
 
-    json_data = json.dumps(sample_data, ensure_ascii=False)
+    json_data = await asyncio.to_thread(json.dumps, sample_data, ensure_ascii=False)
     # print("json_data", json.dumps(sample_data, indent=4, ensure_ascii=False))
 
     # 构建SSE响应
@@ -557,6 +557,9 @@ async def generate_no_stream_response(timestamp, model, content=None, tools_id=N
     if function_call_name:
         if not tools_id:
             tools_id = f"call_{random_str}"
+
+        arguments_json = await asyncio.to_thread(json.dumps, function_call_content, ensure_ascii=False)
+
         sample_data = {
             "id": f"chatcmpl-{random_str}",
             "object": "chat.completion",
@@ -574,7 +577,7 @@ async def generate_no_stream_response(timestamp, model, content=None, tools_id=N
                             "type": "function",
                             "function": {
                                 "name": function_call_name,
-                                "arguments": json.dumps(function_call_content, ensure_ascii=False)
+                                "arguments": arguments_json
                             }
                         }
                     ],
@@ -605,7 +608,7 @@ async def generate_no_stream_response(timestamp, model, content=None, tools_id=N
     if total_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)
+    json_data = await asyncio.to_thread(json.dumps, sample_data, ensure_ascii=False)
     # print("json_data", json.dumps(sample_data, indent=4, ensure_ascii=False))
 
     return json_data

aient/models/chatgpt.py

@@ -33,6 +33,18 @@ class ModelNotFoundError(Exception):
     """Custom exception for model not found (404) errors."""
     pass
 
+class RateLimitError(Exception):
+    """Custom exception for rate limit (429) errors."""
+    pass
+
+class ConfigurationError(Exception):
+    """Custom exception for configuration errors."""
+    pass
+
+class RetryFailedError(Exception):
+    """Custom exception for retry failures."""
+    pass
+
 class TaskComplete(Exception):
     """Exception-like signal to indicate the task is complete."""
     def __init__(self, message):
@@ -727,7 +739,7 @@ class chatgpt(BaseLLM):
         need_done_prompt = False
 
         # 发送请求并处理响应
-        for i in range(30):
+        for i in range(10):
             tmp_post_json = copy.deepcopy(json_post)
             if need_done_prompt:
                 tmp_post_json["messages"].extend(need_done_prompt)
@@ -778,19 +790,24 @@ class chatgpt(BaseLLM):
                             raise APITimeoutError("Response timeout")
                         if "HTTP Error', 'status_code': 404" in processed_chunk:
                             raise ModelNotFoundError(f"Model: {model or self.engine} not found!")
+                        if "HTTP Error', 'status_code': 429" in processed_chunk:
+                            raise RateLimitError(f"Rate limit exceeded for model: {model or self.engine}")
                     yield processed_chunk
                     index += 1
 
                 # 成功处理，跳出重试循环
                 break
-            except (httpx.ConnectError, httpx.ReadTimeout):
-                self.logger.error("连接或读取超时错误，请检查服务器状态或网络连接。")
+            except (httpx.ConnectError, httpx.ReadTimeout, httpx.PoolTimeout):
+                self.logger.error("Connection or read timeout.")
                 return # Stop iteration
             except httpx.RemoteProtocolError:
                 continue
             except APITimeoutError:
                 self.logger.warning("API response timeout (524), retrying...")
                 continue
+            except RateLimitError as e:
+                self.logger.warning(f"{e}, retrying...")
+                continue
             except ValidationError as e:
                 self.logger.warning(f"Validation failed: {e}. Retrying with corrective prompt.")
                 need_done_prompt = [
@@ -811,10 +828,10 @@ class chatgpt(BaseLLM):
                 self.logger.error(traceback.format_exc())
                 if "Invalid URL" in str(e):
                     error_message = "您输入了无效的API URL，请使用正确的URL并使用`/start`命令重新设置API URL。具体错误如下：\n\n" + str(e)
-                    raise Exception(json.dumps({"type": "configuration_error", "message": error_message}, ensure_ascii=False))
+                    raise ConfigurationError(error_message)
                 # 最后一次重试失败，向上抛出异常
-                if i == 11:
-                    raise Exception(json.dumps({"type": "retry_failed", "message": str(e)}, ensure_ascii=False))
+                if i == 10:
+                    raise RetryFailedError(str(e))
 
     def ask_stream(
         self,
@@ -941,4 +958,4 @@ class chatgpt(BaseLLM):
             {"role": "system", "content": self.system_prompt},
         ]
         self.tokens_usage[convo_id] = 0
-        self.current_tokens[convo_id] = 0
+        self.current_tokens[convo_id] = 0

{aient-1.1.76.dist-info → aient-1.1.78.dist-info}/METADATA

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

{aient-1.1.76.dist-info → aient-1.1.78.dist-info}/RECORD

@@ -2,9 +2,9 @@ aient/__init__.py,sha256=SRfF7oDVlOOAi6nGKiJIUK6B_arqYLO9iSMp-2IZZps,21
 aient/core/__init__.py,sha256=NxjebTlku35S4Dzr16rdSqSTWUvvwEeACe8KvHJnjPg,34
 aient/core/log_config.py,sha256=kz2_yJv1p-o3lUQOwA3qh-LSc3wMHv13iCQclw44W9c,274
 aient/core/models.py,sha256=KMlCRLjtq1wQHZTJGqnbWhPS2cHq6eLdnk7peKDrzR8,7490
-aient/core/request.py,sha256=vfwi3ZGYp2hQzSJ6mPXJVgcV_uu5AJ_NAL84mLfF8WA,76674
-aient/core/response.py,sha256=vQFuc3amHiD1hv_OiINRJnh33n79PnbdzMSBSRlqR5E,34309
-aient/core/utils.py,sha256=D98d5Cy1h4ejKtuxS0EEDtL4YqpaZLB5tuXoVP0IBWQ,28462
+aient/core/request.py,sha256=nSLKNI-T6lI4yv873eQUElz5_fkJVy1LbOy27NiJPNk,77512
+aient/core/response.py,sha256=ye6Ie5HevXVcH3X5V5BoOC5yDJMBKTKopWQzsCNs008,34977
+aient/core/utils.py,sha256=wCPw4upG4-ciBxGyEXeTEHzxTbm38rVHrr4PIu8hDW8,28579
 aient/core/test/test_base_api.py,sha256=pWnycRJbuPSXKKU9AQjWrMAX1wiLC_014Qc9hh5C2Pw,524
 aient/core/test/test_geminimask.py,sha256=HFX8jDbNg_FjjgPNxfYaR-0-roUrOO-ND-FVsuxSoiw,13254
 aient/core/test/test_image.py,sha256=_T4peNGdXKBHHxyQNx12u-NTyFE8TlYI6NvvagsG2LE,319
@@ -12,7 +12,7 @@ aient/core/test/test_payload.py,sha256=8jBiJY1uidm1jzL-EiK0s6UGmW9XkdsuuKFGrwFhF
 aient/models/__init__.py,sha256=ZTiZgbfBPTjIPSKURE7t6hlFBVLRS9lluGbmqc1WjxQ,43
 aient/models/audio.py,sha256=kRd-8-WXzv4vwvsTGwnstK-WR8--vr9CdfCZzu8y9LA,1934
 aient/models/base.py,sha256=-nnihYnx-vHZMqeVO9ljjt3k4FcD3n-iMk4tT-10nRQ,7232
-aient/models/chatgpt.py,sha256=q62B6cbtHqKrqsQjM24k_1wi_5-UiuxkXa7e2yG_Clg,44661
+aient/models/chatgpt.py,sha256=MLoDnKERg_5GkLG3q0Bzh-rneHcTbuBhHNph03YOVnw,45108
 aient/plugins/__init__.py,sha256=p3KO6Aa3Lupos4i2SjzLQw1hzQTigOAfEHngsldrsyk,986
 aient/plugins/arXiv.py,sha256=yHjb6PS3GUWazpOYRMKMzghKJlxnZ5TX8z9F6UtUVow,1461
 aient/plugins/config.py,sha256=TGgZ5SnNKZ8MmdznrZ-TEq7s2ulhAAwTSKH89bci3dA,7079
@@ -30,8 +30,8 @@ aient/plugins/write_file.py,sha256=Jt8fOEwqhYiSWpCbwfAr1xoi_BmFnx3076GMhuL06uI,3
 aient/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aient/utils/prompt.py,sha256=UcSzKkFE4-h_1b6NofI6xgk3GoleqALRKY8VBaXLjmI,11311
 aient/utils/scripts.py,sha256=VqtK4RFEx7KxkmcqG3lFDS1DxoNlFFGErEjopVcc8IE,40974
-aient-1.1.76.dist-info/licenses/LICENSE,sha256=XNdbcWldt0yaNXXWB_Bakoqnxb3OVhUft4MgMA_71ds,1051
-aient-1.1.76.dist-info/METADATA,sha256=nOBPFlGsNRfFqblnwjC4Z36Dq8TkUMcsdTDrI9Gcm8E,4842
-aient-1.1.76.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-aient-1.1.76.dist-info/top_level.txt,sha256=3oXzrP5sAVvyyqabpeq8A2_vfMtY554r4bVE-OHBrZk,6
-aient-1.1.76.dist-info/RECORD,,
+aient-1.1.78.dist-info/licenses/LICENSE,sha256=XNdbcWldt0yaNXXWB_Bakoqnxb3OVhUft4MgMA_71ds,1051
+aient-1.1.78.dist-info/METADATA,sha256=0U1QX6xBqPt0B26kwN340cm3TG3K8GWdjQs81eThnbA,4842
+aient-1.1.78.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+aient-1.1.78.dist-info/top_level.txt,sha256=3oXzrP5sAVvyyqabpeq8A2_vfMtY554r4bVE-OHBrZk,6
+aient-1.1.78.dist-info/RECORD,,