aient 1.0.92__py3-none-any.whl → 1.0.94__py3-none-any.whl

aient/core/request.py

@@ -31,7 +31,10 @@ async def get_gemini_payload(request, engine, provider, api_key=None):
     model_dict = get_model_dict(provider)
     original_model = model_dict[request.model]
 
-    gemini_stream = "streamGenerateContent"
+    if request.stream:
+        gemini_stream = "streamGenerateContent"
+    else:
+        gemini_stream = "generateContent"
     url = provider['base_url']
     parsed_url = urllib.parse.urlparse(url)
     if "/v1beta" in parsed_url.path:
@@ -295,7 +298,10 @@ async def get_vertex_gemini_payload(request, engine, provider, api_key=None):
     if provider.get("project_id"):
         project_id = provider.get("project_id")
 
-    gemini_stream = "streamGenerateContent"
+    if request.stream:
+        gemini_stream = "streamGenerateContent"
+    else:
+        gemini_stream = "generateContent"
     model_dict = get_model_dict(provider)
     original_model = model_dict[request.model]
     search_tool = None
@@ -941,10 +947,6 @@ async def get_gpt_payload(request, engine, provider, api_key=None):
         payload.pop("tools", None)
         payload.pop("tool_choice", None)
 
-    # if "models.inference.ai.azure.com" in url:
-    #     payload["stream"] = False
-    #     payload.pop("stream_options", None)
-
     if "api.x.ai" in url:
         payload.pop("stream_options", None)
         payload.pop("presence_penalty", None)

aient/core/response.py

@@ -35,6 +35,7 @@ async def fetch_gemini_response_stream(client, url, headers, payload, model):
         promptTokenCount = 0
         candidatesTokenCount = 0
         totalTokenCount = 0
+        parts_json = ""
         # line_index = 0
         # last_text_line = 0
         # if "thinking" in model:
@@ -63,24 +64,25 @@ async def fetch_gemini_response_stream(client, url, headers, payload, model):
                     json_data = parse_json_safely( "{" + line + "}")
                     totalTokenCount = json_data.get('totalTokenCount', 0)
 
-                # print(line)
-                if line and '\"text\": \"' in line and is_finish == False:
+                if (line and '"parts": [' in line or parts_json != "") and is_finish == False:
+                    parts_json += line
+                if parts_json != "" and line and '],' == line.strip():
+                    parts_json =  "{" + parts_json.strip().rstrip(",} ]}") + "}]}"
                     try:
-                        json_data = json.loads( "{" + line + "}")
-                        content = json_data.get('text', '')
-                        # content = content.replace("\n", "\n\n")
-                        # if last_text_line == 0 and is_thinking:
-                        #     content = "> " + content.lstrip()
-                        # if is_thinking:
-                        #     content = content.replace("\n", "\n> ")
-                        # if last_text_line == line_index - 3:
-                        #     is_thinking = False
-                        #     content = "\n\n\n" + content.lstrip()
-                        sse_string = await generate_sse_response(timestamp, model, content=content)
-                        yield sse_string
+                        json_data = json.loads(parts_json)
+
+                        content = safe_get(json_data, "parts", 0, "text", default="")
+
+                        is_thinking = safe_get(json_data, "parts", 0, "thought", default=False)
+                        if is_thinking:
+                            sse_string = await generate_sse_response(timestamp, model, reasoning_content=content)
+                            yield sse_string
+                        else:
+                            sse_string = await generate_sse_response(timestamp, model, content=content)
+                            yield sse_string
                     except json.JSONDecodeError:
-                        logger.error(f"无法解析JSON: {line}")
-                    # last_text_line = line_index
+                        logger.error(f"无法解析JSON: {parts_json}")
+                    parts_json = ""
 
                 if line and ('\"functionCall\": {' in line or revicing_function_call):
                     revicing_function_call = True
@@ -142,7 +144,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 + "}")
+                        json_data = json.loads( "{" + line.strip().rstrip(",") + "}")
                         content = json_data.get('text', '')
                         sse_string = await generate_sse_response(timestamp, model, content=content)
                         yield sse_string
@@ -525,21 +527,28 @@ async def fetch_response(client, url, headers, payload, engine, model):
             parsed_data = ast.literal_eval(str(response_json))
         elif isinstance(response_json, list):
             parsed_data = response_json
+        elif isinstance(response_json, dict):
+            parsed_data = [response_json]
         else:
             logger.error(f"error fetch_response: Unknown response_json type: {type(response_json)}")
             parsed_data = response_json
         # print("parsed_data", json.dumps(parsed_data, indent=4, ensure_ascii=False))
         content = ""
+        reasoning_content = ""
         for item in parsed_data:
             chunk = safe_get(item, "candidates", 0, "content", "parts", 0, "text")
+            is_think = safe_get(item, "candidates", 0, "content", "parts", 0, "thought", default=False)
             # logger.info(f"chunk: {repr(chunk)}")
             if chunk:
-                content += chunk
+                if is_think:
+                    reasoning_content += chunk
+                else:
+                    content += chunk
 
         usage_metadata = safe_get(parsed_data, -1, "usageMetadata")
-        prompt_tokens = usage_metadata.get("promptTokenCount", 0)
-        candidates_tokens = usage_metadata.get("candidatesTokenCount", 0)
-        total_tokens = usage_metadata.get("totalTokenCount", 0)
+        prompt_tokens = safe_get(usage_metadata, "promptTokenCount", default=0)
+        candidates_tokens = safe_get(usage_metadata, "candidatesTokenCount", default=0)
+        total_tokens = safe_get(usage_metadata, "totalTokenCount", default=0)
 
         role = safe_get(parsed_data, -1, "candidates", 0, "content", "role")
         if role == "model":
@@ -552,7 +561,7 @@ async def fetch_response(client, url, headers, payload, engine, model):
         function_call_content = safe_get(parsed_data, -1, "candidates", 0, "content", "parts", 0, "functionCall", "args", default=None)
 
         timestamp = int(datetime.timestamp(datetime.now()))
-        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)
+        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)
 
     elif engine == "claude":
         response_json = response.json()

aient/core/test/test_geminimask.py

@@ -0,0 +1,330 @@
+import httpx
+import base64
+import os # 用于处理 API 密钥
+import asyncio
+import json
+import re
+# python -m core.test.test_geminimask
+from ..utils import get_image_message, safe_get
+
+# --- 请替换为您的实际值 ---
+MODEL_NAME = "gemini-2.5-flash-preview-04-17"
+API_KEY = os.environ.get("GEMINI_API_KEY") # 从环境变量读取密钥
+IMAGE_PATH = os.environ.get("IMAGE_PATH")
+# --------------------------
+
+# 检查 API 密钥是否存在
+if not API_KEY:
+    raise ValueError("请设置 GOOGLE_API_KEY 环境变量或在代码中提供 API 密钥。")
+
+# 确定图片的 MIME 类型
+# 您可以根据文件扩展名进行猜测，或者使用更可靠的库如 python-magic
+if IMAGE_PATH.lower().endswith(".png"):
+    IMAGE_MIME_TYPE = "image/png"
+elif IMAGE_PATH.lower().endswith(".jpg") or IMAGE_PATH.lower().endswith(".jpeg"):
+    IMAGE_MIME_TYPE = "image/jpeg"
+# 添加其他您需要支持的图片类型
+else:
+    raise ValueError(f"不支持的图片格式: {IMAGE_PATH}")
+
+# 读取图片文件并进行 Base64 编码
+try:
+    with open(IMAGE_PATH, "rb") as image_file:
+        image_data = image_file.read()
+        base64_encoded_image = base64.b64encode(image_data).decode("utf-8")
+        # print(base64_encoded_image)
+except FileNotFoundError:
+    print(f"错误：找不到图片文件 '{IMAGE_PATH}'")
+    exit()
+
+# image_message = get_image_message(base64_encoded_image, "gemini")
+image_message = asyncio.run(get_image_message(f"data:{IMAGE_MIME_TYPE};base64," + base64_encoded_image, "gemini"))
+
+# 构建请求 URL
+url = f"https://generativelanguage.googleapis.com/v1beta/models/{MODEL_NAME}:generateContent"
+
+prompt = "Give the segmentation masks for the Search box. Output a JSON list of segmentation masks where each entry contains the 2D bounding box in \"box_2d\" and the mask in \"mask\"."
+# 定义查询参数 (API Key)
+params = {"key": API_KEY}
+
+# 定义请求头
+headers = {"Content-Type": "application/json"}
+
+# 定义请求体 (JSON payload)
+payload = {
+    "contents": [
+        {
+            "parts": [
+                {
+                    "text": prompt
+                },
+                image_message
+            ]
+        }
+    ],
+    "generationConfig": {"thinkingConfig": {"thinkingBudget": 0}},
+}
+
+# 发送 POST 请求
+try:
+    with httpx.Client() as client:
+        response = client.post(url, params=params, headers=headers, json=payload, timeout=60.0) # 增加超时时间
+        response.raise_for_status() # 如果状态码不是 2xx，则抛出异常
+
+    # 您可以在这里添加代码来解析 response.json() 并提取分割掩码
+    text = safe_get(response.json(), "candidates", 0, "content", "parts", 0, "text")
+    # print(text)
+    # 例如: segmentation_masks = response.json().get('candidates', [{}])[0].get('content', {}).get('parts', [{}])[0].get('text')
+
+except httpx.HTTPStatusError as exc:
+    print(f"HTTP 错误发生: {exc.response.status_code} - {exc.response.text}")
+except httpx.RequestError as exc:
+    print(f"发送请求时出错: {exc}")
+except Exception as e:
+    print(f"发生意外错误: {e}")
+
+
+regex_pattern = r'(\[\s*\{.*?\}\s*\])' # 匹配包含至少一个对象的数组
+
+# 使用 re.search 查找第一个匹配项，re.DOTALL 使点号能匹配换行符
+match = re.search(regex_pattern, text, re.DOTALL)
+
+if match:
+    # 提取匹配到的整个 JSON 数组字符串 (group 1 因为模式中有括号)
+    json_string = match.group(1)
+
+    try:
+        # 使用 json.loads() 解析字符串
+        parsed_data = json.loads(json_string)
+        # 使用 json.dumps 美化打印输出
+        print(json.dumps(parsed_data, indent=2, ensure_ascii=False))
+
+        # 例如，获取第一个元素的 label
+        if isinstance(parsed_data, list) and len(parsed_data) > 0:
+            first_item = parsed_data[0]
+            if isinstance(first_item, dict):
+                label = first_item.get('label')
+                print(f"\n第一个元素的 label 是: {label}")
+
+    except json.JSONDecodeError as e:
+        print(f"JSON 解析错误: {e}")
+        print(f"出错的字符串是: {json_string}")
+else:
+    print("在文本中未找到匹配的 JSON 数组。")
+
+
+import io
+from PIL import Image, ImageDraw, ImageFont # pip install Pillow
+
+def extract_box_and_mask_py(parsed_data):
+    """
+    从已解析的 JSON 数据中提取边界框 (box_2d) 和 Base64 编码的掩码 (mask) 数据。
+
+    Args:
+        parsed_data (list): 一个包含字典的列表，每个字典至少包含 'box_2d' 和 'mask' 键。
+                            例如: [{'box_2d': [y1, x1, y2, x2], 'mask': 'data:image/png;base64,...'}, ...]
+
+    Returns:
+        list: 一个包含字典的列表，每个字典包含 'box' (坐标列表)
+              和 'mask_base64' (Base64 字符串)。
+              例如: [{'box': [y1, x1, y2, x2], 'mask_base64': '...'}, ...]
+              坐标系假定为 0-1000 范围。
+    """
+    # 不再需要正则表达式
+    results = []
+    # 检查 parsed_data 是否为列表
+    if not isinstance(parsed_data, list):
+        print(f"Error: Input data is not a list. Received type: {type(parsed_data)}")
+        return results
+
+    for item in parsed_data:
+        if not isinstance(item, dict):
+            print(f"Skipping non-dictionary item in list: {item}")
+            continue
+
+        try:
+            box = item.get('box_2d')
+            mask_data_uri = item.get('mask')
+
+            # 检查 'box_2d' 和 'mask' 是否存在且不为 None
+            if box is None or mask_data_uri is None:
+                print(f"Skipping item due to missing 'box_2d' or 'mask': {item}")
+                continue
+
+            # 从 mask 数据 URI 中提取 Base64 部分
+            # 格式: "data:image/[^;]+;base64,..."
+            if isinstance(mask_data_uri, str) and mask_data_uri.startswith('data:image/') and ';base64,' in mask_data_uri:
+                mask_b64 = mask_data_uri.split(';base64,', 1)[1]
+            else:
+                print(f"Skipping item due to invalid mask format: {mask_data_uri}")
+                continue
+
+            # 验证 box 数据
+            if isinstance(box, list) and len(box) == 4 and all(isinstance(n, int) for n in box):
+                 results.append({"box": box, "mask_base64": mask_b64})
+            else:
+                 print(f"Skipping invalid box format: {box}")
+
+        # 捕捉可能的 KeyError 或其他在字典访问/处理中发生的错误
+        except Exception as e:
+            print(f"Error processing item: {item}. Error: {e}")
+
+    return results
+
+def display_image_with_bounding_boxes_and_masks_py(
+    original_image_path,
+    box_and_mask_data,
+    output_overlay_path="overlay_image.png",
+    output_compare_dir="comparison_outputs"
+):
+    """
+    在原始图像上绘制边界框和掩码，并生成裁剪区域与掩码的对比图。
+
+    Args:
+        original_image_path (str): 原始图像的文件路径。
+        box_and_mask_data (list): extract_box_and_mask_py 的输出列表。
+        output_overlay_path (str): 保存带有叠加效果的图像的路径。
+        output_compare_dir (str): 保存对比图像的目录路径。
+    """
+    try:
+        img_original = Image.open(original_image_path).convert("RGBA")
+        img_width, img_height = img_original.size
+    except FileNotFoundError:
+        print(f"Error: Original image not found at {original_image_path}")
+        return
+    except Exception as e:
+        print(f"Error opening original image: {e}")
+        return
+
+    # 创建一个副本用于绘制叠加效果
+    img_overlay = img_original.copy()
+    draw = ImageDraw.Draw(img_overlay, "RGBA") # 使用 RGBA 模式以支持透明度
+
+    # 定义颜色列表
+    colors_hex = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#FF00FF', '#00FFFF']
+    # 将十六进制颜色转换为 RGBA 元组 (用于绘制)
+    colors_rgba = []
+    for hex_color in colors_hex:
+        h = hex_color.lstrip('#')
+        rgb = tuple(int(h[i:i+2], 16) for i in (0, 2, 4))
+        colors_rgba.append(rgb + (255,)) # (R, G, B, Alpha) - 边框完全不透明
+
+    # 创建输出目录（如果不存在）
+    import os
+    os.makedirs(output_compare_dir, exist_ok=True)
+
+    print(f"Found {len(box_and_mask_data)} box/mask pairs to process.")
+
+    for i, data in enumerate(box_and_mask_data):
+        box_0_1000 = data['box'] # [ymin, xmin, ymax, xmax] in 0-1000 range
+        mask_b64 = data['mask_base64']
+        color_index = i % len(colors_rgba)
+        outline_color = colors_rgba[color_index]
+        # 叠加掩码时使用半透明颜色
+        mask_fill_color = outline_color[:3] + (int(255 * 0.7),) # 70% Alpha
+
+        # --- 1. 坐标转换与验证 ---
+        # 将 0-1000 坐标转换为图像像素坐标 (left, top, right, bottom)
+        # 假设 box 是 [ymin, xmin, ymax, xmax]
+        try:
+            ymin_norm, xmin_norm, ymax_norm, xmax_norm = [c / 1000.0 for c in box_0_1000]
+
+            left   = int(xmin_norm * img_width)
+            top    = int(ymin_norm * img_height)
+            right  = int(xmax_norm * img_width)
+            bottom = int(ymax_norm * img_height)
+
+            # 确保坐标在图像范围内且有效
+            left = max(0, left)
+            top = max(0, top)
+            right = min(img_width, right)
+            bottom = min(img_height, bottom)
+
+            box_width_px = right - left
+            box_height_px = bottom - top
+
+            if box_width_px <= 0 or box_height_px <= 0:
+                print(f"Skipping box {i+1} due to zero or negative dimensions after conversion.")
+                continue
+
+        except Exception as e:
+            print(f"Error processing coordinates for box {i+1}: {box_0_1000}. Error: {e}")
+            continue
+
+        print(f"Processing Box {i+1}: Pixels(L,T,R,B)=({left},{top},{right},{bottom}) Color={colors_hex[color_index]}")
+
+        # --- 2. 在叠加图像上绘制边界框 ---
+        try:
+            draw.rectangle([left, top, right, bottom], outline=outline_color, width=5)
+        except Exception as e:
+             print(f"Error drawing rectangle for box {i+1}: {e}")
+             continue
+
+        # --- 3. 处理并绘制掩码 ---
+        try:
+            # 解码 Base64 掩码数据
+            mask_bytes = base64.b64decode(mask_b64)
+            mask_img_raw = Image.open(io.BytesIO(mask_bytes)).convert("RGBA")
+
+            # 将掩码图像缩放到边界框的像素尺寸
+            mask_img_resized = mask_img_raw.resize((box_width_px, box_height_px), Image.Resampling.NEAREST)
+
+            # 创建一个纯色块，应用掩码的 alpha 通道
+            color_block = Image.new('RGBA', mask_img_resized.size, mask_fill_color)
+
+            # 将带有透明度的颜色块粘贴到叠加图像上，使用掩码的 alpha 通道作为粘贴蒙版
+            # mask_img_resized.split()[-1] 提取 alpha 通道
+            img_overlay.paste(color_block, (left, top), mask=mask_img_resized.split()[-1])
+
+        except base64.binascii.Error:
+             print(f"Error: Invalid Base64 data for mask {i+1}.")
+             continue
+        except Exception as e:
+             print(f"Error processing or drawing mask for box {i+1}: {e}")
+             continue
+
+        # --- 4. 生成对比图 ---
+        try:
+            # 从原始图像中裁剪出边界框区域
+            img_crop = img_original.crop((left, top, right, bottom))
+
+            # 准备掩码预览图（使用原始解码后的掩码，调整大小以匹配裁剪区域）
+            # 这里直接使用缩放后的 mask_img_resized 的 RGB 部分可能更直观
+            mask_preview = mask_img_resized.convert("RGB") # 转换为 RGB 以便保存为常见格式
+
+            # 保存裁剪图和掩码预览图
+            crop_filename = os.path.join(output_compare_dir, f"compare_{i+1}_crop.png")
+            mask_filename = os.path.join(output_compare_dir, f"compare_{i+1}_mask.png")
+            img_crop.save(crop_filename)
+            mask_preview.save(mask_filename)
+            print(f" - Saved comparison: {crop_filename}, {mask_filename}")
+
+        except Exception as e:
+            print(f"Error creating or saving comparison images for box {i+1}: {e}")
+
+    # --- 5. 保存最终的叠加图像 ---
+    try:
+        img_overlay.save(output_overlay_path)
+        print(f"\nOverlay image saved to: {output_overlay_path}")
+        print(f"Comparison images saved in: {output_compare_dir}")
+    except Exception as e:
+        print(f"Error saving the final overlay image: {e}")
+
+
+extracted_data = extract_box_and_mask_py(parsed_data)
+
+if extracted_data:
+    # 确保原始图像存在
+    import os
+    if os.path.exists(IMAGE_PATH):
+            display_image_with_bounding_boxes_and_masks_py(
+                IMAGE_PATH,
+                extracted_data,
+                output_overlay_path="python_overlay_output.png", # 输出带叠加效果的图片名
+                output_compare_dir="python_comparison_outputs" # 输出对比图的文件夹名
+            )
+    else:
+            print(f"Error: Cannot proceed with visualization, image file not found: {IMAGE_PATH}")
+            print("Please update the 'IMAGE_PATH' variable in the script.")
+else:
+    print("No valid box and mask data found in the response text.")

aient/core/utils.py

@@ -70,7 +70,7 @@ def get_engine(provider, endpoint=None, original_model=""):
         engine = "gemini"
     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'):
+    elif parsed_url.netloc.rstrip('/').endswith('azure.com'):
         engine = "azure"
     elif parsed_url.netloc == 'api.cloudflare.com':
         engine = "cloudflare"
@@ -127,6 +127,9 @@ def get_engine(provider, endpoint=None, original_model=""):
         engine = "tts"
         stream = False
 
+    if "stream" in safe_get(provider, "preferences", "post_body_parameter_overrides", default={}):
+        stream = safe_get(provider, "preferences", "post_body_parameter_overrides", "stream")
+
     return engine, stream
 
 from httpx_socks import AsyncProxyTransport
@@ -484,9 +487,17 @@ async def generate_sse_response(timestamp, model, content=None, tools_id=None, f
 
     return sse_response
 
-async def generate_no_stream_response(timestamp, model, content=None, tools_id=None, function_call_name=None, function_call_content=None, role=None, total_tokens=0, prompt_tokens=0, completion_tokens=0):
+async def generate_no_stream_response(timestamp, model, content=None, tools_id=None, function_call_name=None, function_call_content=None, role=None, total_tokens=0, prompt_tokens=0, completion_tokens=0, reasoning_content=None):
     random.seed(timestamp)
     random_str = ''.join(random.choices(string.ascii_letters + string.digits, k=29))
+    message = {
+        "role": role,
+        "content": content,
+        "refusal": None
+    }
+    if reasoning_content:
+        message["reasoning_content"] = reasoning_content
+
     sample_data = {
         "id": f"chatcmpl-{random_str}",
         "object": "chat.completion",
@@ -495,11 +506,7 @@ async def generate_no_stream_response(timestamp, model, content=None, tools_id=N
         "choices": [
             {
                 "index": 0,
-                "message": {
-                    "role": role,
-                    "content": content,
-                    "refusal": None
-                },
+                "message": message,
                 "logprobs": None,
                 "finish_reason": "stop"
             }

aient/models/claude.py

@@ -2,7 +2,6 @@ import os
 import re
 import json
 import copy
-import tiktoken
 import requests
 
 from .base import BaseLLM
@@ -65,39 +64,6 @@ class claude(BaseLLM):
         self.conversation[convo_id] = claudeConversation()
         self.system_prompt = system_prompt or self.system_prompt
 
-    def __truncate_conversation(self, convo_id: str = "default") -> None:
-        """
-        Truncate the conversation
-        """
-        while True:
-            if (
-                self.get_token_count(convo_id) > self.truncate_limit
-                and len(self.conversation[convo_id]) > 1
-            ):
-                # Don't remove the first message
-                self.conversation[convo_id].pop(1)
-            else:
-                break
-
-    def get_token_count(self, convo_id: str = "default") -> int:
-        """
-        Get token count
-        """
-        tiktoken.model.MODEL_TO_ENCODING["claude-2.1"] = "cl100k_base"
-        encoding = tiktoken.encoding_for_model(self.engine)
-
-        num_tokens = 0
-        for message in self.conversation[convo_id]:
-            # every message follows <im_start>{role/name}\n{content}<im_end>\n
-            num_tokens += 5
-            for key, value in message.items():
-                if value:
-                    num_tokens += len(encoding.encode(value))
-                if key == "name":  # if there's a name, the role is omitted
-                    num_tokens += 5  # role is always required and always 1 token
-        num_tokens += 5  # every reply is primed with <im_start>assistant
-        return num_tokens
-
     def ask_stream(
         self,
         prompt: str,
@@ -267,39 +233,6 @@ class claude3(BaseLLM):
         self.conversation[convo_id] = list()
         self.system_prompt = system_prompt or self.system_prompt
 
-    def __truncate_conversation(self, convo_id: str = "default") -> None:
-        """
-        Truncate the conversation
-        """
-        while True:
-            if (
-                self.get_token_count(convo_id) > self.truncate_limit
-                and len(self.conversation[convo_id]) > 1
-            ):
-                # Don't remove the first message
-                self.conversation[convo_id].pop(1)
-            else:
-                break
-
-    def get_token_count(self, convo_id: str = "default") -> int:
-        """
-        Get token count
-        """
-        tiktoken.model.MODEL_TO_ENCODING["claude-2.1"] = "cl100k_base"
-        encoding = tiktoken.encoding_for_model(self.engine)
-
-        num_tokens = 0
-        for message in self.conversation[convo_id]:
-            # every message follows <im_start>{role/name}\n{content}<im_end>\n
-            num_tokens += 5
-            for key, value in message.items():
-                if value:
-                    num_tokens += len(encoding.encode(value))
-                if key == "name":  # if there's a name, the role is omitted
-                    num_tokens += 5  # role is always required and always 1 token
-        num_tokens += 5  # every reply is primed with <im_start>assistant
-        return num_tokens
-
     def ask_stream(
         self,
         prompt: str,

aient/models/groq.py

@@ -1,7 +1,6 @@
 import os
 import json
 import requests
-import tiktoken
 
 from .base import BaseLLM
 
@@ -52,39 +51,6 @@ class groq(BaseLLM):
         self.conversation[convo_id] = list()
         self.system_prompt = system_prompt or self.system_prompt
 
-    def __truncate_conversation(self, convo_id: str = "default") -> None:
-        """
-        Truncate the conversation
-        """
-        while True:
-            if (
-                self.get_token_count(convo_id) > self.truncate_limit
-                and len(self.conversation[convo_id]) > 1
-            ):
-                # Don't remove the first message
-                self.conversation[convo_id].pop(1)
-            else:
-                break
-
-    def get_token_count(self, convo_id: str = "default") -> int:
-        """
-        Get token count
-        """
-        # tiktoken.model.MODEL_TO_ENCODING["mixtral-8x7b-32768"] = "cl100k_base"
-        encoding = tiktoken.get_encoding("cl100k_base")
-
-        num_tokens = 0
-        for message in self.conversation[convo_id]:
-            # every message follows <im_start>{role/name}\n{content}<im_end>\n
-            num_tokens += 5
-            for key, value in message.items():
-                if value:
-                    num_tokens += len(encoding.encode(value))
-                if key == "name":  # if there's a name, the role is omitted
-                    num_tokens += 5  # role is always required and always 1 token
-        num_tokens += 5  # every reply is primed with <im_start>assistant
-        return num_tokens
-
     def ask_stream(
         self,
         prompt: str,

aient/plugins/config.py

@@ -3,8 +3,7 @@ import json
 import inspect
 
 from .registry import registry
-from ..utils.scripts import cut_message
-from ..utils.prompt import search_key_word_prompt, arxiv_doc_user_prompt
+from ..utils.prompt import search_key_word_prompt
 
 async def get_tools_result_async(function_call_name, function_full_response, function_call_max_tokens, engine, robot, api_key, api_url, use_plugins, model, add_message, convo_id, language):
     function_response = ""
@@ -26,10 +25,7 @@ async def get_tools_result_async(function_call_name, function_full_response, fun
                 yield chunk
             else:
                 function_response = "\n\n".join(chunk)
-            # function_response = yield chunk
-        # function_response = yield from eval(function_call_name)(prompt, keywords)
-        function_call_max_tokens = 32000
-        function_response, text_len = cut_message(function_response, function_call_max_tokens, engine)
+
         if function_response:
             function_response = (
                 f"You need to response the following question: {prompt}. Search results is provided inside <Search_results></Search_results> XML tags. Your task is to think about the question step by step and then answer the above question in {language} based on the Search results provided. Please response in {language} and adopt a style that is logical, in-depth, and detailed. Note: In order to make the answer appear highly professional, you should be an expert in textual analysis, aiming to make the answer precise and comprehensive. Directly response markdown format, without using markdown code blocks. For each sentence quoting search results, a markdown ordered superscript number url link must be used to indicate the source, e.g., [¹](https://www.example.com)"
@@ -40,18 +36,13 @@ async def get_tools_result_async(function_call_name, function_full_response, fun
             ).format(function_response)
         else:
             function_response = "无法找到相关信息，停止使用 tools"
-        # user_prompt = f"You need to response the following question: {prompt}. Search results is provided inside <Search_results></Search_results> XML tags. Your task is to think about the question step by step and then answer the above question in {config.language} based on the Search results provided. Please response in {config.language} and adopt a style that is logical, in-depth, and detailed. Note: In order to make the answer appear highly professional, you should be an expert in textual analysis, aiming to make the answer precise and comprehensive. Directly response markdown format, without using markdown code blocks"
-        # self.add_to_conversation(user_prompt, "user", convo_id=convo_id)
+
     elif function_to_call:
         prompt = json.loads(function_full_response)
         if inspect.iscoroutinefunction(function_to_call):
             function_response = await function_to_call(**prompt)
         else:
             function_response = function_to_call(**prompt)
-        function_response, text_len = cut_message(function_response, function_call_max_tokens, engine)
-
-    # if function_call_name == "download_read_arxiv_pdf":
-    #     add_message(arxiv_doc_user_prompt, "user", convo_id=convo_id)
 
     function_response = (
         f"function_response:{function_response}"

aient/plugins/excute_command.py

@@ -1,6 +1,7 @@
 import subprocess
 from .registry import register_tool
 
+import re
 import html
 
 def unescape_html(input_string: str) -> str:
@@ -15,39 +16,137 @@ def unescape_html(input_string: str) -> str:
   """
   return html.unescape(input_string)
 
+def get_python_executable(command: str) -> str:
+    """
+    获取 Python 可执行文件的路径。
+
+    Returns:
+        str: Python 可执行文件的路径。
+    """
+    cmd_parts = command.split(None, 1)
+    if cmd_parts:
+        executable = cmd_parts[0]
+        args_str = cmd_parts[1] if len(cmd_parts) > 1 else ""
+
+        # 检查是否是 python 可执行文件 (如 python, python3, pythonX.Y)
+        is_python_exe = False
+        if executable == "python" or re.match(r"^python[23]?(\.\d+)?$", executable):
+            is_python_exe = True
+
+        if is_python_exe:
+            # 检查参数中是否已经有 -u 选项
+            args_list = args_str.split()
+            has_u_option = "-u" in args_list
+            if not has_u_option:
+                if args_str:
+                    command = f"{executable} -u {args_str}"
+    return command
+
 # 执行命令
 @register_tool()
 def excute_command(command):
     """
-执行命令并返回输出结果
+执行命令并返回输出结果 (标准输出会实时打印到控制台)
 禁止用于查看pdf，禁止使用 pdftotext 命令
 
 参数:
     command: 要执行的命令，可以克隆仓库，安装依赖，运行代码等
 
 返回:
-    命令执行的输出结果或错误信息
+    命令执行的最终状态和收集到的输出/错误信息
     """
     try:
-        # 使用subprocess.run捕获命令输出
-        command = unescape_html(command)
-        result = subprocess.run(command, shell=True, check=True, capture_output=True, text=True)
-        # 返回命令的标准输出
-        if "pip install" in command:
-            stdout_log = "\n".join([x for x in result.stdout.split('\n') if '━━' not in x])
-        else:
-            stdout_log = result.stdout
-        return f"执行命令成功:\n{stdout_log}"
-    except subprocess.CalledProcessError as e:
-        if "pip install" in command:
-            stdout_log = "\n".join([x for x in e.stdout.split('\n') if '━━' not in x])
+        command = unescape_html(command) # 保留 HTML 解码
+
+        command = get_python_executable(command)
+
+
+        # 使用 Popen 以便实时处理输出
+        # bufsize=1 表示行缓冲, universal_newlines=True 与 text=True 效果类似，用于文本模式
+        process = subprocess.Popen(
+            command,
+            shell=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            bufsize=1,
+            universal_newlines=True
+        )
+
+        stdout_lines = []
+
+        # 实时打印 stdout
+        # print(f"--- 开始执行命令: {command} ---")
+        if process.stdout:
+            for line in iter(process.stdout.readline, ''):
+                # 对 pip install 命令的输出进行过滤，去除进度条相关的行
+                if "pip install" in command and '━━' in line:
+                    continue
+                print(line, end='', flush=True) # 实时打印到控制台，并刷新缓冲区
+                stdout_lines.append(line) # 收集行以供后续返回
+            process.stdout.close()
+        # print(f"\n--- 命令实时输出结束 ---")
+
+        # 等待命令完成
+        process.wait()
+
+        # 获取 stderr (命令完成后一次性读取)
+        stderr_output = ""
+        if process.stderr:
+            stderr_output = process.stderr.read()
+            process.stderr.close()
+
+        # 组合最终的 stdout 日志 (已经过 pip install 过滤)
+        final_stdout_log = "".join(stdout_lines)
+
+        if process.returncode == 0:
+            return f"执行命令成功:\n{final_stdout_log}"
         else:
-            stdout_log = e.stdout
-        # 如果命令执行失败，返回错误信息和错误输出
-        return f"执行命令失败 (退出码 {e.returncode}):\n错误: {e.stderr}\n输出: {stdout_log}"
+            return f"执行命令失败 (退出码 {process.returncode}):\n错误: {stderr_output}\n输出: {final_stdout_log}"
+
+    except FileNotFoundError:
+        # 当 shell=True 时，命令未找到通常由 shell 处理，并返回非零退出码。
+        # 此处捕获 FileNotFoundError 主要用于 Popen 自身无法启动命令的场景 (例如 shell 本身未找到)。
+        return f"执行命令失败: 命令或程序未找到 ({command})"
     except Exception as e:
+        # 其他未知异常
         return f"执行命令时发生异常: {e}"
 
 if __name__ == "__main__":
-    print(excute_command("ls -l && echo 'Hello, World!'"))
-    print(excute_command("ls -l && echo 'Hello, World!'"))
+    # print(excute_command("ls -l && echo 'Hello, World!'"))
+    # print(excute_command("ls -l && echo 'Hello, World!'"))
+
+#     tqdm_script = """
+# import time
+# from tqdm import tqdm
+
+# for i in range(10):
+#     print(f"TQDM 进度条测试: {i}")
+#     time.sleep(1)
+# print('\\n-------TQDM 任务完成.')
+# """
+#     processed_tqdm_script = tqdm_script.replace('"', '\\"')
+#     tqdm_command = f"python -u -u -c \"{processed_tqdm_script}\""
+#     # print(f"执行: {tqdm_command}")
+#     print(excute_command(tqdm_command))
+
+
+    # long_running_command_unix = "echo '开始长时间任务...' && for i in 1 2 3; do echo \"正在处理步骤 $i/3...\"; sleep 1; done && echo '长时间任务完成!'"
+    # print(f"执行: {long_running_command_unix}")
+    # print(excute_command(long_running_command_unix))
+
+
+    # long_running_command_unix = "pip install torch"
+    # print(f"执行: {long_running_command_unix}")
+    # print(excute_command(long_running_command_unix))
+
+
+#     python_long_task_command = """
+# python -c "import time; print('Python 长时间任务启动...'); [print(f'Python 任务进度: {i+1}/3', flush=True) or time.sleep(1) for i in range(3)]; print('Python 长时间任务完成.')"
+# """
+#     python_long_task_command = python_long_task_command.strip() # 移除可能的前后空白
+#     print(f"执行: {python_long_task_command}")
+#     print(excute_command(python_long_task_command))
+
+    print(get_python_executable("python -c 'print(123)'"))
+# python -m beswarm.aient.src.aient.plugins.excute_command

aient/utils/scripts.py

@@ -1,33 +1,10 @@
 import os
 import json
-import base64
-import tiktoken
 import requests
 import urllib.parse
 
 from ..core.utils import get_image_message
 
-def get_encode_text(text, model_name):
-    tiktoken.get_encoding("cl100k_base")
-    model_name = "gpt-3.5-turbo"
-    encoding = tiktoken.encoding_for_model(model_name)
-    encode_text = encoding.encode(text, disallowed_special=())
-    return encoding, encode_text
-
-def get_text_token_len(text, model_name):
-    encoding, encode_text = get_encode_text(text, model_name)
-    return len(encode_text)
-
-def cut_message(message: str, max_tokens: int, model_name: str):
-    if type(message) != str:
-        message = str(message)
-    encoding, encode_text = get_encode_text(message, model_name)
-    if len(encode_text) > max_tokens:
-        encode_text = encode_text[:max_tokens]
-        message = encoding.decode(encode_text)
-    encode_text = encoding.encode(message, disallowed_special=())
-    return message, len(encode_text)
-
 def get_doc_from_url(url):
     filename = urllib.parse.unquote(url.split("/")[-1])
     response = requests.get(url, stream=True)

{aient-1.0.92.dist-info → aient-1.0.94.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aient
-Version: 1.0.92
+Version: 1.0.94
 Summary: Aient: The Awakening of Agent.
 Description-Content-Type: text/markdown
 License-File: LICENSE
@@ -15,7 +15,6 @@ Requires-Dist: html2text
 Requires-Dist: httpx-socks
 Requires-Dist: fake-useragent
 Requires-Dist: beautifulsoup4
-Requires-Dist: tiktoken==0.6.0
 Requires-Dist: lxml-html-clean
 Requires-Dist: pdfminer.six==20240706
 Requires-Dist: duckduckgo-search==5.3.1

{aient-1.0.92.dist-info → aient-1.0.94.dist-info}/RECORD

@@ -4,25 +4,26 @@ aient/core/.gitignore,sha256=5JRRlYYsqt_yt6iFvvzhbqh2FTUQMqwo6WwIuFzlGR8,13
 aient/core/__init__.py,sha256=NxjebTlku35S4Dzr16rdSqSTWUvvwEeACe8KvHJnjPg,34
 aient/core/log_config.py,sha256=kz2_yJv1p-o3lUQOwA3qh-LSc3wMHv13iCQclw44W9c,274
 aient/core/models.py,sha256=_1wYZg_n9kb2A3C8xCboyqleH2iHc9scwOvtx9DPeok,7582
-aient/core/request.py,sha256=U0SDf_dOE5EEhzyJA14XMQCiCFUtcivYcISveLBsK64,61405
-aient/core/response.py,sha256=6fo3GKvTKio8nf4cZdizDIYYq7SnBnFeQ6ROvdAIW9k,30959
-aient/core/utils.py,sha256=W-PDhwoJIbmlt4xfyrV9GXHu9TwRk4pivBOcDVXjgsc,26163
+aient/core/request.py,sha256=RChzDuH49gaJE-o5g65h3nCh-OsuHPwLkq8yuyYEcbo,61431
+aient/core/response.py,sha256=8bS1nAoP6QOMDeDvJvZDVAt34kZ1DpWBI3PUGyza0ZU,31447
+aient/core/utils.py,sha256=CAFqWzICaKVysH9GLHBcp-VeOShisLjWGhEsh6-beWo,26365
 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
 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=z-Z0pJfTN2x0cuwfvu0BdMRY9O-RmLwHEnBIJN1x4Fg,6719
 aient/models/chatgpt.py,sha256=-NWkkKxTCyraPYT0YN37NA2rUfOaDNXtvFSQmIE5tS8,45066
-aient/models/claude.py,sha256=thK9P8qkaaoUN3OOJ9Shw4KDs-pAGKPoX4FOPGFXva8,28597
+aient/models/claude.py,sha256=JezghW7y0brl4Y5qiSHvnYR5prQCFywX4RViHt39pGI,26037
 aient/models/duckduckgo.py,sha256=1l7vYCs9SG5SWPCbcl7q6pCcB5AUF_r-a4l9frz3Ogo,8115
 aient/models/gemini.py,sha256=chGLc-8G_DAOxr10HPoOhvVFW1RvMgHd6mt--VyAW98,14730
-aient/models/groq.py,sha256=2JCB0QE1htOprJHI5fZ11R2RtOhsHlsTjbmFyzc8oSM,10084
+aient/models/groq.py,sha256=eXfSOaPxgQEtk4U8qseArN8rFYOFBfMsPwRcDW1nERo,8790
 aient/models/vertex.py,sha256=qVD5l1Q538xXUPulxG4nmDjXE1VoV4yuAkTCpIeJVw0,16795
 aient/plugins/__init__.py,sha256=p3KO6Aa3Lupos4i2SjzLQw1hzQTigOAfEHngsldrsyk,986
 aient/plugins/arXiv.py,sha256=yHjb6PS3GUWazpOYRMKMzghKJlxnZ5TX8z9F6UtUVow,1461
-aient/plugins/config.py,sha256=KnZ5xtb5o41FI2_qvxTEQhssdd3WJc7lIAFNR85INQw,7817
-aient/plugins/excute_command.py,sha256=9bCKFSQCRO2OYEYT-C-Vqmme0TRO7tc7bFjOxlT03Fk,1784
+aient/plugins/config.py,sha256=Vp6CG9ocdC_FAlCMEGtKj45xamir76DFxdJVvURNtog,6539
+aient/plugins/excute_command.py,sha256=u-JOZ21dDcDx1j3O0KVIHAsa6MNuOxHFBdV3iCnTih0,5413
 aient/plugins/get_time.py,sha256=Ih5XIW5SDAIhrZ9W4Qe5Hs1k4ieKPUc_LAd6ySNyqZk,654
 aient/plugins/image.py,sha256=ZElCIaZznE06TN9xW3DrSukS7U3A5_cjk1Jge4NzPxw,2072
 aient/plugins/list_directory.py,sha256=5ubm-mfrj-tanGSDp4M_Tmb6vQb3dx2-XVfQ2yL2G8A,1394
@@ -35,9 +36,9 @@ aient/prompt/__init__.py,sha256=GBtn6-JDT8KHFCcuPpfSNE_aGddg5p4FEyMCy4BfwGs,20
 aient/prompt/agent.py,sha256=3VycHGnUq9OdR5pd_RM0AeLESlpAgBcmzrsesfq82X0,23856
 aient/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aient/utils/prompt.py,sha256=UcSzKkFE4-h_1b6NofI6xgk3GoleqALRKY8VBaXLjmI,11311
-aient/utils/scripts.py,sha256=PPwaJEigPkpciJHUXOag483iq1GjvaLReHDHjkinv6c,26780
-aient-1.0.92.dist-info/licenses/LICENSE,sha256=XNdbcWldt0yaNXXWB_Bakoqnxb3OVhUft4MgMA_71ds,1051
-aient-1.0.92.dist-info/METADATA,sha256=a4NqdE--2nDBpEy7OHvbuGOPdNTU7rlIhKf72PXagHo,5000
-aient-1.0.92.dist-info/WHEEL,sha256=ooBFpIzZCPdw3uqIQsOo4qqbA4ZRPxHnOH7peeONza0,91
-aient-1.0.92.dist-info/top_level.txt,sha256=3oXzrP5sAVvyyqabpeq8A2_vfMtY554r4bVE-OHBrZk,6
-aient-1.0.92.dist-info/RECORD,,
+aient/utils/scripts.py,sha256=JbYHsU3LLtxBcuO_2MWbSgpHpCgtVQe9FGEFJpUyejc,25926
+aient-1.0.94.dist-info/licenses/LICENSE,sha256=XNdbcWldt0yaNXXWB_Bakoqnxb3OVhUft4MgMA_71ds,1051
+aient-1.0.94.dist-info/METADATA,sha256=tTyW2ICb7VYVE7cOdOixBiA9Fj3qFw3-lVGxmRLCHT4,4969
+aient-1.0.94.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
+aient-1.0.94.dist-info/top_level.txt,sha256=3oXzrP5sAVvyyqabpeq8A2_vfMtY554r4bVE-OHBrZk,6
+aient-1.0.94.dist-info/RECORD,,

{aient-1.0.92.dist-info → aient-1.0.94.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.0.1)
+Generator: setuptools (80.3.1)
 Root-Is-Purelib: true
 Tag: py3-none-any