aient 1.0.49__py3-none-any.whl → 1.0.51__py3-none-any.whl

aient/core/request.py

@@ -96,7 +96,7 @@ async def get_gemini_payload(request, engine, provider, api_key=None):
             content[0]["text"] = re.sub(r"_+", "_", content[0]["text"])
             systemInstruction = {"parts": content}
 
-    off_models = ["gemini-2.0-flash-exp", "gemini-1.5", "gemini-2.5-pro"]
+    off_models = ["gemini-2.0-flash", "gemini-1.5", "gemini-2.5-pro"]
     if any(off_model in original_model for off_model in off_models):
         safety_settings = "OFF"
     else:
@@ -294,12 +294,20 @@ async def get_vertex_gemini_payload(request, engine, provider, api_key=None):
         location = gemini1
         search_tool = {"googleSearchRetrieval": {}}
 
-    url = "https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/{MODEL_ID}:{stream}".format(
-        LOCATION=await location.next(),
-        PROJECT_ID=project_id,
-        MODEL_ID=original_model,
-        stream=gemini_stream
-    )
+    if "google-vertex-ai" in provider.get("base_url", ""):
+        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,
+            MODEL_ID=original_model,
+            stream=gemini_stream
+        )
+    else:
+        url = "https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/{MODEL_ID}:{stream}".format(
+            LOCATION=await location.next(),
+            PROJECT_ID=project_id,
+            MODEL_ID=original_model,
+            stream=gemini_stream
+        )
 
     messages = []
     systemInstruction = None

aient/core/response.py

@@ -5,7 +5,7 @@ from datetime import datetime
 
 from .log_config import logger
 
-from .utils import safe_get, generate_sse_response, generate_no_stream_response, end_of_line
+from .utils import safe_get, generate_sse_response, generate_no_stream_response, end_of_line, parse_json_safely
 
 async def check_response(response, error_log):
     if response and not (200 <= response.status_code < 300):
@@ -30,6 +30,9 @@ async def fetch_gemini_response_stream(client, url, headers, payload, model):
         function_full_response = "{"
         need_function_call = False
         is_finish = False
+        promptTokenCount = 0
+        candidatesTokenCount = 0
+        totalTokenCount = 0
         # line_index = 0
         # last_text_line = 0
         # if "thinking" in model:
@@ -42,9 +45,19 @@ async def fetch_gemini_response_stream(client, url, headers, payload, model):
             while "\n" in buffer:
                 line, buffer = buffer.split("\n", 1)
                 # line_index += 1
+
                 if line and '\"finishReason\": \"' in line:
                     is_finish = True
-                    break
+                if is_finish and '\"promptTokenCount\": ' in line:
+                    json_data = parse_json_safely( "{" + line + "}")
+                    promptTokenCount = json_data.get('promptTokenCount', 0)
+                if is_finish and '\"candidatesTokenCount\": ' in line:
+                    json_data = parse_json_safely( "{" + line + "}")
+                    candidatesTokenCount = json_data.get('candidatesTokenCount', 0)
+                if is_finish and '\"totalTokenCount\": ' in line:
+                    json_data = parse_json_safely( "{" + line + "}")
+                    totalTokenCount = json_data.get('totalTokenCount', 0)
+
                 # print(line)
                 if line and '\"text\": \"' in line:
                     try:
@@ -73,9 +86,6 @@ async def fetch_gemini_response_stream(client, url, headers, payload, model):
 
                     function_full_response += line
 
-            if is_finish:
-                break
-
         if need_function_call:
             function_call = json.loads(function_full_response)
             function_call_name = function_call["functionCall"]["name"]
@@ -84,6 +94,10 @@ async def fetch_gemini_response_stream(client, url, headers, payload, model):
             function_full_response = json.dumps(function_call["functionCall"]["args"])
             sse_string = await generate_sse_response(timestamp, model, content=None, tools_id="chatcmpl-9inWv0yEtgn873CxMBzHeCeiHctTV", function_call_name=None, function_call_content=function_full_response)
             yield sse_string
+
+        sse_string = await generate_sse_response(timestamp, model, None, None, None, None, None, totalTokenCount, promptTokenCount, candidatesTokenCount)
+        yield sse_string
+
     yield "data: [DONE]" + end_of_line
 
 async def fetch_vertex_claude_response_stream(client, url, headers, payload, model):
@@ -98,11 +112,29 @@ async def fetch_vertex_claude_response_stream(client, url, headers, payload, mod
         revicing_function_call = False
         function_full_response = "{"
         need_function_call = False
+        is_finish = False
+        promptTokenCount = 0
+        candidatesTokenCount = 0
+        totalTokenCount = 0
+
         async for chunk in response.aiter_text():
             buffer += chunk
             while "\n" in buffer:
                 line, buffer = buffer.split("\n", 1)
                 # logger.info(f"{line}")
+
+                if line and '\"finishReason\": \"' in line:
+                    is_finish = True
+                if is_finish and '\"promptTokenCount\": ' in line:
+                    json_data = parse_json_safely( "{" + line + "}")
+                    promptTokenCount = json_data.get('promptTokenCount', 0)
+                if is_finish and '\"candidatesTokenCount\": ' in line:
+                    json_data = parse_json_safely( "{" + line + "}")
+                    candidatesTokenCount = json_data.get('candidatesTokenCount', 0)
+                if is_finish and '\"totalTokenCount\": ' in line:
+                    json_data = parse_json_safely( "{" + line + "}")
+                    totalTokenCount = json_data.get('totalTokenCount', 0)
+
                 if line and '\"text\": \"' in line:
                     try:
                         json_data = json.loads( "{" + line + "}")
@@ -130,6 +162,10 @@ async def fetch_vertex_claude_response_stream(client, url, headers, payload, mod
             function_full_response = json.dumps(function_call["input"])
             sse_string = await generate_sse_response(timestamp, model, content=None, tools_id=function_call_id, function_call_name=None, function_call_content=function_full_response)
             yield sse_string
+
+        sse_string = await generate_sse_response(timestamp, model, None, None, None, None, None, totalTokenCount, promptTokenCount, candidatesTokenCount)
+        yield sse_string
+
     yield "data: [DONE]" + end_of_line
 
 async def fetch_gpt_response_stream(client, url, headers, payload):

aient/core/utils.py

@@ -1,6 +1,7 @@
 import re
 import io
 import os
+import ast
 import json
 import httpx
 import base64
@@ -65,7 +66,7 @@ def get_engine(provider, endpoint=None, original_model=""):
     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'):
+    elif parsed_url.netloc.rstrip('/').endswith('aiplatform.googleapis.com') or (parsed_url.netloc.rstrip('/').endswith('gateway.ai.cloudflare.com') and "google-vertex-ai" in parsed_url.path):
         engine = "vertex"
     elif parsed_url.netloc.rstrip('/').endswith('openai.azure.com') or parsed_url.netloc.rstrip('/').endswith('services.ai.azure.com'):
         engine = "azure"
@@ -680,4 +681,34 @@ async def get_text_message(message, engine = None):
         return message
     if engine == "cohere":
         return message
-    raise ValueError("Unknown engine")
+    raise ValueError("Unknown engine")
+
+def parse_json_safely(json_str):
+    """
+    尝试解析JSON字符串，先使用ast.literal_eval，失败则使用json.loads
+
+    Args:
+        json_str: 要解析的JSON字符串
+
+    Returns:
+        解析后的Python对象
+
+    Raises:
+        Exception: 当两种方法都失败时抛出异常
+    """
+    try:
+        # 首先尝试使用ast.literal_eval解析
+        return ast.literal_eval(json_str)
+    except (SyntaxError, ValueError):
+        try:
+            # 如果失败，尝试使用json.loads解析
+            return json.loads(json_str, strict=False)
+        except json.JSONDecodeError as e:
+            # 两种方法都失败，抛出异常
+            raise Exception(f"无法解析JSON字符串: {e}")
+
+if __name__ == "__main__":
+    provider = {
+        "base_url": "https://gateway.ai.cloudflare.com/v1/%7Baccount_id%7D/%7Bgateway_id%7D/google-vertex-ai",
+    }
+    print(get_engine(provider))

aient/models/base.py

@@ -81,7 +81,6 @@ class BaseLLM:
                 },
             ],
         }
-        self.truncate_limit: int = 100000
         self.tokens_usage = defaultdict(int)
         self.function_calls_counter = {}
         self.function_call_max_loop = 10

aient/models/chatgpt.py

@@ -445,7 +445,7 @@ class chatgpt(BaseLLM):
                         async for chunk in get_tools_result_async(
                             tool_name, tool_args, function_call_max_tokens,
                             model or self.engine, chatgpt, kwargs.get('api_key', self.api_key),
-                            self.api_url, use_plugins=False, model=model or self.engine,
+                            kwargs.get('api_url', self.api_url.chat_url), use_plugins=False, model=model or self.engine,
                             add_message=self.add_to_conversation, convo_id=convo_id, language=language
                         ):
                             yield chunk
@@ -454,7 +454,7 @@ class chatgpt(BaseLLM):
                             async for chunk in get_tools_result_async(
                                 tool_name, tool_args, function_call_max_tokens,
                                 model or self.engine, chatgpt, kwargs.get('api_key', self.api_key),
-                                self.api_url, use_plugins=False, model=model or self.engine,
+                                kwargs.get('api_url', self.api_url.chat_url), use_plugins=False, model=model or self.engine,
                                 add_message=self.add_to_conversation, convo_id=convo_id, language=language
                             ):
                                 yield chunk

aient/plugins/config.py

@@ -13,7 +13,7 @@ async def get_tools_result_async(function_call_name, function_full_response, fun
     if function_call_name == "get_search_results":
         prompt = json.loads(function_full_response)["query"]
         yield "message_search_stage_1"
-        llm = robot(api_key=api_key, api_url=api_url.source_api_url, engine=engine, use_plugins=use_plugins)
+        llm = robot(api_key=api_key, api_url=api_url, engine=engine, use_plugins=use_plugins)
         keywords = (await llm.ask_async(search_key_word_prompt.format(source=prompt), model=model)).split("\n")
         print("keywords", keywords)
         keywords = [item.replace("三行关键词是：", "") for item in keywords if "\\x" not in item if item != ""]

{aient-1.0.49.dist-info → aient-1.0.51.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aient
-Version: 1.0.49
+Version: 1.0.51
 Summary: Aient: The Awakening of Agent.
 Description-Content-Type: text/markdown
 License-File: LICENSE

{aient-1.0.49.dist-info → aient-1.0.51.dist-info}/RECORD

@@ -3,16 +3,16 @@ aient/core/.git,sha256=lrAcW1SxzRBUcUiuKL5tS9ykDmmTXxyLP3YYU-Y-Q-I,45
 aient/core/__init__.py,sha256=NxjebTlku35S4Dzr16rdSqSTWUvvwEeACe8KvHJnjPg,34
 aient/core/log_config.py,sha256=kz2_yJv1p-o3lUQOwA3qh-LSc3wMHv13iCQclw44W9c,274
 aient/core/models.py,sha256=H3_XuWA7aS25MWZPK1c-5RBiiuxWJbTfE3RAk0Pkc9A,7504
-aient/core/request.py,sha256=OlMkjGMcFAH-ItA1PgPuf2HT-RbI-Ca4JXncWApc3gM,49088
-aient/core/response.py,sha256=7RVSFfGHisejv2SlsHvp0t-N_8OpTS4edQU_NOi5BGU,25822
-aient/core/utils.py,sha256=I0u3WLWaMd4j1ShqKg_tz67m-1wr_uXlWgxGeUjIIiE,25098
+aient/core/request.py,sha256=c963OuUEBe7j1jxiiwipUyzGrbsCwXQIw_XGF_KdL-4,49491
+aient/core/response.py,sha256=uK6a--vHDk61yB4a-0og36S-d4FSO5X6cLeYSxY9G-A,27726
+aient/core/utils.py,sha256=DFpFU8Y-8lzgQlhaDUnao8HmviGoh3-oN8jZR3Dha7E,26150
 aient/core/test/test_base_api.py,sha256=CjfFzMG26r8C4xCPoVkKb3Ac6pp9gy5NUCbZJHoSSsM,393
 aient/core/test/test_image.py,sha256=_T4peNGdXKBHHxyQNx12u-NTyFE8TlYI6NvvagsG2LE,319
 aient/core/test/test_payload.py,sha256=8jBiJY1uidm1jzL-EiK0s6UGmW9XkdsuuKFGrwFhFkw,2755
 aient/models/__init__.py,sha256=ouNDNvoBBpIFrLsk09Q_sq23HR0GbLAKfGLIFmfEuXE,219
 aient/models/audio.py,sha256=kRd-8-WXzv4vwvsTGwnstK-WR8--vr9CdfCZzu8y9LA,1934
-aient/models/base.py,sha256=osN6f1vkO2Dsponq2MzTH-8lABixYfowj46Ky9q12Ps,6855
-aient/models/chatgpt.py,sha256=2W5vjKaMVsxWrMV92kCfucAzB8l8feAK_XCy1A3MMZg,42232
+aient/models/base.py,sha256=PBfQTQI73OE2OuHvw-XN6vD-_y6k_vLjmT_8J4Gelvo,6813
+aient/models/chatgpt.py,sha256=TRvUpgBsQW0oJpxOpT0__fYooy2hbkMSYVdM7CqusnY,42296
 aient/models/claude.py,sha256=thK9P8qkaaoUN3OOJ9Shw4KDs-pAGKPoX4FOPGFXva8,28597
 aient/models/duckduckgo.py,sha256=1l7vYCs9SG5SWPCbcl7q6pCcB5AUF_r-a4l9frz3Ogo,8115
 aient/models/gemini.py,sha256=chGLc-8G_DAOxr10HPoOhvVFW1RvMgHd6mt--VyAW98,14730
@@ -20,7 +20,7 @@ aient/models/groq.py,sha256=2JCB0QE1htOprJHI5fZ11R2RtOhsHlsTjbmFyzc8oSM,10084
 aient/models/vertex.py,sha256=qVD5l1Q538xXUPulxG4nmDjXE1VoV4yuAkTCpIeJVw0,16795
 aient/plugins/__init__.py,sha256=KrCM6kFD1NB96hfhwUZIG8vJcdZVnfpACMew5YOWxSo,956
 aient/plugins/arXiv.py,sha256=yHjb6PS3GUWazpOYRMKMzghKJlxnZ5TX8z9F6UtUVow,1461
-aient/plugins/config.py,sha256=J1x1newErZ44-IzUtm8gT9Rsx0SRrQxIY__p911dJZM,7543
+aient/plugins/config.py,sha256=1eOFH0KVDV5qa-n_r-nyTx4sCstlcmoiGoaHqPjaSpI,7528
 aient/plugins/excute_command.py,sha256=eAoBR6OmEbP7nzUScfRHHK3UwypuE5lxamUro8HmBMk,911
 aient/plugins/get_time.py,sha256=Ih5XIW5SDAIhrZ9W4Qe5Hs1k4ieKPUc_LAd6ySNyqZk,654
 aient/plugins/image.py,sha256=ZElCIaZznE06TN9xW3DrSukS7U3A5_cjk1Jge4NzPxw,2072
@@ -34,8 +34,8 @@ aient/prompt/agent.py,sha256=fYaLRcMZHgM35IIJkeYGeOpJvIhzfXLhpPR4Q3CYIRU,22258
 aient/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aient/utils/prompt.py,sha256=UcSzKkFE4-h_1b6NofI6xgk3GoleqALRKY8VBaXLjmI,11311
 aient/utils/scripts.py,sha256=XCXMRdpWRJb34Znk4t9JkFnvzDzGHVA5Vv5WpUgP2_0,27152
-aient-1.0.49.dist-info/licenses/LICENSE,sha256=XNdbcWldt0yaNXXWB_Bakoqnxb3OVhUft4MgMA_71ds,1051
-aient-1.0.49.dist-info/METADATA,sha256=2y_fd-2FKGrrJXitd5VBnPuZdkhwwMWMbIVq7QIJOSI,4973
-aient-1.0.49.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-aient-1.0.49.dist-info/top_level.txt,sha256=3oXzrP5sAVvyyqabpeq8A2_vfMtY554r4bVE-OHBrZk,6
-aient-1.0.49.dist-info/RECORD,,
+aient-1.0.51.dist-info/licenses/LICENSE,sha256=XNdbcWldt0yaNXXWB_Bakoqnxb3OVhUft4MgMA_71ds,1051
+aient-1.0.51.dist-info/METADATA,sha256=t_zvFpUBihcSS_ke8pZk0DltHeTBj_r8af3f7sO5gAc,4973
+aient-1.0.51.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+aient-1.0.51.dist-info/top_level.txt,sha256=3oXzrP5sAVvyyqabpeq8A2_vfMtY554r4bVE-OHBrZk,6
+aient-1.0.51.dist-info/RECORD,,