beswarm 0.2.59__py3-none-any.whl → 0.2.61__py3-none-any.whl

beswarm/agents/planact.py

@@ -4,6 +4,7 @@ import copy
 import json
 import difflib
 import asyncio
+import tomllib
 import platform
 from pathlib import Path
 from datetime import datetime
@@ -11,6 +12,7 @@ from typing import List, Dict, Union
 
 from ..broker import MessageBroker
 from ..aient.aient.models import chatgpt
+from ..aient.aient.models.chatgpt import ModelNotFoundError, TaskComplete
 from ..aient.aient.plugins import get_function_call_list, registry
 from ..prompt import worker_system_prompt, instruction_system_prompt
 from ..utils import extract_xml_content, get_current_screen_image_message, replace_xml_content, register_mcp_tools, setup_logger
@@ -112,12 +114,6 @@ class InstructionAgent(BaseAgent):
     async def handle_message(self, message: Dict):
         """Receives a worker response, generates the next instruction, and publishes it."""
 
-        if len(message["conversation"]) > 1 and message["conversation"][-2]["role"] == "user" \
-        and "<task_complete_message>" in message["conversation"][-2]["content"]:
-            task_complete_message = extract_xml_content(message["conversation"][-2]["content"], "task_complete_message")
-            self.broker.publish({"status": "finished", "result": task_complete_message}, self.status_topic)
-            return
-
         instruction_prompt = "".join([
                 "</work_agent_conversation_end>\n\n",
                 f"任务目标: {self.goal}\n\n",
@@ -137,10 +133,11 @@ class InstructionAgent(BaseAgent):
         if "find_and_click_element" in json.dumps(self.tools_json):
             instruction_prompt = await get_current_screen_image_message(instruction_prompt)
 
-        raw_response = await self.agent.ask_async(instruction_prompt)
+        try:
+            raw_response = await self.agent.ask_async(instruction_prompt)
+        except ModelNotFoundError as e:
+            raise Exception(str(e))
 
-        if "HTTP Error', 'status_code': 404" in raw_response:
-            raise Exception(f"Model: {self.config['engine']} not found!")
         if "'status_code': 413" in raw_response or \
         "'status_code': 400" in raw_response:
             self.broker.publish({"status": "error", "result": "The request body is too long, please try again."}, self.status_topic)
@@ -205,14 +202,16 @@ class WorkerAgent(BaseAgent):
         instruction = message["instruction"]
         if "find_and_click_element" in json.dumps(self.tools_json):
             instruction = await get_current_screen_image_message(instruction)
-        response = await self.agent.ask_async(instruction)
+
+        try:
+            response = await self.agent.ask_async(instruction)
+        except TaskComplete as e:
+            self.broker.publish({"status": "finished", "result": e.completion_message}, self.status_topic)
+            return
 
         if response.strip() == '':
             self.logger.error("\n❌ 工作智能体回复为空，请重新生成指令。")
             self.broker.publish(message, self.error_topic)
-        elif "HTTP Error', 'status_code': 524" in response:
-            self.logger.error("\n❌ 工作智能体回复超时 100 秒，请重新生成指令。")
-            self.broker.publish(message, self.error_topic)
         else:
             self.broker.publish({"status": "new_message", "result": "\n✅ 工作智能体:\n" + response}, self.status_topic)
             self.broker.publish({
@@ -258,6 +257,12 @@ class BrokerWorker:
         self.logger = setup_logger(f"task_{self.work_dir.name}", log_file_path)
         self.logger.info(f"Logger for task '{self.goal}' initialized. Log file: {log_file_path}")
 
+        beswarm_dir = Path(__file__).parent.parent.parent
+        with open(beswarm_dir / "pyproject.toml", "rb") as f:
+            pyproject_data = tomllib.load(f)
+            version = pyproject_data["project"]["version"]
+            self.logger.info(f"beswarm version: {version}")
+
     async def _configure_tools(self):
         mcp_list = [item for item in self.tools if isinstance(item, dict)]
         if mcp_list:
@@ -275,6 +280,7 @@ class BrokerWorker:
     def _task_status_subscriber(self, message: Dict):
         """Subscriber for task status changes."""
         if message.get("status") == "finished":
+            self.logger.info("Task completed: " + message.get("result"))
             self.final_result = message.get("result")
             self.task_completion_event.set()
 

beswarm/aient/aient/core/models.py

@@ -1,4 +1,3 @@
-import json
 from io import IOBase
 from pydantic import BaseModel, Field, model_validator, ConfigDict
 from typing import List, Dict, Optional, Union, Tuple, Literal, Any

beswarm/aient/aient/core/request.py

@@ -2,6 +2,7 @@ import re
 import json
 import httpx
 import base64
+import asyncio
 import urllib.parse
 from io import IOBase
 from typing import Tuple
@@ -336,11 +337,11 @@ def create_jwt(client_email, private_key):
     segments.append(base64.urlsafe_b64encode(signature).rstrip(b'='))
     return b'.'.join(segments).decode()
 
-def get_access_token(client_email, private_key):
-    jwt = create_jwt(client_email, private_key)
+async def get_access_token(client_email, private_key):
+    jwt = await asyncio.to_thread(create_jwt, client_email, private_key)
 
-    with httpx.Client() as client:
-        response = client.post(
+    async with httpx.AsyncClient() as client:
+        response = await client.post(
             "https://oauth2.googleapis.com/token",
             data={
                 "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
@@ -356,7 +357,7 @@ async def get_vertex_gemini_payload(request, engine, provider, api_key=None):
         'Content-Type': 'application/json'
     }
     if provider.get("client_email") and provider.get("private_key"):
-        access_token = get_access_token(provider['client_email'], provider['private_key'])
+        access_token = await get_access_token(provider['client_email'], provider['private_key'])
         headers['Authorization'] = f"Bearer {access_token}"
     if provider.get("project_id"):
         project_id = provider.get("project_id")
@@ -596,7 +597,7 @@ async def get_vertex_claude_payload(request, engine, provider, api_key=None):
         'Content-Type': 'application/json',
     }
     if provider.get("client_email") and provider.get("private_key"):
-        access_token = get_access_token(provider['client_email'], provider['private_key'])
+        access_token = await get_access_token(provider['client_email'], provider['private_key'])
         headers['Authorization'] = f"Bearer {access_token}"
     if provider.get("project_id"):
         project_id = provider.get("project_id")
@@ -972,7 +973,9 @@ async def get_aws_payload(request, engine, provider, api_key=None):
 
     if provider.get("aws_access_key") and provider.get("aws_secret_key"):
         ACCEPT_HEADER = "application/vnd.amazon.bedrock.payload+json" # 指定接受 Bedrock 流格式
-        amz_date, payload_hash, authorization_header = get_signature(payload, original_model, provider.get("aws_access_key"), provider.get("aws_secret_key"), AWS_REGION, HOST, CONTENT_TYPE, ACCEPT_HEADER)
+        amz_date, payload_hash, authorization_header = await asyncio.to_thread(
+            get_signature, payload, original_model, provider.get("aws_access_key"), provider.get("aws_secret_key"), AWS_REGION, HOST, CONTENT_TYPE, ACCEPT_HEADER
+        )
         headers = {
             'Accept': ACCEPT_HEADER,
             'Content-Type': CONTENT_TYPE,

beswarm/aient/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:
@@ -53,9 +54,9 @@ def gemini_json_poccess(response_str):
 
     return is_thinking, reasoning_content, content, image_base64, function_call_name, function_full_response, finishReason, blockReason, promptTokenCount, candidatesTokenCount, totalTokenCount
 
-async def fetch_gemini_response_stream(client, url, headers, payload, model):
+async def fetch_gemini_response_stream(client, url, headers, payload, model, timeout):
     timestamp = int(datetime.timestamp(datetime.now()))
-    async with client.stream('POST', url, headers=headers, json=payload) as response:
+    async with client.stream('POST', url, headers=headers, json=payload, timeout=timeout) as response:
         error_message = await check_response(response, "fetch_gemini_response_stream")
         if error_message:
             yield error_message
@@ -75,7 +76,7 @@ async def fetch_gemini_response_stream(client, url, headers, payload, model):
                 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):
                     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)
@@ -122,9 +123,9 @@ async def fetch_gemini_response_stream(client, url, headers, payload, model):
 
     yield "data: [DONE]" + end_of_line
 
-async def fetch_vertex_claude_response_stream(client, url, headers, payload, model):
+async def fetch_vertex_claude_response_stream(client, url, headers, payload, model, timeout):
     timestamp = int(datetime.timestamp(datetime.now()))
-    async with client.stream('POST', url, headers=headers, json=payload) as response:
+    async with client.stream('POST', url, headers=headers, json=payload, timeout=timeout) as response:
         error_message = await check_response(response, "fetch_vertex_claude_response_stream")
         if error_message:
             yield error_message
@@ -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)
@@ -190,14 +191,14 @@ async def fetch_vertex_claude_response_stream(client, url, headers, payload, mod
 
     yield "data: [DONE]" + end_of_line
 
-async def fetch_gpt_response_stream(client, url, headers, payload):
+async def fetch_gpt_response_stream(client, url, headers, payload, timeout):
     timestamp = int(datetime.timestamp(datetime.now()))
     random.seed(timestamp)
     random_str = ''.join(random.choices(string.ascii_letters + string.digits, k=29))
     is_thinking = False
     has_send_thinking = False
     ark_tag = False
-    async with client.stream('POST', url, headers=headers, json=payload) as response:
+    async with client.stream('POST', url, headers=headers, json=payload, timeout=timeout) as response:
         error_message = await check_response(response, "fetch_gpt_response_stream")
         if error_message:
             yield error_message
@@ -213,7 +214,7 @@ async def fetch_gpt_response_stream(client, url, headers, payload):
                 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> 标签
@@ -306,12 +307,12 @@ async def fetch_gpt_response_stream(client, url, headers, payload):
                         yield "data: " + json.dumps(line).strip() + end_of_line
     yield "data: [DONE]" + end_of_line
 
-async def fetch_azure_response_stream(client, url, headers, payload):
+async def fetch_azure_response_stream(client, url, headers, payload, timeout):
     timestamp = int(datetime.timestamp(datetime.now()))
     is_thinking = False
     has_send_thinking = False
     ark_tag = False
-    async with client.stream('POST', url, headers=headers, json=payload) as response:
+    async with client.stream('POST', url, headers=headers, json=payload, timeout=timeout) as response:
         error_message = await check_response(response, "fetch_azure_response_stream")
         if error_message:
             yield error_message
@@ -327,7 +328,7 @@ async def fetch_azure_response_stream(client, url, headers, payload):
                 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="")
 
@@ -362,9 +363,9 @@ async def fetch_azure_response_stream(client, url, headers, payload):
                         yield "data: " + json.dumps(line).strip() + end_of_line
     yield "data: [DONE]" + end_of_line
 
-async def fetch_cloudflare_response_stream(client, url, headers, payload, model):
+async def fetch_cloudflare_response_stream(client, url, headers, payload, model, timeout):
     timestamp = int(datetime.timestamp(datetime.now()))
-    async with client.stream('POST', url, headers=headers, json=payload) as response:
+    async with client.stream('POST', url, headers=headers, json=payload, timeout=timeout) as response:
         error_message = await check_response(response, "fetch_cloudflare_response_stream")
         if error_message:
             yield error_message
@@ -380,16 +381,16 @@ 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)
                         yield sse_string
     yield "data: [DONE]" + end_of_line
 
-async def fetch_cohere_response_stream(client, url, headers, payload, model):
+async def fetch_cohere_response_stream(client, url, headers, payload, model, timeout):
     timestamp = int(datetime.timestamp(datetime.now()))
-    async with client.stream('POST', url, headers=headers, json=payload) as response:
+    async with client.stream('POST', url, headers=headers, json=payload, timeout=timeout) as response:
         error_message = await check_response(response, "fetch_cohere_response_stream")
         if error_message:
             yield error_message
@@ -401,7 +402,7 @@ async def fetch_cohere_response_stream(client, url, headers, payload, model):
             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":
@@ -410,9 +411,9 @@ async def fetch_cohere_response_stream(client, url, headers, payload, model):
                     yield sse_string
     yield "data: [DONE]" + end_of_line
 
-async def fetch_claude_response_stream(client, url, headers, payload, model):
+async def fetch_claude_response_stream(client, url, headers, payload, model, timeout):
     timestamp = int(datetime.timestamp(datetime.now()))
-    async with client.stream('POST', url, headers=headers, json=payload) as response:
+    async with client.stream('POST', url, headers=headers, json=payload, timeout=timeout) as response:
         error_message = await check_response(response, "fetch_claude_response_stream")
         if error_message:
             yield error_message
@@ -427,7 +428,7 @@ async def fetch_claude_response_stream(client, url, headers, payload, model):
                 # 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)
@@ -463,9 +464,9 @@ async def fetch_claude_response_stream(client, url, headers, payload, model):
 
     yield "data: [DONE]" + end_of_line
 
-async def fetch_aws_response_stream(client, url, headers, payload, model):
+async def fetch_aws_response_stream(client, url, headers, payload, model, timeout):
     timestamp = int(datetime.timestamp(datetime.now()))
-    async with client.stream('POST', url, headers=headers, json=payload) as response:
+    async with client.stream('POST', url, headers=headers, json=payload, timeout=timeout) as response:
         error_message = await check_response(response, "fetch_aws_response_stream")
         if error_message:
             yield error_message
@@ -486,7 +487,7 @@ async def fetch_aws_response_stream(client, url, headers, payload, model):
                 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):
                     # 解码 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,13 +515,13 @@ async def fetch_aws_response_stream(client, url, headers, payload, model):
 
     yield "data: [DONE]" + end_of_line
 
-async def fetch_response(client, url, headers, payload, engine, model):
+async def fetch_response(client, url, headers, payload, engine, model, timeout=200):
     response = None
     if payload.get("file"):
         file = payload.pop("file")
-        response = await client.post(url, headers=headers, data=payload, files={"file": file})
+        response = await client.post(url, headers=headers, data=payload, files={"file": file}, timeout=timeout)
     else:
-        response = await client.post(url, headers=headers, json=payload)
+        response = await client.post(url, headers=headers, json=payload, timeout=timeout)
     error_message = await check_response(response, "fetch_response")
     if error_message:
         yield error_message
@@ -530,7 +531,8 @@ async def fetch_response(client, url, headers, payload, engine, model):
         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):
         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):
         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,34 +622,36 @@ async def fetch_response(client, url, headers, payload, engine, model):
         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):
+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):
+        async for chunk in fetch_gemini_response_stream(client, url, headers, payload, model, timeout):
             yield chunk
     elif engine == "claude" or engine == "vertex-claude":
-        async for chunk in fetch_claude_response_stream(client, url, headers, payload, model):
+        async for chunk in fetch_claude_response_stream(client, url, headers, payload, model, timeout):
             yield chunk
     elif engine == "aws":
-        async for chunk in fetch_aws_response_stream(client, url, headers, payload, model):
+        async for chunk in fetch_aws_response_stream(client, url, headers, payload, model, timeout):
             yield chunk
     elif engine == "gpt" or engine == "openrouter" or engine == "azure-databricks":
-        async for chunk in fetch_gpt_response_stream(client, url, headers, payload):
+        async for chunk in fetch_gpt_response_stream(client, url, headers, payload, timeout):
             yield chunk
     elif engine == "azure":
-        async for chunk in fetch_azure_response_stream(client, url, headers, payload):
+        async for chunk in fetch_azure_response_stream(client, url, headers, payload, timeout):
             yield chunk
     elif engine == "cloudflare":
-        async for chunk in fetch_cloudflare_response_stream(client, url, headers, payload, model):
+        async for chunk in fetch_cloudflare_response_stream(client, url, headers, payload, model, timeout):
             yield chunk
     elif engine == "cohere":
-        async for chunk in fetch_cohere_response_stream(client, url, headers, payload, model):
+        async for chunk in fetch_cohere_response_stream(client, url, headers, payload, model, timeout):
             yield chunk
     else:
-        raise ValueError("Unknown response")
+        raise ValueError("Unknown response")

beswarm/aient/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):

beswarm/taskmanager.py

@@ -23,6 +23,7 @@ class TaskManager:
     它管理任务的生命周期，并通过一个固定大小的工作者池来控制并发执行的任务数量。
     """
     def __init__(self, concurrency_limit=None):
+        self.raw_concurrency_limit = concurrency_limit
         self.concurrency_limit = concurrency_limit or int(os.getenv("BESWARM_CONCURRENCY_LIMIT", "3"))
 
         if self.concurrency_limit <= 0:
@@ -38,8 +39,6 @@ class TaskManager:
         self.cache_dir = None
         self.task_cache_file = None
 
-        print(f"TaskManager 初始化，并发限制为: {self.concurrency_limit}")
-
     def set_root_path(self, root_path):
         """设置工作根目录并加载持久化的任务状态。"""
         if self.root_path is not None:
@@ -52,6 +51,9 @@ class TaskManager:
         self._load_tasks_from_cache()
         self.set_task_cache("root_path", str(self.root_path))
 
+        if not self.raw_concurrency_limit:
+            self.concurrency_limit = int(os.getenv("BESWARM_CONCURRENCY_LIMIT", "3"))
+
         # 启动工作者池
         self.start()
         # 恢复中断的任务
@@ -86,7 +88,6 @@ class TaskManager:
 
     async def _worker_loop(self, worker_name: str):
         """每个工作者的主循环，从队列中拉取并执行任务。"""
-        print(f"[{worker_name}] 已就绪，等待任务...")
         while self._is_running:
             try:
                 task_id, coro = await self._pending_queue.get()

beswarm/tools/completion.py

@@ -14,4 +14,4 @@ def task_complete(message: str) -> str:
     Returns:
         str: 传入的任务完成信息。
     """
-    return "<task_complete_message>" + message + "</task_complete_message>"
+    return message

{beswarm-0.2.59.dist-info → beswarm-0.2.61.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: beswarm
-Version: 0.2.59
+Version: 0.2.61
 Summary: MAS
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown

{beswarm-0.2.59.dist-info → beswarm-0.2.61.dist-info}/RECORD

@@ -3,16 +3,16 @@ beswarm/broker.py,sha256=64Y-djrKYaZfBQ8obwHOmr921QgZeu9BtScZWaYLfDo,9887
 beswarm/core.py,sha256=htssaaeIBZ_yOqvX9VtANoVWaZHt_7oWcxyDI1z0paQ,310
 beswarm/knowledge_graph.py,sha256=oiOMknAJzGrOHc2AyQgvrCcZAkGLhFnsnvSBdfFBWMw,14831
 beswarm/prompt.py,sha256=INVRWQZP6lysvGUcPOYI_er5-bi1gGe_qa6BTov7PmY,32362
-beswarm/taskmanager.py,sha256=ErgZa9_aBeWdmt5neRw6sDVdwSWewxip458gAjeQhS4,12188
+beswarm/taskmanager.py,sha256=vMmcoZ4FlNvjEliRkv3AniPji50NcY4Q1_2HETzR0DU,12226
 beswarm/utils.py,sha256=0J-b38P5QGT-A_38co7FjzaUNJykaskI7mbbcQ4w_68,8215
 beswarm/agents/chatgroup.py,sha256=PzrmRcDKAbB7cxL16nMod_CzPosDV6bfTmXxQVuv-AQ,12012
-beswarm/agents/planact.py,sha256=ndsisPbGhHN04O85Htc12bC1UWPdgq_CNA1077nxeD0,20147
+beswarm/agents/planact.py,sha256=y6NtiiQC1SIkJCWeK4REoR8yRoMllRda5aiVeXhDZdE,20082
 beswarm/aient/aient/__init__.py,sha256=SRfF7oDVlOOAi6nGKiJIUK6B_arqYLO9iSMp-2IZZps,21
 beswarm/aient/aient/core/__init__.py,sha256=NxjebTlku35S4Dzr16rdSqSTWUvvwEeACe8KvHJnjPg,34
 beswarm/aient/aient/core/log_config.py,sha256=kz2_yJv1p-o3lUQOwA3qh-LSc3wMHv13iCQclw44W9c,274
-beswarm/aient/aient/core/models.py,sha256=d4MISNezTSe0ls0-fjuToI2SoT-sk5fWqAJuKVinIlo,7502
-beswarm/aient/aient/core/request.py,sha256=4FFCwQ7h7b6bqtrA8qw-DPJVXZTj2i1CkYccFeEwUPw,76552
-beswarm/aient/aient/core/response.py,sha256=tYKWOeexYEhWK4napIAfYjCn2rQ1zpxRvknsBBWrv2M,33206
+beswarm/aient/aient/core/models.py,sha256=KMlCRLjtq1wQHZTJGqnbWhPS2cHq6eLdnk7peKDrzR8,7490
+beswarm/aient/aient/core/request.py,sha256=vfwi3ZGYp2hQzSJ6mPXJVgcV_uu5AJ_NAL84mLfF8WA,76674
+beswarm/aient/aient/core/response.py,sha256=vQFuc3amHiD1hv_OiINRJnh33n79PnbdzMSBSRlqR5E,34309
 beswarm/aient/aient/core/utils.py,sha256=D98d5Cy1h4ejKtuxS0EEDtL4YqpaZLB5tuXoVP0IBWQ,28462
 beswarm/aient/aient/core/test/test_base_api.py,sha256=pWnycRJbuPSXKKU9AQjWrMAX1wiLC_014Qc9hh5C2Pw,524
 beswarm/aient/aient/core/test/test_geminimask.py,sha256=HFX8jDbNg_FjjgPNxfYaR-0-roUrOO-ND-FVsuxSoiw,13254
@@ -21,7 +21,7 @@ beswarm/aient/aient/core/test/test_payload.py,sha256=8jBiJY1uidm1jzL-EiK0s6UGmW9
 beswarm/aient/aient/models/__init__.py,sha256=ZTiZgbfBPTjIPSKURE7t6hlFBVLRS9lluGbmqc1WjxQ,43
 beswarm/aient/aient/models/audio.py,sha256=kRd-8-WXzv4vwvsTGwnstK-WR8--vr9CdfCZzu8y9LA,1934
 beswarm/aient/aient/models/base.py,sha256=-nnihYnx-vHZMqeVO9ljjt3k4FcD3n-iMk4tT-10nRQ,7232
-beswarm/aient/aient/models/chatgpt.py,sha256=S_6s_3epO9UPWea9zyf4-cy8ekMZXNct3AuvSJtO9Pg,43056
+beswarm/aient/aient/models/chatgpt.py,sha256=q62B6cbtHqKrqsQjM24k_1wi_5-UiuxkXa7e2yG_Clg,44661
 beswarm/aient/aient/plugins/__init__.py,sha256=p3KO6Aa3Lupos4i2SjzLQw1hzQTigOAfEHngsldrsyk,986
 beswarm/aient/aient/plugins/arXiv.py,sha256=yHjb6PS3GUWazpOYRMKMzghKJlxnZ5TX8z9F6UtUVow,1461
 beswarm/aient/aient/plugins/config.py,sha256=TGgZ5SnNKZ8MmdznrZ-TEq7s2ulhAAwTSKH89bci3dA,7079
@@ -104,7 +104,7 @@ beswarm/queries/tree-sitter-languages/scala-tags.scm,sha256=UxQjz80JIrrJ7Pm56uUn
 beswarm/queries/tree-sitter-languages/typescript-tags.scm,sha256=OMdCeedPiA24ky82DpgTMKXK_l2ySTuF2zrQ2fJAi9E,1253
 beswarm/tools/__init__.py,sha256=CPFj04Lm6TEol_2BFX-mVrpzvVIsmPcfTWhRg3xex7Q,2077
 beswarm/tools/click.py,sha256=wu-Ov5U2ZZLcU0gennDVh_2w_Td7F4dbVJcmu_dfHV4,20872
-beswarm/tools/completion.py,sha256=-ozGZoHDBB50bbdnwenb7p_KRmi3iXmS7z6Wfu85lBo,601
+beswarm/tools/completion.py,sha256=AIHtEHfSp6fs4Xa_nTcw9fLgXtYHtYGRDTSy7_x_Fak,544
 beswarm/tools/edit_file.py,sha256=ZTJvbpsfRlp2t98kTn9XQ5qZBTdsWJVWv9t0lvK4RfU,9147
 beswarm/tools/graph.py,sha256=IJ-wExgSEDLGYXTRtCtyxMM2wkA1gg40Z-KI3iyQThE,5084
 beswarm/tools/planner.py,sha256=vsHd7rE8RQHJrZ7BQ0ZXhbt4Fjh3DeyxU4piA5R-VPM,1253
@@ -116,7 +116,7 @@ beswarm/tools/search_web.py,sha256=NYrb5KL_WUGPm-fOKT8Cyjon04lxBU-gaLdrVjeYgGo,1
 beswarm/tools/subtasks.py,sha256=mIjA2QrRy9Fos4rYm8fCfu2QrsE_MGnQI9IR8dOxsGs,9885
 beswarm/tools/worker.py,sha256=_cSkRUKRJMAiZiTfnBze_e9Kc7k7KvbB5hdxdvp4FW4,2009
 beswarm/tools/write_csv.py,sha256=u0Hq18Ksfheb52MVtyLNCnSDHibITpsYBPs2ub7USYA,1466
-beswarm-0.2.59.dist-info/METADATA,sha256=4rzgb1d_Aoyfl6zco7t0UHm7OuC3TFEcXIA0mtcXb9U,3878
-beswarm-0.2.59.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-beswarm-0.2.59.dist-info/top_level.txt,sha256=pJw4O87wvt5882smuSO6DfByJz7FJ8SxxT8h9fHCmpo,8
-beswarm-0.2.59.dist-info/RECORD,,
+beswarm-0.2.61.dist-info/METADATA,sha256=s14zUG7R6_2LWhXRMX-UZF0MQeh06d0bmX84pH5KlhU,3878
+beswarm-0.2.61.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+beswarm-0.2.61.dist-info/top_level.txt,sha256=pJw4O87wvt5882smuSO6DfByJz7FJ8SxxT8h9fHCmpo,8
+beswarm-0.2.61.dist-info/RECORD,,