aient 1.0.48__py3-none-any.whl → 1.0.50__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/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/prompt/agent.py

@@ -237,7 +237,7 @@ When making code changes, NEVER output code to the USER, unless requested. Inste
 Answer the user's request using the relevant tool(s), if they are available. Check that all the required parameters for each tool call are provided or can reasonably be inferred from context. IF there are no relevant tools or there are missing values for required parameters, ask the user to supply these values. If the user provides a specific value for a parameter (for example provided in quotes), make sure to use that value EXACTLY. DO NOT make up values for or ask about optional parameters. Carefully analyze descriptive terms in the request as they may indicate required parameter values that should be included even if not explicitly quoted.
 
 <user_info>
-The user's OS version is win32 10.0.22631. The absolute path of the user's workspace is /d%3A/CodeBase/private/autojs6. The user's shell is C:\WINDOWS\System32\WindowsPowerShell\v1.0\powershell.exe.
+The user's OS version is win32 10.0.22631. The absolute path of the user's workspace is /d%3A/CodeBase/private/autojs6. The user's shell is C:\\WINDOWS\\System32\\WindowsPowerShell\\v1.0\\powershell.exe.
 </user_info>
 
 <tools>

{aient-1.0.48.dist-info → aient-1.0.50.dist-info}/METADATA

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

{aient-1.0.48.dist-info → aient-1.0.50.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/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
@@ -30,12 +30,12 @@ aient/plugins/registry.py,sha256=YknzhieU_8nQ3oKlUSSWDB4X7t2Jx0JnqT2Jd9Xsvfk,357
 aient/plugins/run_python.py,sha256=dgcUwBunMuDkaSKR5bToudVzSdrXVewktDDFUz_iIOQ,4589
 aient/plugins/websearch.py,sha256=yiBzqXK5X220ibR-zko3VDsn4QOnLu1k6E2YOygCeTQ,15185
 aient/prompt/__init__.py,sha256=GBtn6-JDT8KHFCcuPpfSNE_aGddg5p4FEyMCy4BfwGs,20
-aient/prompt/agent.py,sha256=9-0Y-u5_V2xXpse7e_p-_FhIu7Ulcmy77_9RhfWSjUA,22253
+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.48.dist-info/licenses/LICENSE,sha256=XNdbcWldt0yaNXXWB_Bakoqnxb3OVhUft4MgMA_71ds,1051
-aient-1.0.48.dist-info/METADATA,sha256=OaHoXNvNvEDSjsCDOS1TXWGxG_f9lsoGxh9CFsS5T6A,4973
-aient-1.0.48.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-aient-1.0.48.dist-info/top_level.txt,sha256=3oXzrP5sAVvyyqabpeq8A2_vfMtY554r4bVE-OHBrZk,6
-aient-1.0.48.dist-info/RECORD,,
+aient-1.0.50.dist-info/licenses/LICENSE,sha256=XNdbcWldt0yaNXXWB_Bakoqnxb3OVhUft4MgMA_71ds,1051
+aient-1.0.50.dist-info/METADATA,sha256=RmWJ3S5Q7occhCb0_snb3ltvV9MorT0Z51VK2lNvJ7I,4973
+aient-1.0.50.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+aient-1.0.50.dist-info/top_level.txt,sha256=3oXzrP5sAVvyyqabpeq8A2_vfMtY554r4bVE-OHBrZk,6
+aient-1.0.50.dist-info/RECORD,,