aient 1.1.76__tar.gz → 1.1.78__tar.gz

{aient-1.1.76 → aient-1.1.78}/PKG-INFO

@@ -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 → aient-1.1.78}/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-1.1.76 → aient-1.1.78}/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-1.1.76 → aient-1.1.78}/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-1.1.76 → aient-1.1.78}/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 → aient-1.1.78}/aient.egg-info/PKG-INFO

@@ -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 → aient-1.1.78}/pyproject.toml

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