beswarm 0.2.39__py3-none-any.whl → 0.2.41__py3-none-any.whl

beswarm/taskmanager.py

@@ -17,210 +17,257 @@ class TaskStatus(Enum):
 
 
 class TaskManager:
-    """一个简单的异步任务管理器"""
-    def __init__(self):
-        self.tasks = {}  # 使用字典来存储任务，key是task_id, value是task对象
-        self.results_queue = asyncio.Queue()
+    """
+    一个带并发控制的异步任务管理器。
+    它管理任务的生命周期，并通过一个固定大小的工作者池来控制并发执行的任务数量。
+    """
+    def __init__(self, concurrency_limit=3):
+        if concurrency_limit <= 0:
+            raise ValueError("并发限制必须大于0")
+
+        self.concurrency_limit = concurrency_limit
+        self.tasks_cache = {}          # 存储所有任务的状态和元数据, key: task_id
+
+        self._pending_queue = asyncio.Queue() # 内部待办任务队列
+        self._results_queue = asyncio.Queue() # 内部已完成任务结果队列
+        self._workers = []                    # 持有工作者任务的引用
+        self._is_running = False              # 标记工作者池是否在运行
         self.root_path = None
-        self.tasks_cache = {}
+        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:
+        """设置工作根目录并加载持久化的任务状态。"""
+        if self.root_path is not None:
             return
         self.root_path = Path(root_path)
         self.cache_dir = self.root_path / ".beswarm"
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
         self.task_cache_file = self.cache_dir / "tasks.json"
-        self.task_cache_file.touch(exist_ok=True)
-        self.read_tasks_cache()
+
+        self._load_tasks_from_cache()
         self.set_task_cache("root_path", str(self.root_path))
-        self.resume_all_running_task()
 
-    def set_task_cache(self, *keys_and_value):
-        """
-        设置可嵌套的任务缓存。
-        接受无限个键和一个值，例如 set_task_cache('a', 'b', 'c', value)
-        会转换为 tasks_cache['a']['b']['c'] = value
-        """
-        if len(keys_and_value) < 2:
-            return  # 至少需要一个键和一个值
+        # 启动工作者池
+        self.start()
+        # 恢复中断的任务
+        self.resume_interrupted_tasks()
+
+    def start(self):
+        """启动并发工作者池。"""
+        if self._is_running:
+            return
+
+        self._is_running = True
+        for i in range(self.concurrency_limit):
+            worker = asyncio.create_task(self._worker_loop(f"Worker-{i+1}"))
+            self._workers.append(worker)
+        print(f"已启动 {self.concurrency_limit} 个并发工作者。")
+
+    async def stop(self):
+        """优雅地停止所有工作者。"""
+        if not self._is_running:
+            return
+
+        print("\n正在停止 TaskManager...")
+        await self._pending_queue.join()
+
+        for worker in self._workers:
+            worker.cancel()
+
+        await asyncio.gather(*self._workers, return_exceptions=True)
+
+        self._is_running = False
+        print("所有工作者已停止。")
+
+    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()
+
+                print(f"[{worker_name}] 领到了任务 <{task_id[:8]}>，开始执行...")
+                self._update_task_status(task_id, TaskStatus.RUNNING)
 
-        keys = keys_and_value[:-1]
-        value = keys_and_value[-1]
+                try:
+                    result = await coro
+                    self._handle_task_completion(task_id, TaskStatus.DONE, result)
+                except Exception as e:
+                    self._handle_task_completion(task_id, TaskStatus.ERROR, e)
+                finally:
+                    self._pending_queue.task_done()
 
+            except asyncio.CancelledError:
+                print(f"[{worker_name}] 被取消，正在退出...")
+                break
+            except Exception as e:
+                print(f"[{worker_name}] 循环中遇到严重错误: {e}")
+
+    def _handle_task_completion(self, task_id, status, result):
+        """统一处理任务完成的内部函数。"""
+        if status == TaskStatus.DONE:
+            print(f"✅ 任务 <{task_id[:8]}> 执行成功。")
+        else: # ERROR
+            print(f"❌ 任务 <{task_id[:8]}> 执行失败: {result}")
+
+        self._update_task_status(task_id, status, result=str(result))
+        self._results_queue.put_nowait((task_id, status, result))
+
+    def set_task_cache(self, *keys_and_value):
+        """设置可嵌套的任务缓存。"""
+        if len(keys_and_value) < 2: return
+        keys, value = keys_and_value[:-1], keys_and_value[-1]
         d = self.tasks_cache
-        # 遍历到倒数第二个键，确保路径存在
         for key in keys[:-1]:
             d = d.setdefault(key, {})
-
-        # 在最后一个键上设置值
         d[keys[-1]] = value
-        self.save_tasks_cache()
+        self._save_tasks_to_cache()
 
-    def save_tasks_cache(self):
-        self.task_cache_file.write_text(json.dumps(self.tasks_cache, ensure_ascii=False, indent=4), encoding="utf-8")
-
-    def read_tasks_cache(self):
-        content = self.task_cache_file.read_text(encoding="utf-8")
+    def _save_tasks_to_cache(self):
+        """将任务缓存持久化到文件。"""
+        if not self.task_cache_file: return
         try:
-            self.tasks_cache = json.loads(content) if content else {}
-        except json.JSONDecodeError:
-            raise ValueError("任务缓存文件格式错误")
+            with self.task_cache_file.open('w', encoding='utf-8') as f:
+                json.dump(self.tasks_cache, f, indent=4, ensure_ascii=False)
+        except Exception as e:
+            print(f"警告：无法将任务状态持久化到文件: {e}")
 
-    def create_tasks(self, task_coro, tasks_params):
-        """
-        批量创建并注册任务。
+    def _load_tasks_from_cache(self):
+        """从文件加载任务缓存。"""
+        if not self.task_cache_file or not self.task_cache_file.exists():
+            self.tasks_cache = {}
+            return
+        try:
+            content = self.task_cache_file.read_text(encoding='utf-8')
+            if content:
+                self.tasks_cache = json.loads(content)
+            else:
+                self.tasks_cache = {}
+        except (FileNotFoundError, json.JSONDecodeError):
+            self.tasks_cache = {}
+            print("警告：任务缓存文件不存在或格式错误，将使用空缓存。")
 
-        Args:
-            task_coro: 用于创建任务的协程函数。
-            tasks_params (list): 包含任务参数的列表。
+    async def get_next_result(self):
+        """异步获取下一个完成的任务结果。"""
+        return await self._results_queue.get()
 
-        Returns:
-            list: 创建的任务ID列表。
+    def create_tasks_batch(self, task_coro_func, tasks_params_list):
+        """
+        批量创建任务，但不是立即执行，而是将它们放入待处理队列。
         """
+        if not self._is_running:
+            raise RuntimeError("TaskManager尚未启动。请先调用 start() 方法。")
+
         task_ids = []
-        for args in tasks_params:
-            coro = task_coro(**args)
-            task_id = self.create_task(coro)
+        for params in tasks_params_list:
+            task_id = str(uuid.uuid4())
+            coro = task_coro_func(**params)
+
+            # 初始化任务状态为 PENDING
+            self._update_task_status(task_id, TaskStatus.PENDING, args=params)
+
+            # 将任务定义放入队列
+            self._pending_queue.put_nowait((task_id, coro))
             task_ids.append(task_id)
-            self.set_task_cache(task_id, "args", args)
-            self.set_task_cache(task_id, "status", TaskStatus.RUNNING.value)
+
+        print(f"已将 {len(task_ids)} 个新任务加入待处理队列。队列当前大小: {self._pending_queue.qsize()}")
         return task_ids
 
-    def resume_all_running_task(self):
-        running_task_id_list = [task_id for task_id, task in self.tasks_cache.items() if task_id != "root_path" and task.get("status") == "RUNNING"]
-        for task_id in running_task_id_list:
-            tasks_params = self.tasks_cache[task_id]["args"]
-            task_id = self.resume_task(task_id, registry.tools["worker"], tasks_params)
+    def create_tasks(self, task_coro_func, tasks_params_list):
+        """批量将任务放入待处理队列。"""
+        if not self._is_running:
+            raise RuntimeError("TaskManager尚未启动。请先在 set_root_path 后确保其已启动。")
 
-    def resume_task(self, task_id, task_coro, args):
-        """
-        恢复一个任务。
-        """
-        task = self.tasks_cache.get(task_id)
-        if not task:
-            return TaskStatus.NOT_FOUND
+        task_ids = []
+        for params in tasks_params_list:
+            task_id = str(uuid.uuid4())
+            coro = task_coro_func(**params)
 
-        coro = task_coro(**args)
-        task_id = self.create_task(coro, task_id)
-        self.set_task_cache(task_id, "args", args)
-        self.set_task_cache(task_id, "status", TaskStatus.RUNNING.value)
-        print(f"任务已恢复: ID={task_id}, Name={task_id}")
-        print(f"args: {args}")
-        print(f"self.tasks_cache: {json.dumps(self.tasks_cache, ensure_ascii=False, indent=4)}")
-        return task_id
+            self._update_task_status(task_id, TaskStatus.PENDING, args=params)
+            self._pending_queue.put_nowait((task_id, coro))
+            task_ids.append(task_id)
 
-    def create_task(self, coro, task_id=None):
-        """
-        创建并注册一个新任务。
+        print(f"已将 {len(task_ids)} 个新任务加入待处理队列。队列当前大小: {self._pending_queue.qsize()}")
+        return task_ids
 
-        Args:
-            coro: 要执行的协程。
-            name (str, optional): 任务的可读名称。 Defaults to None.
+    def resume_interrupted_tasks(self):
+        """在启动时，恢复所有处于 PENDING 或 RUNNING 状态的旧任务。"""
+        interrupted_tasks = [
+            (tid, info) for tid, info in self.tasks_cache.items()
+            if tid != "root_path" and info.get("status") in [TaskStatus.PENDING.value, TaskStatus.RUNNING.value]
+        ]
 
-        Returns:
-            str: 任务的唯一ID。
-        """
-        if task_id == None:
-            task_id = str(uuid.uuid4())
-        task_name = f"Task-{task_id[:8]}"
+        if not interrupted_tasks:
+            return
 
-        # 使用 asyncio.create_task() 创建任务
-        task = asyncio.create_task(coro, name=task_name)
+        print(f"检测到 {len(interrupted_tasks)} 个中断的任务，正在恢复...")
+        worker_fun = registry.tools["worker"]
 
-        # 将任务存储在管理器中
-        # 当任务完成时，通过回调函数将结果放入队列
-        task.add_done_callback(
-            lambda t: self._on_task_done(task_id, t)
-        )
-        self.tasks[task_id] = task
-        print(f"任务已创建: ID={task_id}, Name={task_name}")
-        return task_id
+        for task_id, task_info in interrupted_tasks:
+            args = task_info.get("args")
+            if not args:
+                print(f"警告：任务 <{task_id[:8]}> 缺少参数，无法恢复。")
+                self._update_task_status(task_id, TaskStatus.ERROR, result="缺少参数，无法恢复")
+                continue
 
-    def get_task_status(self, task_id):
-        """
-        查询特定任务的状态。
+            coro = worker_fun(**args)
+            self._update_task_status(task_id, TaskStatus.PENDING)
+            self._pending_queue.put_nowait((task_id, coro))
 
-        Args:
-            task_id (str): 要查询的任务ID。
+        print(f"{len(interrupted_tasks)} 个中断的任务已重新加入队列。")
 
-        Returns:
-            TaskStatus: 任务的当前状态。
-        """
-        task = self.tasks.get(task_id)
-        if not task:
-            return TaskStatus.NOT_FOUND
+    def resume_task(self, task_id, goal):
+        """恢复一个指定的任务，实质上是创建一个新任务并替换旧的记录，但ID保持不变。"""
+        if task_id not in self.tasks_cache:
+            return f"任务 {task_id} 不存在"
 
-        if task.done():
-            if task.cancelled():
-                return TaskStatus.CANCELLED
-            elif task.exception() is not None:
-                return TaskStatus.ERROR
-            else:
-                return TaskStatus.DONE
+        old_task_info = self.tasks_cache.get(task_id, {})
+        tasks_params = old_task_info.get("args", {})
+        if not tasks_params:
+             return f"<tool_error>任务 {task_id} 缺少参数信息，无法恢复。</tool_error>"
 
-        # asyncio.Task 没有直接的 'RUNNING' 状态。
-        # 如果任务还没有完成，它要么是等待执行（PENDING），要么是正在执行（RUNNING）。
-        # 这里我们简化处理，认为未完成的就是运行中。
-        return TaskStatus.RUNNING
+        tasks_params["goal"] = goal
+        tasks_params["cache_messages"] = True # 恢复时强制使用缓存
 
-    def get_task_result(self, task_id):
-        """获取已完成任务的结果，如果任务未完成或出错则返回相应信息。"""
-        task = self.tasks.get(task_id)
-        if self.get_task_status(task_id) == TaskStatus.DONE:
-            return task.result()
-        elif self.get_task_status(task_id) == TaskStatus.ERROR:
-            return task.exception()
-        return None
-
-    def _on_task_done(self, task_id, task):
-        """私有回调函数，在任务完成时将结果放入队列。"""
-        try:
-            # 将元组 (task_id, status, result) 放入队列
-            self.results_queue.put_nowait(
-                (task_id, TaskStatus.DONE, task.result())
-            )
-            self.set_task_cache(task_id, "status", TaskStatus.DONE.value)
-            self.set_task_cache(task_id, "result", task.result())
-        except asyncio.CancelledError:
-            self.results_queue.put_nowait(
-                (task_id, TaskStatus.CANCELLED, None)
-            )
-            self.set_task_cache(task_id, "status", TaskStatus.CANCELLED.value)
-        except Exception as e:
-            self.results_queue.put_nowait(
-                (task_id, TaskStatus.ERROR, e)
-            )
-            self.set_task_cache(task_id, "status", TaskStatus.ERROR.value)
-            self.set_task_cache(task_id, "result", str(e))
+        worker_fun = registry.tools["worker"]
+        coro = worker_fun(**tasks_params)
 
-    async def get_next_result(self):
-        """
-        等待并返回下一个完成的任务结果。
+        self._update_task_status(task_id, TaskStatus.PENDING, args=tasks_params)
+        self._pending_queue.put_nowait((task_id, coro))
 
-        如果所有任务都已提交，但没有任务完成，此方法将异步等待。
+        print(f"任务 <{task_id[:8]}> 已被重新加入队列等待恢复执行。")
+        return f"任务 {task_id} 已恢复"
 
-        Returns:
-            tuple: 一个包含 (task_id, status, result) 的元组。
-        """
-        return await self.results_queue.get()
+    def _update_task_status(self, task_id, status: TaskStatus, args=None, result=None):
+        """统一更新任务状态缓存并持久化。"""
+        if task_id not in self.tasks_cache:
+            self.tasks_cache[task_id] = {}
 
-    def get_task_index(self, task_id):
-        """
-        获取任务在任务字典中的插入顺序索引。
+        current_task = self.tasks_cache[task_id]
+        current_task['status'] = status.value
+        if args is not None:
+            current_task['args'] = args
+        if result is not None:
+            current_task['result'] = result
 
-        Args:
-            task_id (str): 要查询的任务ID。
+        self._save_tasks_to_cache()
 
-        Returns:
-            int: 任务的索引（从0开始），如果未找到则返回-1。
-        """
-        try:
-            # 将字典的键转换为列表并查找索引
-            task_ids_list = list(self.tasks.keys())
-            return task_ids_list.index(task_id)
-        except ValueError:
-            # 如果任务ID不存在，则返回-1
-            return -1
+    def get_task_status(self, task_id):
+        """查询特定任务的状态。"""
+        task_info = self.tasks_cache.get(task_id)
+        if not task_info:
+            return TaskStatus.NOT_FOUND
+        return TaskStatus(task_info.get("status", "NOT_FOUND"))
+
+    def get_task_result(self, task_id):
+        """获取已完成任务的结果。"""
+        task_info = self.tasks_cache.get(task_id)
+        if not task_info or task_info.get("status") not in [TaskStatus.DONE.value, TaskStatus.ERROR.value]:
+            return None
+        return task_info.get("result")
 
 async def main():
     manager = TaskManager()

beswarm/tools/__init__.py

@@ -3,11 +3,12 @@ from .search_web import search_web
 from .completion import task_complete
 from .search_arxiv import search_arxiv
 from .repomap import get_code_repo_map
+from .write_csv import append_row_to_csv
 from .request_input import request_admin_input
 from .screenshot import save_screenshot_to_file
 from .worker import worker, worker_gen, chatgroup
 from .click import find_and_click_element, scroll_screen
-from .subtasks import create_task, resume_task, get_all_tasks_status, get_task_result
+from .subtasks import create_task, resume_task, get_all_tasks_status, get_task_result, create_tasks_from_csv
 
 #显式导入 aient.plugins 中的所需内容
 from ..aient.src.aient.plugins import (
@@ -47,12 +48,14 @@ __all__ = [
     "list_directory",
     "get_task_result",
     "get_url_content",
+    "append_row_to_csv",
     "set_readonly_path",
     "get_code_repo_map",
     "run_python_script",
     "get_search_results",
     "request_admin_input",
     "get_all_tasks_status",
+    "create_tasks_from_csv",
     "find_and_click_element",
     "download_read_arxiv_pdf",
     "save_screenshot_to_file",

beswarm/tools/completion.py

@@ -9,7 +9,7 @@ def task_complete(message: str) -> str:
     它标志着一个任务的成功结束，并将最终的输出传递给用户或调用者。
 
     Args:
-        message (str): 任务完成的信息或最终结果。
+        message (str): 任务完成的信息或最终结果。必填字段。
 
     Returns:
         str: 传入的任务完成信息。

beswarm/tools/search_arxiv.py

@@ -1,4 +1,5 @@
 import requests
+import csv
 from datetime import datetime
 from ..aient.src.aient.plugins import register_tool
 
@@ -43,10 +44,11 @@ NoProp: Training Neural Networks without Back-propagation or Forward-propagation
     包含搜索结果的字典列表，每个字典包含论文的标题、作者、摘要、发布日期、最后更新日期、arXiv ID、类别和PDF链接等信息
     """
     try:
-        base_url = "http://export.arxiv.org/api/query"
+        base_url = "https://export.arxiv.org/api/query"
 
         # 构建查询参数
-        search_query = f"all:{query}"
+        search_query = query
+        # search_query = f"all:{query}"
 
         # 添加类别过滤
         if categories:
@@ -64,6 +66,8 @@ NoProp: Training Neural Networks without Back-propagation or Forward-propagation
             else:
                 search_query += f" AND au:\"{authors}\""
 
+        print(search_query)
+
         # 添加日期过滤
         # arXiv API不直接支持日期范围过滤，需要在结果中过滤
 
@@ -130,6 +134,27 @@ NoProp: Training Neural Networks without Back-propagation or Forward-propagation
                 if cat_term not in categories_list:
                     categories_list.append(cat_term)
 
+            # 应用严格的类别过滤，确保论文的所有类别都符合用户的要求
+            if categories:
+                user_specified_categories = categories if isinstance(categories, list) else [categories]
+
+                allowed_prefixes = []
+                for pattern in user_specified_categories:
+                    if pattern.endswith('*'):
+                        allowed_prefixes.append(pattern[:-1])
+                    else:
+                        allowed_prefixes.append(pattern)
+
+                all_paper_categories_match = True
+                for paper_cat in categories_list:
+                    # 检查当前论文的每个分类是否至少匹配一个用户指定的模式前缀
+                    if not any(paper_cat.startswith(prefix) for prefix in allowed_prefixes):
+                        all_paper_categories_match = False
+                        break
+
+                if not all_paper_categories_match:
+                    continue  # 如果有任何一个分类不匹配，就跳过这篇论文
+
             # 获取摘要
             abstract = ""
             if include_abstract:
@@ -167,9 +192,11 @@ if __name__ == '__main__':
     # python -m beswarm.tools.search_arxiv
     test_query = "NoProp"
     test_query = '"Attention Is All You Need"'
+    test_query = '(all:"sparse autoencoders" OR all:"sparse autoencoder" OR (all:SAE AND NOT au:SAE))'
+
     print(f"使用关键词 '{test_query}' 测试搜索...")
 
-    search_results = search_arxiv(query=test_query, max_results=50, sort_by='lastUpdatedDate')
+    search_results = search_arxiv(query=test_query, max_results=1000, categories='cs*', sort_by='lastUpdatedDate')
 
     if isinstance(search_results, str):
         # 如果返回的是错误信息字符串，则打印错误
@@ -183,9 +210,37 @@ if __name__ == '__main__':
                 print(f"  作者: {', '.join(paper['authors'])}")
                 print(f"  发布日期: {paper['published_date']}")
                 print(f"  arXiv ID: {paper['arxiv_id']}")
+                print(f"  领域: {paper['categories']}")
                 print(f"  PDF链接: {paper['pdf_url']}")
                 print(f"  摘要: {paper['abstract'][:150]}...") # 打印摘要前150个字符
                 print("-" * 20)
+
+            # 将结果保存到CSV文件
+            csv_filename = 'arxiv_search_results.csv'
+            print(f"\n正在将 {len(search_results)} 条结果保存到 {csv_filename}...")
+
+            try:
+                with open(csv_filename, mode='w', newline='', encoding='utf-8') as csv_file:
+                    # 使用第一个数据项的键作为CSV文件的标题
+                    fieldnames = search_results[0].keys()
+                    writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
+
+                    writer.writeheader()
+                    for paper in search_results:
+                        # 转换列表为字符串以便写入CSV
+                        paper_for_csv = paper.copy()
+                        if 'authors' in paper_for_csv and isinstance(paper_for_csv['authors'], list):
+                            paper_for_csv['authors'] = ', '.join(paper_for_csv['authors'])
+                        if 'categories' in paper_for_csv and isinstance(paper_for_csv['categories'], list):
+                            paper_for_csv['categories'] = ', '.join(paper_for_csv['categories'])
+
+                        writer.writerow(paper_for_csv)
+
+                print(f"结果已成功保存到 {csv_filename}")
+
+            except IOError as e:
+                print(f"错误：无法写入文件 {csv_filename}: {e}")
+
         else:
             print("未找到相关论文。")
     else: