aient 1.0.49__tar.gz → 1.0.50__tar.gz

{aient-1.0.49/src/aient.egg-info → aient-1.0.50}/PKG-INFO

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

{aient-1.0.49 → aient-1.0.50}/setup.py

@@ -4,7 +4,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="aient",
-    version="1.0.49",
+    version="1.0.50",
     description="Aient: The Awakening of Agent.",
     long_description=Path.open(Path("README.md"), encoding="utf-8").read(),
     long_description_content_type="text/markdown",

{aient-1.0.49 → aient-1.0.50}/src/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-1.0.49 → aient-1.0.50}/src/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-1.0.49 → aient-1.0.50}/src/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-1.0.49 → aient-1.0.50}/src/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-1.0.49 → aient-1.0.50}/src/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 → aient-1.0.50/src/aient.egg-info}/PKG-INFO

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