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

aient/core/response.py

@@ -3,6 +3,7 @@ import json
 import random
 import string
 import base64
+import asyncio
 from datetime import datetime
 
 from .log_config import logger
@@ -14,19 +15,19 @@ async def check_response(response, error_log):
         error_message = await response.aread()
         error_str = error_message.decode('utf-8', errors='replace')
         try:
-            error_json = json.loads(error_str)
+            error_json = await asyncio.to_thread(json.loads, error_str)
         except json.JSONDecodeError:
             error_json = error_str
         return {"error": f"{error_log} HTTP Error", "status_code": response.status_code, "details": error_json}
     return None
 
-def gemini_json_poccess(response_str):
+async def gemini_json_poccess(response_str):
     promptTokenCount = 0
     candidatesTokenCount = 0
     totalTokenCount = 0
     image_base64 = None
 
-    response_json = json.loads(response_str)
+    response_json = await asyncio.to_thread(json.loads, response_str)
     json_data = safe_get(response_json, "candidates", 0, "content", default=None)
     finishReason = safe_get(response_json, "candidates", 0 , "finishReason", default=None)
     if finishReason:
@@ -75,7 +76,7 @@ async def fetch_gemini_response_stream(client, url, headers, payload, model, tim
                 if line.startswith("data: "):
                     parts_json = line.lstrip("data: ").strip()
                     try:
-                        json.loads(parts_json)
+                        await asyncio.to_thread(json.loads, parts_json)
                     except json.JSONDecodeError:
                         logger.error(f"JSON decode error: {parts_json}")
                         continue
@@ -83,12 +84,12 @@ async def fetch_gemini_response_stream(client, url, headers, payload, model, tim
                     parts_json += line
                     parts_json = parts_json.lstrip("[,")
                     try:
-                        json.loads(parts_json)
+                        await asyncio.to_thread(json.loads, parts_json)
                     except json.JSONDecodeError:
                         continue
 
                 # https://ai.google.dev/api/generate-content?hl=zh-cn#FinishReason
-                is_thinking, reasoning_content, content, image_base64, function_call_name, function_full_response, finishReason, blockReason, promptTokenCount, candidatesTokenCount, totalTokenCount = gemini_json_poccess(parts_json)
+                is_thinking, reasoning_content, content, image_base64, function_call_name, function_full_response, finishReason, blockReason, promptTokenCount, candidatesTokenCount, totalTokenCount = await gemini_json_poccess(parts_json)
 
                 if is_thinking:
                     sse_string = await generate_sse_response(timestamp, model, reasoning_content=reasoning_content)
@@ -159,7 +160,7 @@ async def fetch_vertex_claude_response_stream(client, url, headers, payload, mod
 
                 if line and '\"text\": \"' in line and is_finish == False:
                     try:
-                        json_data = json.loads( "{" + line.strip().rstrip(",") + "}")
+                        json_data = await asyncio.to_thread(json.loads, "{" + line.strip().rstrip(",") + "}")
                         content = json_data.get('text', '')
                         sse_string = await generate_sse_response(timestamp, model, content=content)
                         yield sse_string
@@ -176,7 +177,7 @@ async def fetch_vertex_claude_response_stream(client, url, headers, payload, mod
                     function_full_response += line
 
         if need_function_call:
-            function_call = json.loads(function_full_response)
+            function_call = await asyncio.to_thread(json.loads, function_full_response)
             function_call_name = function_call["name"]
             function_call_id = function_call["id"]
             sse_string = await generate_sse_response(timestamp, model, content=None, tools_id=function_call_id, function_call_name=function_call_name)
@@ -213,7 +214,7 @@ async def fetch_gpt_response_stream(client, url, headers, payload, timeout):
                 if line and not line.startswith(":") and (result:=line.lstrip("data: ").strip()):
                     if result.strip() == "[DONE]":
                         break
-                    line = json.loads(result)
+                    line = await asyncio.to_thread(json.loads, result)
                     line['id'] = f"chatcmpl-{random_str}"
 
                     # 处理 <think> 标签
@@ -327,7 +328,7 @@ async def fetch_azure_response_stream(client, url, headers, payload, timeout):
                 if line and not line.startswith(":") and (result:=line.lstrip("data: ").strip()):
                     if result.strip() == "[DONE]":
                         break
-                    line = json.loads(result)
+                    line = await asyncio.to_thread(json.loads, result)
                     no_stream_content = safe_get(line, "choices", 0, "message", "content", default="")
                     content = safe_get(line, "choices", 0, "delta", "content", default="")
 
@@ -380,7 +381,7 @@ async def fetch_cloudflare_response_stream(client, url, headers, payload, model,
                     line = line.lstrip("data: ")
                     if line == "[DONE]":
                         break
-                    resp: dict = json.loads(line)
+                    resp: dict = await asyncio.to_thread(json.loads, line)
                     message = resp.get("response")
                     if message:
                         sse_string = await generate_sse_response(timestamp, model, content=message)
@@ -401,7 +402,7 @@ async def fetch_cohere_response_stream(client, url, headers, payload, model, tim
             while "\n" in buffer:
                 line, buffer = buffer.split("\n", 1)
                 # logger.info("line: %s", repr(line))
-                resp: dict = json.loads(line)
+                resp: dict = await asyncio.to_thread(json.loads, line)
                 if resp.get("is_finished") == True:
                     break
                 if resp.get("event_type") == "text-generation":
@@ -427,7 +428,7 @@ async def fetch_claude_response_stream(client, url, headers, payload, model, tim
                 # logger.info(line)
 
                 if line.startswith("data:") and (line := line.lstrip("data: ")):
-                    resp: dict = json.loads(line)
+                    resp: dict = await asyncio.to_thread(json.loads, line)
 
                     input_tokens = input_tokens or safe_get(resp, "message", "usage", "input_tokens", default=0)
                     # cache_creation_input_tokens = safe_get(resp, "message", "usage", "cache_creation_input_tokens", default=0)
@@ -486,7 +487,7 @@ async def fetch_aws_response_stream(client, url, headers, payload, model, timeou
                 if not json_match:
                     continue
                 try:
-                    chunk_data = json.loads(json_match.group(0).lstrip('event'))
+                    chunk_data = await asyncio.to_thread(json.loads, json_match.group(0).lstrip('event'))
                 except json.JSONDecodeError:
                     logger.error(f"DEBUG json.JSONDecodeError: {json_match.group(0).lstrip('event')!r}")
                     continue
@@ -496,7 +497,7 @@ async def fetch_aws_response_stream(client, url, headers, payload, model, timeou
                     # 解码 Base64 编码的字节
                     decoded_bytes = base64.b64decode(chunk_data["bytes"])
                     # 将解码后的字节再次解析为 JSON
-                    payload_chunk = json.loads(decoded_bytes.decode('utf-8'))
+                    payload_chunk = await asyncio.to_thread(json.loads, decoded_bytes.decode('utf-8'))
                     # print(f"DEBUG payload_chunk: {payload_chunk!r}")
 
                     text = safe_get(payload_chunk, "delta", "text", default="")
@@ -514,7 +515,7 @@ async def fetch_aws_response_stream(client, url, headers, payload, model, timeou
 
     yield "data: [DONE]" + end_of_line
 
-async def fetch_response(client, url, headers, payload, engine, model, timeout):
+async def fetch_response(client, url, headers, payload, engine, model, timeout=200):
     response = None
     if payload.get("file"):
         file = payload.pop("file")
@@ -530,7 +531,8 @@ async def fetch_response(client, url, headers, payload, engine, model, timeout):
         yield response.read()
 
     elif engine == "gemini" or engine == "vertex-gemini" or engine == "aws":
-        response_json = response.json()
+        response_bytes = await response.aread()
+        response_json = await asyncio.to_thread(json.loads, response_bytes)
         # print("response_json", json.dumps(response_json, indent=4, ensure_ascii=False))
 
         if isinstance(response_json, str):
@@ -585,7 +587,8 @@ async def fetch_response(client, url, headers, payload, engine, model, timeout):
         yield await generate_no_stream_response(timestamp, model, content=content, tools_id=None, function_call_name=function_call_name, function_call_content=function_call_content, role=role, total_tokens=total_tokens, prompt_tokens=prompt_tokens, completion_tokens=candidates_tokens, reasoning_content=reasoning_content, image_base64=image_base64)
 
     elif engine == "claude":
-        response_json = response.json()
+        response_bytes = await response.aread()
+        response_json = await asyncio.to_thread(json.loads, response_bytes)
         # print("response_json", json.dumps(response_json, indent=4, ensure_ascii=False))
 
         content = safe_get(response_json, "content", 0, "text")
@@ -604,7 +607,8 @@ async def fetch_response(client, url, headers, payload, engine, model, timeout):
         yield await generate_no_stream_response(timestamp, model, content=content, tools_id=tools_id, function_call_name=function_call_name, function_call_content=function_call_content, role=role, total_tokens=total_tokens, prompt_tokens=prompt_tokens, completion_tokens=output_tokens)
 
     elif engine == "azure":
-        response_json = response.json()
+        response_bytes = await response.aread()
+        response_json = await asyncio.to_thread(json.loads, response_bytes)
         # 删除 content_filter_results
         if "choices" in response_json:
             for choice in response_json["choices"]:
@@ -618,14 +622,16 @@ async def fetch_response(client, url, headers, payload, engine, model, timeout):
         yield response_json
 
     elif "dashscope.aliyuncs.com" in url and "multimodal-generation" in url:
-        response_json = response.json()
+        response_bytes = await response.aread()
+        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
     else:
-        response_json = response.json()
+        response_bytes = await response.aread()
+        response_json = await asyncio.to_thread(json.loads, response_bytes)
         yield response_json
 
-async def fetch_response_stream(client, url, headers, payload, engine, model, timeout):
+async def fetch_response_stream(client, url, headers, payload, engine, model, timeout=200):
     if engine == "gemini" or engine == "vertex-gemini":
         async for chunk in fetch_gemini_response_stream(client, url, headers, payload, model, timeout):
             yield chunk

aient/models/chatgpt.py

@@ -15,6 +15,31 @@ from ..utils.scripts import safe_get, async_generator_to_sync, parse_function_xm
 from ..core.request import prepare_request_payload
 from ..core.response import fetch_response_stream, fetch_response
 
+class APITimeoutError(Exception):
+    """Custom exception for API timeout errors."""
+    pass
+
+class ValidationError(Exception):
+    """Custom exception for response validation errors."""
+    def __init__(self, message, response_text):
+        super().__init__(message)
+        self.response_text = response_text
+
+class EmptyResponseError(Exception):
+    """Custom exception for empty API responses."""
+    pass
+
+class ModelNotFoundError(Exception):
+    """Custom exception for model not found (404) errors."""
+    pass
+
+class TaskComplete(Exception):
+    """Exception-like signal to indicate the task is complete."""
+    def __init__(self, message):
+        self.completion_message = message
+        super().__init__(f"Task completed with message: {message}")
+
+
 class chatgpt(BaseLLM):
     """
     Official ChatGPT API
@@ -436,7 +461,7 @@ class chatgpt(BaseLLM):
                 yield chunk
 
         if not full_response.strip():
-            raise Exception(json.dumps({"type": "response_empty_error", "message": "Response is empty"}, ensure_ascii=False))
+            raise EmptyResponseError("Response is empty")
 
         if self.print_log:
             self.logger.info(f"total_tokens: {total_tokens}")
@@ -450,7 +475,7 @@ class chatgpt(BaseLLM):
             if self.check_done:
                 # self.logger.info(f"worker Response: {full_response}")
                 if not full_response.strip().endswith('[done]'):
-                    raise Exception(json.dumps({"type": "validation_error", "message": "Response is not ended with [done]", "response": full_response}, ensure_ascii=False))
+                    raise ValidationError("Response is not ended with [done]", response_text=full_response)
                 else:
                     full_response = full_response.strip().rstrip('[done]')
             full_response = full_response.replace("<tool_code>", "").replace("</tool_code>", "")
@@ -494,6 +519,8 @@ class chatgpt(BaseLLM):
                 # 删除 task_complete 跟其他工具一起调用的情况，因为 task_complete 必须单独调用
                 if len(function_parameter) > 1:
                     function_parameter = [tool_dict for tool_dict in function_parameter if tool_dict.get("function_name", "") != "task_complete"]
+                if len(function_parameter) == 1 and function_parameter[0].get("function_name", "") == "task_complete":
+                    raise TaskComplete(safe_get(function_parameter, 0, "parameter", "message", default="The task has been completed."))
 
                 if self.print_log and invalid_tools:
                     self.logger.error(f"invalid_tools: {invalid_tools}")
@@ -739,13 +766,20 @@ class chatgpt(BaseLLM):
                         )
 
                 # 处理正常响应
+                index = 0
                 async for processed_chunk in self._process_stream_response(
                     generator, convo_id=convo_id, function_name=function_name,
                     total_tokens=total_tokens, function_arguments=function_arguments,
                     function_call_id=function_call_id, model=model, language=language,
                     system_prompt=system_prompt, pass_history=pass_history, is_async=True, stream=stream, **kwargs
                 ):
+                    if index == 0:
+                        if "HTTP Error', 'status_code': 524" in processed_chunk:
+                            raise APITimeoutError("Response timeout")
+                        if "HTTP Error', 'status_code': 404" in processed_chunk:
+                            raise ModelNotFoundError(f"Model: {model or self.engine} not found!")
                     yield processed_chunk
+                    index += 1
 
                 # 成功处理，跳出重试循环
                 break
@@ -754,17 +788,25 @@ class chatgpt(BaseLLM):
                 return # Stop iteration
             except httpx.RemoteProtocolError:
                 continue
+            except APITimeoutError:
+                self.logger.warning("API response timeout (524), retrying...")
+                continue
+            except ValidationError as e:
+                self.logger.warning(f"Validation failed: {e}. Retrying with corrective prompt.")
+                need_done_prompt = [
+                    {"role": "assistant", "content": e.response_text},
+                    {"role": "user", "content": "你的消息没有以[done]结尾，请重新输出"}
+                ]
+                continue
+            except EmptyResponseError as e:
+                self.logger.warning(f"{e}, retrying...")
+                continue
+            except TaskComplete as e:
+                raise
+            except ModelNotFoundError as e:
+                raise
             except Exception as e:
                 self.logger.error(f"{e}")
-                if "validation_error" in str(e):
-                    bad_assistant_message = json.loads(str(e))["response"]
-                    need_done_prompt = [
-                        {"role": "assistant", "content": bad_assistant_message},
-                        {"role": "user", "content": "你的消息没有以[done]结尾，请重新输出"}
-                    ]
-                    continue
-                if "response_empty_error" in str(e):
-                    continue
                 import traceback
                 self.logger.error(traceback.format_exc())
                 if "Invalid URL" in str(e):

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

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

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

@@ -3,7 +3,7 @@ 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=2TzRG-CAZsNg66VGIjeFMLSM0U5F8HcS8JOSWwm3Y0c,33530
+aient/core/response.py,sha256=vQFuc3amHiD1hv_OiINRJnh33n79PnbdzMSBSRlqR5E,34309
 aient/core/utils.py,sha256=D98d5Cy1h4ejKtuxS0EEDtL4YqpaZLB5tuXoVP0IBWQ,28462
 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
@@ -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=S_6s_3epO9UPWea9zyf4-cy8ekMZXNct3AuvSJtO9Pg,43056
+aient/models/chatgpt.py,sha256=q62B6cbtHqKrqsQjM24k_1wi_5-UiuxkXa7e2yG_Clg,44661
 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.74.dist-info/licenses/LICENSE,sha256=XNdbcWldt0yaNXXWB_Bakoqnxb3OVhUft4MgMA_71ds,1051
-aient-1.1.74.dist-info/METADATA,sha256=hvI4hLoMgG3nHQaUkSRQqfR3MKWw2G78O2-Qfhh0HRA,4842
-aient-1.1.74.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-aient-1.1.74.dist-info/top_level.txt,sha256=3oXzrP5sAVvyyqabpeq8A2_vfMtY554r4bVE-OHBrZk,6
-aient-1.1.74.dist-info/RECORD,,
+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,,