refactor(aibox): 抽离工具类与子系统,附单测基线
A 批(tools/basic_tools.py):迁移 7 个无外部依赖工具类(FileToolBase/ReadFileTool/WriteFileTool/ListFilesTool/SimpleWebCrawlerTool/GetCurrentTimeTool/CalculatorTool) Wave1(subsystems/):迁移 5 个叶子子系统(CommandExecutor/SubprocessCommandExecutor/PromptManager/MemoryDatabase/MemoDatabase/SkillManager),aibox.py 减少 ~1440 行 新增 tests/ 单测 113 个用例全绿,依赖方向维持 common -> subsystems -> tools -> aibox 单向
This commit is contained in:
1 parent
8af8dec1b0
commit
eadba81c26
15 files changed
+3216
-1883
No files matched your search
+12
-1880
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,30 @@
|
||||
"""
|
||||
subsystems - aibox 子系统模块(包)
|
||||
|
||||
抽离自 aibox.py 的基础设施子系统,按职责分文件。
|
||||
当前包含(Wave 1):
|
||||
- executors.py : CommandExecutor + SubprocessCommandExecutor
|
||||
- prompt_manager.py : PromptManager(单例)
|
||||
- db.py : MemoryDatabase + MemoDatabase
|
||||
- skill_manager.py : SkillManager
|
||||
|
||||
依赖方向:common <- subsystems <- aibox
|
||||
不反向依赖 aibox 任何符号。
|
||||
|
||||
采用"重导出"策略:外部统一用 `from subsystems import MemoryDatabase`,
|
||||
调用点无需关心具体子模块所在文件。
|
||||
"""
|
||||
|
||||
from .executors import CommandExecutor, SubprocessCommandExecutor
|
||||
from .prompt_manager import PromptManager
|
||||
from .db import MemoryDatabase, MemoDatabase
|
||||
from .skill_manager import SkillManager
|
||||
|
||||
__all__ = [
|
||||
"CommandExecutor",
|
||||
"SubprocessCommandExecutor",
|
||||
"PromptManager",
|
||||
"MemoryDatabase",
|
||||
"MemoDatabase",
|
||||
"SkillManager",
|
||||
]
|
||||
@@ -0,0 +1,639 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
subsystems.db - 数据库子系统
|
||||
|
||||
抽离自 aibox.py,提供两个独立的 SQLite 数据库类:
|
||||
- MemoryDatabase : 记忆数据库(FTS5 全文检索),存储重要对话记忆
|
||||
- MemoDatabase : 备忘录/任务数据库,支持提醒、重复、即时任务
|
||||
|
||||
依赖:common (constants, i18n, Logger) + 标准库 sqlite3
|
||||
不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from common import constants, i18n, Logger
|
||||
|
||||
|
||||
# ==================== 记忆数据库 ====================
|
||||
|
||||
class MemoryDatabase:
|
||||
"""记忆数据库 - 存储重要的对话记忆"""
|
||||
|
||||
def __init__(self, db_path: str = None):
|
||||
if db_path is None:
|
||||
db_path = os.path.join(constants.SAVE_DIR, constants.MEMORY_DB)
|
||||
|
||||
# 确保目录存在
|
||||
os.makedirs(constants.SAVE_DIR, exist_ok=True)
|
||||
self.db_path = db_path
|
||||
self.logger = Logger("MemoryDB")
|
||||
self._init_database()
|
||||
|
||||
def _init_database(self):
|
||||
"""初始化数据库表"""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
CREATE TABLE IF NOT EXISTS memories (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
topic TEXT NOT NULL,
|
||||
summary TEXT NOT NULL,
|
||||
content TEXT,
|
||||
start_time TEXT NOT NULL,
|
||||
end_time TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL,
|
||||
tags TEXT,
|
||||
importance INTEGER DEFAULT 1
|
||||
)
|
||||
''')
|
||||
|
||||
cursor.execute('''
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS memories_fts
|
||||
USING fts5(topic, summary, content, tags, content=memories)
|
||||
''')
|
||||
|
||||
conn.commit()
|
||||
|
||||
def add_memory(self, topic: str, summary: str, content: str = None,
|
||||
start_time: str = None, end_time: str = None,
|
||||
tags: List[str] = None, importance: int = None) -> int:
|
||||
"""添加记忆到数据库"""
|
||||
if start_time is None:
|
||||
start_time = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
if end_time is None:
|
||||
end_time = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
if importance is None:
|
||||
importance = constants.MEMORY_DEFAULT_IMPORTANCE
|
||||
|
||||
tags_str = ",".join(tags) if tags else ""
|
||||
created_at = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
INSERT INTO memories (topic, summary, content, start_time, end_time, created_at, tags, importance)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||||
''', (topic, summary, content, start_time, end_time, created_at, tags_str, importance))
|
||||
|
||||
memory_id = cursor.lastrowid
|
||||
|
||||
cursor.execute('''
|
||||
INSERT INTO memories_fts (rowid, topic, summary, content, tags)
|
||||
VALUES (?, ?, ?, ?, ?)
|
||||
''', (memory_id, topic, summary, content or "", tags_str))
|
||||
|
||||
conn.commit()
|
||||
|
||||
self.logger.info(f"已添加记忆: {topic} (ID: {memory_id})")
|
||||
return memory_id
|
||||
|
||||
def delete_memory(self, memory_id: int) -> bool:
|
||||
"""删除记忆"""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('DELETE FROM memories WHERE id = ?', (memory_id,))
|
||||
deleted = cursor.rowcount > 0
|
||||
|
||||
if deleted:
|
||||
cursor.execute('DELETE FROM memories_fts WHERE rowid = ?', (memory_id,))
|
||||
conn.commit()
|
||||
self.logger.info(f"已删除记忆 ID: {memory_id}")
|
||||
|
||||
return deleted
|
||||
|
||||
def search_memories(self, keywords: List[str], match_all: bool = False, limit: int = None) -> List[Dict]:
|
||||
"""多关键词模糊搜索记忆"""
|
||||
if not keywords:
|
||||
return []
|
||||
|
||||
if limit is None:
|
||||
limit = constants.MEMORY_MAX_SEARCH_RESULTS
|
||||
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
|
||||
if match_all:
|
||||
query = " AND ".join([f'"{kw}"' for kw in keywords])
|
||||
else:
|
||||
query = " OR ".join([f'"{kw}"' for kw in keywords])
|
||||
|
||||
cursor.execute('''
|
||||
SELECT m.*, rank
|
||||
FROM memories m
|
||||
JOIN memories_fts f ON m.rowid = f.rowid
|
||||
WHERE memories_fts MATCH ?
|
||||
ORDER BY m.end_time DESC, rank, m.importance DESC
|
||||
LIMIT ?
|
||||
''', (query, limit))
|
||||
|
||||
results = []
|
||||
for row in cursor.fetchall():
|
||||
memory = dict(row)
|
||||
if memory.get('tags'):
|
||||
memory['tags'] = memory['tags'].split(',')
|
||||
results.append(memory)
|
||||
|
||||
return results
|
||||
|
||||
def get_recent_memories(self, limit: int = 10) -> List[Dict]:
|
||||
"""获取最近的记忆"""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('SELECT * FROM memories ORDER BY end_time DESC LIMIT ?', (limit,))
|
||||
|
||||
results = []
|
||||
for row in cursor.fetchall():
|
||||
memory = dict(row)
|
||||
if memory.get('tags'):
|
||||
memory['tags'] = memory['tags'].split(',')
|
||||
results.append(memory)
|
||||
|
||||
return results
|
||||
|
||||
def get_memory_by_id(self, memory_id: int) -> Optional[Dict]:
|
||||
"""根据ID获取记忆"""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('SELECT * FROM memories WHERE id = ?', (memory_id,))
|
||||
row = cursor.fetchone()
|
||||
|
||||
if row:
|
||||
memory = dict(row)
|
||||
if memory.get('tags'):
|
||||
memory['tags'] = memory['tags'].split(',')
|
||||
return memory
|
||||
return None
|
||||
|
||||
def delete_memories_by_ids(self, memory_ids: List[int]) -> int:
|
||||
"""批量删除记忆"""
|
||||
if not memory_ids:
|
||||
return 0
|
||||
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
placeholders = ','.join(['?'] * len(memory_ids))
|
||||
cursor.execute(f'DELETE FROM memories WHERE id IN ({placeholders})', memory_ids)
|
||||
deleted_count = cursor.rowcount
|
||||
|
||||
if deleted_count > 0:
|
||||
cursor.execute(f'DELETE FROM memories_fts WHERE rowid IN ({placeholders})', memory_ids)
|
||||
conn.commit()
|
||||
self.logger.info(f"已批量删除 {deleted_count} 条记忆")
|
||||
|
||||
return deleted_count
|
||||
|
||||
|
||||
# ==================== 备忘录数据库 ====================
|
||||
|
||||
class MemoDatabase:
|
||||
"""备忘录数据库 - 存储和管理任务"""
|
||||
|
||||
def __init__(self, db_path: str = None):
|
||||
if db_path is None:
|
||||
db_path = os.path.join(constants.SAVE_DIR, constants.MEMORY_DB)
|
||||
|
||||
# 确保目录存在
|
||||
os.makedirs(constants.SAVE_DIR, exist_ok=True)
|
||||
self.db_path = db_path
|
||||
self.logger = Logger("MemoDB")
|
||||
self._init_database()
|
||||
|
||||
def _init_database(self):
|
||||
"""初始化备忘录表"""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute('''
|
||||
CREATE TABLE IF NOT EXISTS memos (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
title TEXT NOT NULL,
|
||||
content TEXT,
|
||||
created_at TEXT NOT NULL,
|
||||
reminder_time TEXT,
|
||||
completed_at TEXT,
|
||||
is_completed INTEGER DEFAULT 0,
|
||||
tags TEXT,
|
||||
priority INTEGER DEFAULT 1,
|
||||
created_by TEXT,
|
||||
source TEXT,
|
||||
repeat_type TEXT DEFAULT 'none',
|
||||
repeat_interval INTEGER DEFAULT 0,
|
||||
repeat_interval_unit TEXT DEFAULT 'days',
|
||||
repeat_end_time TEXT,
|
||||
last_triggered_at TEXT,
|
||||
trigger_count INTEGER DEFAULT 0,
|
||||
is_immediate INTEGER DEFAULT 0
|
||||
)
|
||||
''')
|
||||
|
||||
cursor.execute('''
|
||||
CREATE INDEX IF NOT EXISTS idx_memos_reminder
|
||||
ON memos(reminder_time) WHERE is_completed = 0
|
||||
''')
|
||||
|
||||
cursor.execute('''
|
||||
CREATE INDEX IF NOT EXISTS idx_memos_immediate
|
||||
ON memos(is_immediate) WHERE is_immediate = 1 AND is_completed = 0
|
||||
''')
|
||||
|
||||
conn.commit()
|
||||
|
||||
self.logger.info("任务表初始化完成")
|
||||
|
||||
def add_memo(self, title: str, content: str = None, reminder_time: str = None,
|
||||
tags: List[str] = None, priority: int = 1, created_by: str = None,
|
||||
source: str = None,
|
||||
repeat_type: str = 'none', repeat_interval: int = 0,
|
||||
repeat_interval_unit: str = 'days', repeat_end_time: str = None,
|
||||
is_immediate: bool = False) -> int:
|
||||
"""添加任务 - 所有任务都是对话任务,不再需要metadata和任务类型"""
|
||||
created_at = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
tags_str = ",".join(tags) if tags else ""
|
||||
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
INSERT INTO memos
|
||||
(title, content, created_at, reminder_time, tags, priority, created_by,
|
||||
is_completed, source, repeat_type, repeat_interval, repeat_interval_unit,
|
||||
repeat_end_time, is_immediate)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, 0, ?, ?, ?, ?, ?, ?)
|
||||
''', (title, content, created_at, reminder_time, tags_str, priority, created_by,
|
||||
source, repeat_type, repeat_interval, repeat_interval_unit,
|
||||
repeat_end_time, 1 if is_immediate else 0))
|
||||
|
||||
memo_id = cursor.lastrowid
|
||||
conn.commit()
|
||||
|
||||
repeat_info = f" (重复: {repeat_type}"
|
||||
if repeat_type == 'custom':
|
||||
repeat_info += f", 每{repeat_interval} {repeat_interval_unit}"
|
||||
repeat_info += ")"
|
||||
immediate_info = " (即时任务)" if is_immediate else ""
|
||||
self.logger.info(i18n.get('task_added', title, memo_id) + repeat_info + immediate_info)
|
||||
return memo_id
|
||||
|
||||
def get_memo(self, memo_id: int) -> Optional[Dict]:
|
||||
"""获取单个任务"""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('SELECT * FROM memos WHERE id = ?', (memo_id,))
|
||||
row = cursor.fetchone()
|
||||
|
||||
if row:
|
||||
memo = dict(row)
|
||||
if memo.get('tags'):
|
||||
memo['tags'] = memo['tags'].split(',')
|
||||
return memo
|
||||
return None
|
||||
|
||||
def get_pending_memos(self, limit: int = None, before_time: str = None, include_immediate: bool = False) -> List[Dict]:
|
||||
"""获取待执行的任务 - 严格使用当前时间比较"""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
|
||||
# 使用当前时间作为基准
|
||||
now = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
|
||||
# 如果传入了 before_time,使用它,否则使用当前时间
|
||||
compare_time = before_time if before_time else now
|
||||
|
||||
query = '''
|
||||
SELECT * FROM memos
|
||||
WHERE is_completed = 0
|
||||
AND reminder_time IS NOT NULL
|
||||
AND reminder_time <= ? -- 严格小于等于比较时间
|
||||
AND (repeat_end_time IS NULL OR repeat_end_time >= ?)
|
||||
'''
|
||||
|
||||
if not include_immediate:
|
||||
query += ' AND is_immediate = 0'
|
||||
|
||||
query += ' ORDER BY reminder_time ASC, priority DESC'
|
||||
|
||||
params = [compare_time, compare_time]
|
||||
|
||||
if limit:
|
||||
query += f' LIMIT {limit}'
|
||||
|
||||
cursor.execute(query, params)
|
||||
|
||||
results = []
|
||||
for row in cursor.fetchall():
|
||||
memo = dict(row)
|
||||
if memo.get('tags'):
|
||||
memo['tags'] = memo['tags'].split(',')
|
||||
|
||||
# 【关键修复】二次验证:确保 reminder_time 确实 <= 当前时间
|
||||
if memo.get('reminder_time'):
|
||||
try:
|
||||
reminder_time = datetime.strptime(memo['reminder_time'], constants.DATETIME_FORMAT)
|
||||
current_time = datetime.strptime(compare_time, constants.DATETIME_FORMAT)
|
||||
if reminder_time <= current_time:
|
||||
results.append(memo)
|
||||
else:
|
||||
# 调试日志
|
||||
self.logger.debug(f"过滤掉未到期任务 #{memo['id']}: {memo['reminder_time']} > {compare_time}")
|
||||
except (ValueError, TypeError) as e:
|
||||
self.logger.warning(f"解析任务 #{memo['id']} 的提醒时间失败: {e}")
|
||||
else:
|
||||
# 没有提醒时间的任务不应该在待执行列表中
|
||||
self.logger.warning(f"任务 #{memo['id']} 没有提醒时间,但被返回为待执行")
|
||||
|
||||
return results
|
||||
|
||||
def get_immediate_memos(self, limit: int = None) -> List[Dict]:
|
||||
"""获取待执行的即时任务"""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
|
||||
query = '''
|
||||
SELECT * FROM memos
|
||||
WHERE is_completed = 0
|
||||
AND is_immediate = 1
|
||||
ORDER BY created_at ASC
|
||||
'''
|
||||
|
||||
if limit:
|
||||
query += f' LIMIT {limit}'
|
||||
|
||||
cursor.execute(query)
|
||||
|
||||
results = []
|
||||
for row in cursor.fetchall():
|
||||
memo = dict(row)
|
||||
if memo.get('tags'):
|
||||
memo['tags'] = memo['tags'].split(',')
|
||||
results.append(memo)
|
||||
|
||||
return results
|
||||
|
||||
def get_overdue_memos(self, limit: int = None) -> List[Dict]:
|
||||
"""获取已过期的任务"""
|
||||
now = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
return self.get_pending_memos(limit=limit, before_time=now)
|
||||
|
||||
def get_all_memos(self, include_completed: bool = False, limit: int = 50) -> List[Dict]:
|
||||
"""获取所有任务"""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
|
||||
if include_completed:
|
||||
cursor.execute('''
|
||||
SELECT * FROM memos
|
||||
ORDER BY
|
||||
CASE WHEN is_completed = 0 THEN 0 ELSE 1 END,
|
||||
reminder_time ASC,
|
||||
created_at DESC
|
||||
LIMIT ?
|
||||
''', (limit,))
|
||||
else:
|
||||
cursor.execute('''
|
||||
SELECT * FROM memos
|
||||
WHERE is_completed = 0
|
||||
ORDER BY reminder_time ASC, priority DESC, created_at DESC
|
||||
LIMIT ?
|
||||
''', (limit,))
|
||||
|
||||
results = []
|
||||
for row in cursor.fetchall():
|
||||
memo = dict(row)
|
||||
if memo.get('tags'):
|
||||
memo['tags'] = memo['tags'].split(',')
|
||||
results.append(memo)
|
||||
|
||||
return results
|
||||
|
||||
def mark_as_triggered(self, memo_id: int) -> Optional[str]:
|
||||
"""标记任务已被触发,并根据重复规则计算下一次提醒时间"""
|
||||
memo = self.get_memo(memo_id)
|
||||
if not memo:
|
||||
return None
|
||||
|
||||
now = datetime.now()
|
||||
next_reminder = None
|
||||
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute('''
|
||||
UPDATE memos
|
||||
SET last_triggered_at = ?, trigger_count = trigger_count + 1
|
||||
WHERE id = ?
|
||||
''', (now.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||
|
||||
# 处理重复任务
|
||||
if memo['repeat_type'] != 'none' and not self._should_end_repeat(memo, now):
|
||||
# 使用当前时间作为基准来计算下一次执行时间
|
||||
next_reminder = self._calculate_next_reminder(memo, now)
|
||||
|
||||
if next_reminder:
|
||||
cursor.execute('''
|
||||
UPDATE memos
|
||||
SET reminder_time = ?
|
||||
WHERE id = ?
|
||||
''', (next_reminder.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||
|
||||
self.logger.info(f"任务 #{memo_id} 已更新下一次执行时间: {next_reminder}")
|
||||
else:
|
||||
# 重复结束,标记为完成
|
||||
cursor.execute('''
|
||||
UPDATE memos
|
||||
SET is_completed = 1, completed_at = ?
|
||||
WHERE id = ?
|
||||
''', (now.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||
|
||||
self.logger.info(f"任务 #{memo_id} 重复结束,已标记为完成")
|
||||
elif memo['is_immediate']:
|
||||
# 即时任务执行后标记为完成
|
||||
cursor.execute('''
|
||||
UPDATE memos
|
||||
SET is_completed = 1, completed_at = ?
|
||||
WHERE id = ?
|
||||
''', (now.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||
|
||||
self.logger.info(f"即时任务 #{memo_id} 已完成")
|
||||
else:
|
||||
# 一次性任务,标记为完成
|
||||
cursor.execute('''
|
||||
UPDATE memos
|
||||
SET is_completed = 1, completed_at = ?
|
||||
WHERE id = ?
|
||||
''', (now.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||
|
||||
self.logger.info(f"任务 #{memo_id} 已完成")
|
||||
|
||||
conn.commit()
|
||||
|
||||
return next_reminder.strftime(constants.DATETIME_FORMAT) if next_reminder else None
|
||||
|
||||
def _should_end_repeat(self, memo: Dict, current_time: datetime) -> bool:
|
||||
"""检查是否应该结束重复"""
|
||||
if memo['repeat_end_time']:
|
||||
end_time = datetime.strptime(memo['repeat_end_time'], constants.DATETIME_FORMAT)
|
||||
if current_time > end_time:
|
||||
return True
|
||||
return False
|
||||
|
||||
def _calculate_next_reminder(self, memo: Dict, current_time: datetime) -> Optional[datetime]:
|
||||
"""
|
||||
计算下一次提醒时间,使用当前时间作为基准
|
||||
支持分钟、小时、天为单位
|
||||
"""
|
||||
if memo['repeat_type'] == 'daily':
|
||||
# 每天重复:当前时间 + 1天
|
||||
next_time = current_time + timedelta(days=1)
|
||||
elif memo['repeat_type'] == 'weekly':
|
||||
# 每周重复:当前时间 + 7天
|
||||
next_time = current_time + timedelta(weeks=1)
|
||||
elif memo['repeat_type'] == 'monthly':
|
||||
# 每月重复:尝试在当前时间上加一个月
|
||||
try:
|
||||
# 处理月份进位
|
||||
if current_time.month == 12:
|
||||
next_time = current_time.replace(year=current_time.year + 1, month=1)
|
||||
else:
|
||||
next_time = current_time.replace(month=current_time.month + 1)
|
||||
except ValueError:
|
||||
# 处理月份天数问题(如1月31日加一个月)
|
||||
if current_time.month == 12:
|
||||
next_time = current_time.replace(year=current_time.year + 1, month=1, day=1) - timedelta(days=1)
|
||||
else:
|
||||
next_time = current_time.replace(month=current_time.month + 1, day=1) - timedelta(days=1)
|
||||
elif memo['repeat_type'] == 'custom' and memo['repeat_interval'] > 0:
|
||||
unit = memo.get('repeat_interval_unit', 'days')
|
||||
if unit == 'minutes':
|
||||
# 每N分钟重复:当前时间 + N分钟
|
||||
next_time = current_time + timedelta(minutes=memo['repeat_interval'])
|
||||
elif unit == 'hours':
|
||||
# 每N小时重复:当前时间 + N小时
|
||||
next_time = current_time + timedelta(hours=memo['repeat_interval'])
|
||||
else: # days
|
||||
# 每N天重复:当前时间 + N天
|
||||
next_time = current_time + timedelta(days=memo['repeat_interval'])
|
||||
else:
|
||||
return None
|
||||
|
||||
# 检查是否超过重复结束时间
|
||||
if memo['repeat_end_time']:
|
||||
end_time = datetime.strptime(memo['repeat_end_time'], constants.DATETIME_FORMAT)
|
||||
if next_time > end_time:
|
||||
return None
|
||||
|
||||
return next_time
|
||||
|
||||
def complete_memo(self, memo_id: int, force_complete: bool = False) -> bool:
|
||||
"""完成任务"""
|
||||
memo = self.get_memo(memo_id)
|
||||
if not memo:
|
||||
return False
|
||||
|
||||
if memo['repeat_type'] != 'none' and not force_complete:
|
||||
self.mark_as_triggered(memo_id)
|
||||
return True
|
||||
|
||||
completed_at = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
UPDATE memos
|
||||
SET is_completed = 1, completed_at = ?
|
||||
WHERE id = ? AND is_completed = 0
|
||||
''', (completed_at, memo_id))
|
||||
|
||||
affected = cursor.rowcount
|
||||
conn.commit()
|
||||
|
||||
if affected > 0:
|
||||
self.logger.info(i18n.get('memo_completed', memo_id))
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_memo_stats(self) -> Dict:
|
||||
"""获取任务统计信息"""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute('SELECT COUNT(*) FROM memos')
|
||||
total = cursor.fetchone()[0]
|
||||
|
||||
cursor.execute('SELECT COUNT(*) FROM memos WHERE is_completed = 0')
|
||||
pending = cursor.fetchone()[0]
|
||||
|
||||
now = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
cursor.execute('''
|
||||
SELECT COUNT(*) FROM memos
|
||||
WHERE is_completed = 0
|
||||
AND reminder_time <= ?
|
||||
AND is_immediate = 0
|
||||
''', (now,))
|
||||
overdue = cursor.fetchone()[0]
|
||||
|
||||
cursor.execute('SELECT COUNT(*) FROM memos WHERE is_completed = 1')
|
||||
completed = cursor.fetchone()[0]
|
||||
|
||||
cursor.execute('SELECT COUNT(*) FROM memos WHERE is_immediate = 1 AND is_completed = 0')
|
||||
immediate = cursor.fetchone()[0]
|
||||
|
||||
return {
|
||||
"total": total,
|
||||
"pending": pending,
|
||||
"overdue": overdue,
|
||||
"completed": completed,
|
||||
"immediate": immediate
|
||||
}
|
||||
|
||||
def add_user_task(self, title: str, content: str = None, reminder_time: str = None,
|
||||
repeat_type: str = 'none', repeat_interval: int = 0,
|
||||
repeat_interval_unit: str = 'days', is_immediate: bool = False) -> int:
|
||||
"""添加用户任务(前端直接调用的任务)- 所有任务都是对话任务"""
|
||||
if not constants.MEMO_ALLOW_USER_TASKS:
|
||||
raise Exception("用户任务功能已禁用")
|
||||
|
||||
return self.add_memo(
|
||||
title=title,
|
||||
content=content,
|
||||
reminder_time=reminder_time,
|
||||
repeat_type=repeat_type,
|
||||
repeat_interval=repeat_interval,
|
||||
repeat_interval_unit=repeat_interval_unit,
|
||||
created_by="user",
|
||||
source="websocket",
|
||||
is_immediate=is_immediate
|
||||
)
|
||||
|
||||
def get_pending_tasks_for_user(self) -> List[Dict]:
|
||||
"""获取待执行任务列表(用于前端显示)"""
|
||||
pending = self.get_pending_memos(include_immediate=False)
|
||||
immediate = self.get_immediate_memos()
|
||||
result = []
|
||||
|
||||
now = datetime.now()
|
||||
for memo in pending + immediate:
|
||||
result.append({
|
||||
"id": memo['id'],
|
||||
"title": memo['title'],
|
||||
"content": memo['content'],
|
||||
"reminder_time": memo['reminder_time'],
|
||||
"priority": memo['priority'],
|
||||
"is_overdue": datetime.strptime(memo['reminder_time'], constants.DATETIME_FORMAT) < now if memo['reminder_time'] else False,
|
||||
"repeat_type": memo['repeat_type'],
|
||||
"repeat_interval": memo.get('repeat_interval', 0),
|
||||
"repeat_interval_unit": memo.get('repeat_interval_unit', 'days'),
|
||||
"trigger_count": memo.get('trigger_count', 0),
|
||||
"is_immediate": memo.get('is_immediate', False)
|
||||
})
|
||||
|
||||
return result
|
||||
@@ -0,0 +1,349 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
subsystems.executors - 命令执行器
|
||||
|
||||
抽离自 aibox.py,提供命令执行的抽象基类与 subprocess 实现。
|
||||
- CommandExecutor : 抽象基类,定义 execute / execute_streaming / cancel
|
||||
- SubprocessCommandExecutor : 基于 subprocess.Popen 的实现,支持超时/工作目录/
|
||||
1MB 输出上限/线程安全取消
|
||||
|
||||
依赖:common (ConfigManager, constants, Logger) + 标准库
|
||||
不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
import os
|
||||
import queue
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
from common import ConfigManager, constants, Logger
|
||||
|
||||
|
||||
class CommandExecutor(ABC):
|
||||
"""命令执行器抽象基类 - 负责执行系统命令"""
|
||||
|
||||
def __init__(self, logger: Optional[Logger] = None):
|
||||
self.logger = logger or Logger("CommandExecutor")
|
||||
self.config = ConfigManager()
|
||||
self.process = None
|
||||
|
||||
@abstractmethod
|
||||
def execute(self, command: str, timeout: int = 30, **kwargs) -> str:
|
||||
"""执行命令并返回结果(非流式)"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def execute_streaming(self, command: str, output_queue: queue.Queue, timeout: int = 30, **kwargs):
|
||||
"""流式执行命令,通过队列输出结果"""
|
||||
pass
|
||||
|
||||
def is_enabled(self) -> bool:
|
||||
"""检查命令执行是否启用"""
|
||||
return constants.COMMANDS_ENABLED
|
||||
|
||||
def format_result(self, success: bool, message: str, data: Dict = None) -> str:
|
||||
"""统一结果格式化"""
|
||||
result = []
|
||||
emoji = "✅" if success else "❌"
|
||||
result.append(f"{emoji} {message}")
|
||||
|
||||
if data:
|
||||
for key, value in data.items():
|
||||
if isinstance(value, str) and "\n" in value:
|
||||
result.append(f" {key}:")
|
||||
for line in value.split("\n"):
|
||||
result.append(f" {line}")
|
||||
else:
|
||||
result.append(f" {key}: {value}")
|
||||
|
||||
return "\n".join(result)
|
||||
|
||||
|
||||
class SubprocessCommandExecutor(CommandExecutor):
|
||||
"""基于subprocess的命令执行器 - 支持超时、工作目录、输出大小限制和取消"""
|
||||
|
||||
# 命令输出最大字节数(1MB),超出截断
|
||||
MAX_OUTPUT_BYTES = 1 * 1024 * 1024
|
||||
|
||||
def __init__(self, logger=None):
|
||||
super().__init__(logger)
|
||||
self._current_process = None
|
||||
self._process_lock = threading.Lock()
|
||||
|
||||
def cancel(self):
|
||||
"""取消当前正在运行的命令(kill子进程)。线程安全。"""
|
||||
with self._process_lock:
|
||||
proc = self._current_process
|
||||
if proc is not None:
|
||||
try:
|
||||
if proc.poll() is None:
|
||||
proc.kill()
|
||||
self.logger.info("已终止正在运行的子进程")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _read_with_limit(self, process, timeout: int, start_time: float, output_queue=None) -> Tuple[str, bool, bool]:
|
||||
"""从进程读取输出,限制总大小和超时。
|
||||
返回 (output_str, timeout_occurred, truncated)
|
||||
"""
|
||||
output_chunks = []
|
||||
total_bytes = 0
|
||||
timeout_occurred = False
|
||||
truncated = False
|
||||
|
||||
while True:
|
||||
# 检查超时
|
||||
if timeout > 0 and time.time() - start_time > timeout:
|
||||
process.terminate()
|
||||
timeout_occurred = True
|
||||
break
|
||||
|
||||
# 检查输出大小
|
||||
if total_bytes >= self.MAX_OUTPUT_BYTES:
|
||||
truncated = True
|
||||
process.terminate()
|
||||
break
|
||||
|
||||
# 读取一行
|
||||
line = process.stdout.readline()
|
||||
if not line and process.poll() is not None:
|
||||
break
|
||||
if line:
|
||||
line_bytes = len(line.encode('utf-8', errors='replace'))
|
||||
if total_bytes + line_bytes > self.MAX_OUTPUT_BYTES and not truncated:
|
||||
allowed = self.MAX_OUTPUT_BYTES - total_bytes
|
||||
if allowed > 0:
|
||||
output_chunks.append(line.encode('utf-8', errors='replace')[:allowed].decode('utf-8', errors='replace'))
|
||||
total_bytes = self.MAX_OUTPUT_BYTES
|
||||
truncated = True
|
||||
process.terminate()
|
||||
break
|
||||
|
||||
output_chunks.append(line)
|
||||
total_bytes += line_bytes
|
||||
|
||||
if output_queue:
|
||||
output_queue.put(("line", line.rstrip('\n') + "\n"))
|
||||
|
||||
return ''.join(output_chunks), timeout_occurred, truncated
|
||||
|
||||
def execute(self, command: str, timeout: int = 30, work_dir: str = None, env: Dict = None, **kwargs) -> str:
|
||||
"""执行命令并返回完整结果(带输出大小限制)"""
|
||||
if not self.is_enabled():
|
||||
return "❌ 命令执行功能已禁用"
|
||||
|
||||
# 默认工作目录为workspace,避免在启动目录散落文件
|
||||
if work_dir is None or work_dir == ".":
|
||||
work_dir = constants.WORKSPACE_DIR
|
||||
|
||||
self.logger.info(f"执行命令: {command} (work_dir={work_dir})")
|
||||
|
||||
# 准备环境变量
|
||||
exec_env = os.environ.copy()
|
||||
if env:
|
||||
exec_env.update(env)
|
||||
|
||||
process = None
|
||||
try:
|
||||
process = subprocess.Popen(
|
||||
command,
|
||||
shell=True,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
stdin=subprocess.PIPE,
|
||||
text=True,
|
||||
bufsize=1,
|
||||
universal_newlines=True,
|
||||
env=exec_env,
|
||||
cwd=os.path.expanduser(work_dir)
|
||||
)
|
||||
# 保存进程引用以便cancel()可以kill它
|
||||
with self._process_lock:
|
||||
self._current_process = process
|
||||
|
||||
start_time = time.time()
|
||||
output_str, timeout_occurred, truncated = self._read_with_limit(process, timeout, start_time)
|
||||
|
||||
# 确保进程终止
|
||||
if process.poll() is None:
|
||||
if timeout_occurred or truncated:
|
||||
process.kill()
|
||||
else:
|
||||
process.terminate()
|
||||
try:
|
||||
process.wait(timeout=2)
|
||||
except:
|
||||
process.kill()
|
||||
|
||||
return_code = process.returncode if process.returncode is not None else -1
|
||||
|
||||
data = {
|
||||
"输出": output_str.strip() if output_str.strip() else "(无输出)",
|
||||
"返回码": return_code
|
||||
}
|
||||
|
||||
if timeout_occurred:
|
||||
return self.format_result(False, f"命令执行超时({timeout}秒)", data)
|
||||
|
||||
if truncated:
|
||||
data["警告"] = f"输出超过 {self.MAX_OUTPUT_BYTES/1024/1024:.0f}MB 限制,结果已截断"
|
||||
return self.format_result(False, "命令执行完成(输出已截断)", data)
|
||||
|
||||
return self.format_result(return_code == 0, "命令执行完成", data)
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
return f"❌ 命令执行超时({timeout}秒)"
|
||||
except FileNotFoundError as e:
|
||||
return f"❌ 命令执行错误: 命令未找到 - {str(e)}"
|
||||
except PermissionError as e:
|
||||
return f"❌ 命令执行错误: 权限不足 - {str(e)}"
|
||||
except Exception as e:
|
||||
return f"❌ 执行错误: {str(e)}"
|
||||
finally:
|
||||
if process and process.poll() is None:
|
||||
process.kill()
|
||||
# 清除当前进程引用
|
||||
with self._process_lock:
|
||||
self._current_process = None
|
||||
|
||||
def execute_streaming(self, command: str, output_queue: queue.Queue, timeout: int = 30,
|
||||
work_dir: str = None, env: Dict = None, **kwargs):
|
||||
"""流式执行命令,支持超时控制、工作目录和输出大小限制"""
|
||||
if not self.is_enabled():
|
||||
output_queue.put(("line", "❌ 命令执行功能已禁用\n"))
|
||||
output_queue.put(("complete", "❌ 命令执行功能已禁用"))
|
||||
return
|
||||
|
||||
# 默认工作目录为workspace,避免在启动目录散落文件
|
||||
if work_dir is None or work_dir == ".":
|
||||
work_dir = constants.WORKSPACE_DIR
|
||||
|
||||
self.logger.info(f"流式执行命令: {command} (work_dir={work_dir})")
|
||||
output_queue.put(("line", f"\n🔧 执行命令: {command}\n"))
|
||||
output_queue.put(("line", f"⏱️ 超时设置: {timeout} 秒\n"))
|
||||
output_queue.put(("line", f"📦 输出限制: {self.MAX_OUTPUT_BYTES/1024/1024:.0f}MB\n"))
|
||||
output_queue.put(("line", f"📁 工作目录: {work_dir}\n"))
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
# 准备环境变量
|
||||
exec_env = os.environ.copy()
|
||||
if env:
|
||||
exec_env.update(env)
|
||||
|
||||
process = None
|
||||
try:
|
||||
process = subprocess.Popen(
|
||||
command,
|
||||
shell=True,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
stdin=subprocess.PIPE,
|
||||
text=True,
|
||||
bufsize=1,
|
||||
universal_newlines=True,
|
||||
env=exec_env,
|
||||
cwd=os.path.expanduser(work_dir)
|
||||
)
|
||||
self.process = process
|
||||
# 保存进程引用以便cancel()可以kill它
|
||||
with self._process_lock:
|
||||
self._current_process = process
|
||||
|
||||
start_time = time.time()
|
||||
output_lines = []
|
||||
total_bytes = 0
|
||||
timeout_occurred = False
|
||||
truncated = False
|
||||
|
||||
# 实时读取输出
|
||||
while True:
|
||||
# 检查超时
|
||||
if timeout > 0 and time.time() - start_time > timeout:
|
||||
process.terminate()
|
||||
timeout_occurred = True
|
||||
output_queue.put(("line", f"\n⏰ 命令执行超时({timeout}秒)\n"))
|
||||
break
|
||||
|
||||
# 检查输出大小
|
||||
if total_bytes >= self.MAX_OUTPUT_BYTES:
|
||||
truncated = True
|
||||
process.terminate()
|
||||
output_queue.put(("line", f"\n📦 输出超过 {self.MAX_OUTPUT_BYTES/1024/1024:.0f}MB 限制,结果已截断\n"))
|
||||
break
|
||||
|
||||
# 读取一行输出
|
||||
line = process.stdout.readline()
|
||||
if not line and process.poll() is not None:
|
||||
break
|
||||
if line:
|
||||
line_bytes = len(line.encode('utf-8', errors='replace'))
|
||||
|
||||
if total_bytes + line_bytes > self.MAX_OUTPUT_BYTES and not truncated:
|
||||
allowed = self.MAX_OUTPUT_BYTES - total_bytes
|
||||
if allowed > 0:
|
||||
truncated_line = line.encode('utf-8', errors='replace')[:allowed].decode('utf-8', errors='replace')
|
||||
output_queue.put(("line", truncated_line))
|
||||
output_lines.append(truncated_line)
|
||||
total_bytes = self.MAX_OUTPUT_BYTES
|
||||
truncated = True
|
||||
process.terminate()
|
||||
output_queue.put(("line", f"\n📦 输出超过 {self.MAX_OUTPUT_BYTES/1024/1024:.0f}MB 限制,结果已截断\n"))
|
||||
break
|
||||
|
||||
stripped = line.rstrip('\n')
|
||||
output_queue.put(("line", stripped + "\n"))
|
||||
output_lines.append(line)
|
||||
total_bytes += line_bytes
|
||||
|
||||
# 如果超时或截断,确保进程被终止
|
||||
if (timeout_occurred or truncated) and process.poll() is None:
|
||||
process.kill()
|
||||
try:
|
||||
process.wait(timeout=2)
|
||||
except:
|
||||
pass
|
||||
|
||||
return_code = process.returncode if process.returncode is not None else -1
|
||||
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
if timeout_occurred:
|
||||
output_queue.put(("line", f"⏰ 命令执行超时({timeout}秒),返回码: {return_code}\n"))
|
||||
elif truncated:
|
||||
output_queue.put(("line", f"📦 输出已截断,返回码: {return_code}\n"))
|
||||
else:
|
||||
output_queue.put(("line", f"✅ 命令执行完成,返回码: {return_code}\n"))
|
||||
|
||||
# 准备完整结果
|
||||
full_output = "".join(output_lines)
|
||||
data = {
|
||||
"输出": full_output.strip() if full_output.strip() else "(无输出)",
|
||||
"返回码": return_code
|
||||
}
|
||||
|
||||
if timeout_occurred:
|
||||
result = self.format_result(False, f"命令执行超时({timeout}秒)", data)
|
||||
elif truncated:
|
||||
data["警告"] = f"输出超过 {self.MAX_OUTPUT_BYTES/1024/1024:.0f}MB 限制,结果已截断"
|
||||
result = self.format_result(False, "命令执行完成(输出已截断)", data)
|
||||
else:
|
||||
result = self.format_result(return_code == 0, "命令执行完成", data)
|
||||
|
||||
output_queue.put(("complete", result))
|
||||
|
||||
except Exception as e:
|
||||
output_queue.put(("line", f"❌ 执行错误: {str(e)}\n"))
|
||||
output_queue.put(("complete", f"❌ 执行错误: {str(e)}"))
|
||||
finally:
|
||||
if process and process.poll() is None:
|
||||
process.terminate()
|
||||
try:
|
||||
process.wait(timeout=2)
|
||||
except:
|
||||
process.kill()
|
||||
self.process = None
|
||||
with self._process_lock:
|
||||
self._current_process = None
|
||||
@@ -0,0 +1,308 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
subsystems.prompt_manager - 提示词管理器
|
||||
|
||||
抽离自 aibox.py,负责从 prompts 目录加载系统提示词。
|
||||
单例模式(__new__ + 类变量 _instance),多实例共享同一份 _prompts。
|
||||
|
||||
文件命名规范:
|
||||
- 00_IDENTITY.md : 用户身份信息
|
||||
- N_*.md/txt : 基础技能提示词(N 为数字),按数字升序加载
|
||||
- 99_SOUL.md : 智能体灵魂文件(角色定位、身份定义、性格等)
|
||||
- 100_STATUS.md : 智能体状态记录
|
||||
|
||||
依赖:common (ConfigManager, constants) + 标准库
|
||||
不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
from datetime import datetime
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from common import ConfigManager
|
||||
|
||||
|
||||
class PromptManager:
|
||||
"""提示词管理器 - 负责从 prompts 目录加载系统提示词
|
||||
|
||||
文件命名规范:
|
||||
- 00_IDENTITY.md: 用户身份信息
|
||||
- N_*.md/txt: 基础技能提示词(N 为数字,如 1_memory_search.md、10_task_handling.md),按数字升序加载
|
||||
- 99_SOUL.md: 智能体灵魂文件(角色定位、身份定义、性格等)
|
||||
- 100_STATUS.md: 智能体状态记录
|
||||
"""
|
||||
|
||||
_instance = None
|
||||
_prompts = []
|
||||
_prompts_dir = None
|
||||
|
||||
# 核心文件
|
||||
IDENTITY_FILE = "00_IDENTITY.md"
|
||||
STATUS_FILE = "100_STATUS.md"
|
||||
SOUL_FILE = "99_SOUL.md" # 智能体灵魂文件 - 所有角色定位内容
|
||||
|
||||
# 基础技能文件名模式:N_xxx.md(N 为数字,如 1_memory_search.md、10_task_handling.md)
|
||||
_SKILL_FILE_RE = re.compile(r'^(\d+)_')
|
||||
|
||||
def __new__(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
return cls._instance
|
||||
|
||||
def __init__(self):
|
||||
if not self._prompts:
|
||||
self._load_prompts()
|
||||
|
||||
def _get_prompts_dir(self) -> str:
|
||||
config = ConfigManager()
|
||||
prompts_dir = config.get("system.prompts_dir", "~/.aibox/prompts")
|
||||
return os.path.expanduser(prompts_dir)
|
||||
|
||||
def _load_prompts(self):
|
||||
self._prompts_dir = self._get_prompts_dir()
|
||||
os.makedirs(self._prompts_dir, exist_ok=True)
|
||||
self._prompts = []
|
||||
|
||||
# 确保核心文件存在
|
||||
self._ensure_identity_file()
|
||||
self._ensure_status_file()
|
||||
self._ensure_soul_file() # 智能体灵魂文件
|
||||
|
||||
# 确保基础技能文件存在
|
||||
self._ensure_skill_files()
|
||||
|
||||
prompt_files = []
|
||||
if os.path.exists(self._prompts_dir):
|
||||
# 按文件名排序,确保1_开头的技能文件在99_SOUL.md之前加载
|
||||
prompt_files = sorted([f for f in os.listdir(self._prompts_dir)
|
||||
if os.path.isfile(os.path.join(self._prompts_dir, f))
|
||||
and not f.startswith('.')])
|
||||
|
||||
for filename in prompt_files:
|
||||
filepath = os.path.join(self._prompts_dir, filename)
|
||||
try:
|
||||
with open(filepath, 'r', encoding='utf-8') as f:
|
||||
content = f.read().strip()
|
||||
if content:
|
||||
self._prompts.append({
|
||||
'file': filename,
|
||||
'content': content
|
||||
})
|
||||
except Exception as e:
|
||||
print(f"⚠️ 读取提示词文件 {filename} 失败: {e}")
|
||||
|
||||
essential_files = [self.IDENTITY_FILE, self.STATUS_FILE, self.SOUL_FILE]
|
||||
other_files = [f for f in prompt_files if f not in essential_files] if prompt_files else []
|
||||
|
||||
if not self._prompts:
|
||||
self._create_default_prompts()
|
||||
|
||||
def _get_template_dir(self) -> str:
|
||||
"""获取仓库中的提示词模板目录(aibox/prompts/)"""
|
||||
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
# subsystems/ 上一层是 aibox/,模板在其下的 prompts/
|
||||
aibox_dir = os.path.dirname(script_dir)
|
||||
return os.path.join(aibox_dir, 'prompts')
|
||||
|
||||
def _ensure_from_template(self, filename: str):
|
||||
"""从模板目录复制提示词文件到运行目录。
|
||||
|
||||
模板目录为 aibox/prompts/,若其中存在同名文件则复制过去;
|
||||
若模板不存在则跳过(由调用方保证模板齐全)。
|
||||
已存在的运行时文件不会被覆盖。
|
||||
"""
|
||||
target_path = os.path.join(self._prompts_dir, filename)
|
||||
if os.path.exists(target_path):
|
||||
return # 已存在,不覆盖
|
||||
|
||||
template_path = os.path.join(self._get_template_dir(), filename)
|
||||
if os.path.exists(template_path):
|
||||
try:
|
||||
shutil.copy2(template_path, target_path)
|
||||
print(f"✅ 已从模板创建: {filename}")
|
||||
except Exception as e:
|
||||
print(f"⚠️ 从模板复制 {filename} 失败: {e}")
|
||||
|
||||
def _ensure_identity_file(self):
|
||||
"""确保用户身份文件存在"""
|
||||
self._ensure_from_template(self.IDENTITY_FILE)
|
||||
|
||||
def _ensure_status_file(self):
|
||||
"""确保状态文件存在"""
|
||||
self._ensure_from_template(self.STATUS_FILE)
|
||||
|
||||
def _ensure_soul_file(self):
|
||||
"""确保智能体灵魂文件存在"""
|
||||
self._ensure_from_template(self.SOUL_FILE)
|
||||
|
||||
def _ensure_skill_files(self):
|
||||
"""确保基础技能文件存在 - 从模板目录复制所有 N_*.md 文件"""
|
||||
template_dir = self._get_template_dir()
|
||||
if not os.path.isdir(template_dir):
|
||||
return
|
||||
for filename in sorted(os.listdir(template_dir)):
|
||||
if self._SKILL_FILE_RE.match(filename): # N_ 开头的文件
|
||||
self._ensure_from_template(filename)
|
||||
|
||||
def update_identity(self, content: str) -> bool:
|
||||
"""更新用户身份信息"""
|
||||
identity_path = os.path.join(self._prompts_dir, self.IDENTITY_FILE)
|
||||
try:
|
||||
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
updated_content = f"""# 用户身份信息
|
||||
|
||||
{content}
|
||||
|
||||
【更新记录】
|
||||
- 上次更新:{timestamp}
|
||||
"""
|
||||
with open(identity_path, 'w', encoding='utf-8') as f:
|
||||
f.write(updated_content)
|
||||
self.reload()
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"⚠️ 更新身份文件失败: {e}")
|
||||
return False
|
||||
|
||||
def update_status(self, content: str) -> bool:
|
||||
"""更新智能体状态"""
|
||||
status_path = os.path.join(self._prompts_dir, self.STATUS_FILE)
|
||||
try:
|
||||
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
updated_content = f"""# 智能体状态记录
|
||||
|
||||
{content}
|
||||
|
||||
【更新记录】
|
||||
- 上次更新:{timestamp}
|
||||
"""
|
||||
with open(status_path, 'w', encoding='utf-8') as f:
|
||||
f.write(updated_content)
|
||||
self.reload()
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"⚠️ 更新状态文件失败: {e}")
|
||||
return False
|
||||
|
||||
def update_soul(self, content: str) -> bool:
|
||||
"""更新智能体灵魂定义 - 所有角色定位内容"""
|
||||
soul_path = os.path.join(self._prompts_dir, self.SOUL_FILE)
|
||||
try:
|
||||
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
updated_content = f"""# 智能体灵魂定义 - 我是谁
|
||||
|
||||
{content}
|
||||
|
||||
【更新记录】
|
||||
- 上次更新:{timestamp}
|
||||
"""
|
||||
with open(soul_path, 'w', encoding='utf-8') as f:
|
||||
f.write(updated_content)
|
||||
self.reload()
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"⚠️ 更新灵魂文件失败: {e}")
|
||||
return False
|
||||
|
||||
def get_identity(self) -> str:
|
||||
"""获取用户身份信息"""
|
||||
identity_path = os.path.join(self._prompts_dir, self.IDENTITY_FILE)
|
||||
try:
|
||||
if os.path.exists(identity_path):
|
||||
with open(identity_path, 'r', encoding='utf-8') as f:
|
||||
return f.read().strip()
|
||||
except Exception as e:
|
||||
print(f"⚠️ 读取身份文件失败: {e}")
|
||||
return "尚不清楚 - 无法读取身份信息"
|
||||
|
||||
def get_status(self) -> str:
|
||||
"""获取智能体状态"""
|
||||
status_path = os.path.join(self._prompts_dir, self.STATUS_FILE)
|
||||
try:
|
||||
if os.path.exists(status_path):
|
||||
with open(status_path, 'r', encoding='utf-8') as f:
|
||||
return f.read().strip()
|
||||
except Exception as e:
|
||||
print(f"⚠️ 读取状态文件失败: {e}")
|
||||
return "这是第一次对话 - 无法读取状态信息"
|
||||
|
||||
def get_soul(self) -> str:
|
||||
"""获取智能体灵魂定义 - 角色定位核心"""
|
||||
soul_path = os.path.join(self._prompts_dir, self.SOUL_FILE)
|
||||
try:
|
||||
if os.path.exists(soul_path):
|
||||
with open(soul_path, 'r', encoding='utf-8') as f:
|
||||
return f.read().strip()
|
||||
except Exception as e:
|
||||
print(f"⚠️ 读取灵魂文件失败: {e}")
|
||||
return """我是DeepSeek,一个智能AI助手。
|
||||
请创建 99_SOUL.md 文件来定义我的完整身份、角色、性格和能力边界。"""
|
||||
|
||||
def get_skill_prompts(self) -> List[Dict]:
|
||||
"""获取所有基础技能提示词(以1_开头的文件)"""
|
||||
return [p for p in self._prompts if p['file'].startswith(self.SKILL_PREFIX)]
|
||||
|
||||
def get_soul_prompt(self) -> Optional[Dict]:
|
||||
"""获取灵魂提示词(99_SOUL.md)"""
|
||||
for p in self._prompts:
|
||||
if p['file'] == self.SOUL_FILE:
|
||||
return p
|
||||
return None
|
||||
|
||||
def _create_default_prompts(self):
|
||||
"""创建默认提示词(已由各 _ensure_* 方法处理)"""
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def _skill_sort_key(cls, filename: str):
|
||||
"""技能文件按前缀数字排序,非技能文件返回 None。"""
|
||||
m = cls._SKILL_FILE_RE.match(filename)
|
||||
return int(m.group(1)) if m else None
|
||||
|
||||
def _prompt_sort_key(self, p):
|
||||
"""排序规则:00_IDENTITY → 技能文件(按数字升序) → 99_SOUL → 100_STATUS → 其他"""
|
||||
file = p['file']
|
||||
if file == self.IDENTITY_FILE:
|
||||
return (0,)
|
||||
skill_num = self._skill_sort_key(file)
|
||||
if skill_num is not None:
|
||||
return (1, skill_num)
|
||||
if file == self.SOUL_FILE:
|
||||
return (2,)
|
||||
if file == self.STATUS_FILE:
|
||||
return (4,)
|
||||
return (3, file) # 其他文件按名字排序
|
||||
|
||||
def get_prompts(self, join_with: str = "\n\n") -> str:
|
||||
"""获取所有提示词,按顺序拼接(identity → 技能文件 → soul → status)"""
|
||||
if not self._prompts:
|
||||
return ""
|
||||
sorted_prompts = sorted(self._prompts, key=self._prompt_sort_key)
|
||||
contents = [p['content'] for p in sorted_prompts]
|
||||
return join_with.join(contents)
|
||||
|
||||
def get_sorted_prompts_list(self) -> List[Dict]:
|
||||
"""获取按正确顺序排列的提示词列表(用于 system 消息构建)。"""
|
||||
return sorted(self._prompts, key=self._prompt_sort_key)
|
||||
|
||||
def get_prompts_list(self) -> List[Dict]:
|
||||
"""获取所有提示词列表(原始加载顺序,已弃用,请用 get_sorted_prompts_list)。"""
|
||||
return self._prompts.copy()
|
||||
|
||||
def get_skill_prompts(self) -> List[Dict]:
|
||||
"""获取所有基础技能提示词(N_ 开头的文件)。"""
|
||||
return [p for p in self._prompts if self._skill_sort_key(p['file']) is not None]
|
||||
|
||||
def reload(self):
|
||||
"""重新加载所有提示词"""
|
||||
self._prompts = []
|
||||
self._load_prompts()
|
||||
return len(self._prompts)
|
||||
|
||||
def get_prompt_count(self) -> int:
|
||||
return len(self._prompts)
|
||||
|
||||
def get_prompts_dir_path(self) -> str:
|
||||
return self._prompts_dir
|
||||
@@ -0,0 +1,230 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
subsystems.skill_manager - 扩展技能管理器
|
||||
|
||||
抽离自 aibox.py,收集并管理 ~/.aibox/skills/ 目录下的技能文件。
|
||||
|
||||
支持两种技能格式:
|
||||
1. 可执行脚本/程序:执行 --help,第一行为简介
|
||||
2. Markdown 文件 (.md):第一行为简介(# 标题或普通文本),其余为技能内容
|
||||
|
||||
依赖:common (constants, Logger) + 标准库
|
||||
不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from common import constants, Logger
|
||||
|
||||
|
||||
class SkillManager:
|
||||
"""扩展技能管理器 - 收集并管理 ~/.aibox/skills/ 目录下的技能文件。
|
||||
|
||||
支持两种技能格式:
|
||||
1. 可执行脚本/程序:执行 --help,第一行为简介
|
||||
2. Markdown 文件 (.md):第一行为简介(# 标题或普通文本),其余为技能内容
|
||||
"""
|
||||
|
||||
def __init__(self, logger: Optional[Logger] = None):
|
||||
self.logger = logger or Logger("SkillManager")
|
||||
# name -> {path, file, summary, content, is_markdown}
|
||||
self.skills: Dict[str, Dict] = {}
|
||||
self._load_skills()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 内部辅助
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _run_help(self, skill_path: str) -> str:
|
||||
"""执行 skill_path --help 并返回完整输出(stdout 优先,否则 stderr)。"""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[skill_path, "--help"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=10,
|
||||
env=os.environ.copy()
|
||||
)
|
||||
output = result.stdout.strip()
|
||||
if not output:
|
||||
output = result.stderr.strip()
|
||||
return output
|
||||
except Exception as e:
|
||||
self.logger.warning(f"运行 {skill_path} --help 失败: {e}")
|
||||
return ""
|
||||
|
||||
def _read_markdown_skill(self, skill_path: str) -> Tuple[str, str]:
|
||||
"""读取 Markdown 技能文件,返回 (简介, 完整内容)"""
|
||||
try:
|
||||
with open(skill_path, 'r', encoding='utf-8') as f:
|
||||
content = f.read().strip()
|
||||
|
||||
if not content:
|
||||
return "(空文件)", ""
|
||||
|
||||
# 提取第一行非空内容作为简介
|
||||
lines = content.splitlines()
|
||||
summary = ""
|
||||
for line in lines:
|
||||
stripped = line.strip()
|
||||
if stripped:
|
||||
# 如果是 Markdown 标题,去掉 # 前缀
|
||||
if stripped.startswith('#'):
|
||||
summary = stripped.lstrip('#').strip()
|
||||
else:
|
||||
summary = stripped
|
||||
break
|
||||
|
||||
if not summary:
|
||||
summary = os.path.basename(skill_path).replace('.md', '')
|
||||
|
||||
return summary, content
|
||||
|
||||
except Exception as e:
|
||||
self.logger.warning(f"读取 Markdown 技能文件失败 {skill_path}: {e}")
|
||||
return "(读取失败)", ""
|
||||
|
||||
@staticmethod
|
||||
def _first_nonempty_line(text: str) -> str:
|
||||
"""返回文本中第一行非空内容。"""
|
||||
for line in text.splitlines():
|
||||
stripped = line.strip()
|
||||
if stripped:
|
||||
return stripped
|
||||
return ""
|
||||
|
||||
def _is_markdown_file(self, filename: str) -> bool:
|
||||
"""判断是否为 Markdown 文件"""
|
||||
return filename.lower().endswith('.md')
|
||||
|
||||
def _is_executable(self, filepath: str) -> bool:
|
||||
"""检查文件是否可执行"""
|
||||
if os.name == 'nt': # Windows
|
||||
# Windows 下检查扩展名
|
||||
executable_extensions = {'.exe', '.bat', '.cmd', '.ps1', '.py', '.sh'}
|
||||
ext = os.path.splitext(filepath)[1].lower()
|
||||
return ext in executable_extensions
|
||||
else: # Unix/Linux/Mac
|
||||
return os.access(filepath, os.X_OK)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 加载 / 重载
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _load_skills(self):
|
||||
"""扫描技能目录,加载所有技能文件(可执行文件或 Markdown 文件)。"""
|
||||
skills_dir = constants.SKILLS_DIR
|
||||
self.skills = {}
|
||||
|
||||
if not os.path.exists(skills_dir):
|
||||
self.logger.info(f"技能目录不存在: {skills_dir}")
|
||||
return
|
||||
|
||||
skill_files = [
|
||||
f for f in os.listdir(skills_dir)
|
||||
if os.path.isfile(os.path.join(skills_dir, f))
|
||||
and not f.startswith('_')
|
||||
and not f.startswith('.')
|
||||
]
|
||||
|
||||
for skill_file in skill_files:
|
||||
skill_path = os.path.join(skills_dir, skill_file)
|
||||
skill_name = os.path.splitext(skill_file)[0]
|
||||
is_markdown = self._is_markdown_file(skill_file)
|
||||
|
||||
if is_markdown:
|
||||
# Markdown 格式技能:直接读取内容
|
||||
summary, content = self._read_markdown_skill(skill_path)
|
||||
if content:
|
||||
self.skills[skill_name] = {
|
||||
"path": skill_path,
|
||||
"file": skill_file,
|
||||
"summary": summary,
|
||||
"content": content,
|
||||
"is_markdown": True,
|
||||
"is_executable": False
|
||||
}
|
||||
self.logger.info(f"已加载 Markdown 技能: {skill_name} — {summary}")
|
||||
else:
|
||||
self.logger.warning(f"Markdown 技能文件为空(跳过): {skill_file}")
|
||||
|
||||
else:
|
||||
# 可执行文件格式技能:执行 --help 获取简介
|
||||
if not self._is_executable(skill_path):
|
||||
self.logger.warning(f"技能文件不可执行(跳过): {skill_file}")
|
||||
continue
|
||||
|
||||
help_output = self._run_help(skill_path)
|
||||
summary = self._first_nonempty_line(help_output) or f"扩展技能: {skill_name}"
|
||||
|
||||
self.skills[skill_name] = {
|
||||
"path": skill_path,
|
||||
"file": skill_file,
|
||||
"summary": summary,
|
||||
"content": help_output, # 完整的 --help 输出作为内容
|
||||
"is_markdown": False,
|
||||
"is_executable": True
|
||||
}
|
||||
self.logger.info(f"已加载可执行技能: {skill_name} — {summary}")
|
||||
|
||||
self.logger.info(f"共加载 {len(self.skills)} 个技能({len([s for s in self.skills.values() if s['is_markdown']])} 个 Markdown 技能)")
|
||||
|
||||
def reload(self) -> int:
|
||||
"""重新扫描技能目录,返回加载的技能数量。"""
|
||||
self._load_skills()
|
||||
return len(self.skills)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 查询
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def search(self, keyword: str) -> List[Dict]:
|
||||
"""在技能名称、简介和内容中搜索关键字,返回匹配项列表。"""
|
||||
kw = keyword.lower()
|
||||
results = []
|
||||
for name, info in self.skills.items():
|
||||
match = (
|
||||
kw in name.lower() or
|
||||
kw in info["summary"].lower() or
|
||||
kw in info.get("content", "").lower()
|
||||
)
|
||||
if match:
|
||||
results.append({
|
||||
"name": name,
|
||||
"path": info["path"],
|
||||
"summary": info["summary"],
|
||||
"content": info.get("content", ""),
|
||||
"is_markdown": info.get("is_markdown", False),
|
||||
"is_executable": info.get("is_executable", False)
|
||||
})
|
||||
return results
|
||||
|
||||
def get_skill_content(self, skill_name: str) -> Optional[str]:
|
||||
"""获取技能的内容(Markdown 技能返回完整内容,可执行技能返回 --help 输出)"""
|
||||
info = self.skills.get(skill_name)
|
||||
if info:
|
||||
return info.get("content")
|
||||
return None
|
||||
|
||||
def get_skill(self, skill_name: str) -> Optional[Dict]:
|
||||
"""获取单个技能的完整信息"""
|
||||
return self.skills.get(skill_name)
|
||||
|
||||
def get_all(self) -> List[Dict]:
|
||||
"""返回所有技能的列表(含名称、路径、简介和内容)。"""
|
||||
return [
|
||||
{
|
||||
"name": name,
|
||||
"path": info["path"],
|
||||
"summary": info["summary"],
|
||||
"content": info.get("content", ""),
|
||||
"is_markdown": info.get("is_markdown", False),
|
||||
"is_executable": info.get("is_executable", False)
|
||||
}
|
||||
for name, info in self.skills.items()
|
||||
]
|
||||
|
||||
def count(self) -> int:
|
||||
return len(self.skills)
|
||||
Whitespace-only changes.
@@ -0,0 +1,21 @@
|
||||
"""pytest 公共配置:把 aibox 目录加入 sys.path,桩掉可选第三方依赖。
|
||||
|
||||
aibox 主模块 aibox.py 依赖 websockets,但 tools/basic_tools.py 仅依赖
|
||||
common + 标准库,所以单测 basic_tools 不需要真实 websockets——桩掉即可
|
||||
让 `from tools import ...` 间接 import aibox.py 时不出错(虽然本测试
|
||||
不直接 import aibox,但保持环境干净)。
|
||||
"""
|
||||
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
# 把 aibox 目录加入 sys.path(无需安装为包即可 import tools/common)
|
||||
AIBOX_DIR = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(AIBOX_DIR))
|
||||
|
||||
# 桩掉 websockets(aibox.py 顶部 import websockets,测试环境可能未安装)
|
||||
if "websockets" not in sys.modules:
|
||||
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||
|
||||
import pytest # noqa: E402
|
||||
@@ -0,0 +1,327 @@
|
||||
"""tools.basic_tools 单元测试
|
||||
|
||||
覆盖 A 批抽离的 7 个工具类:
|
||||
- FileToolBase : 路径安全检查
|
||||
- ReadFileTool : 读文件指引
|
||||
- WriteFileTool : 写文件指引
|
||||
- ListFilesTool : 列目录指引
|
||||
- SimpleWebCrawlerTool : 网页访问指引
|
||||
- GetCurrentTimeTool : 时间获取
|
||||
- CalculatorTool : 数学计算
|
||||
|
||||
设计原则:
|
||||
- 不依赖网络/文件系统副作用(除 FileToolBase 路径检查用临时文件)
|
||||
- 不依赖 DeepSeek API、websockets 等外部服务
|
||||
- 验证工具接口契约(name/description/parameters/execute 返回类型)
|
||||
- 验证命令行指引类工具的核心语义(返回包含正确命令关键词 + run_command 提示)
|
||||
"""
|
||||
|
||||
import os
|
||||
import queue
|
||||
import tempfile
|
||||
from datetime import datetime
|
||||
|
||||
import pytest
|
||||
|
||||
from tools import (
|
||||
FileToolBase,
|
||||
ReadFileTool,
|
||||
WriteFileTool,
|
||||
ListFilesTool,
|
||||
SimpleWebCrawlerTool,
|
||||
GetCurrentTimeTool,
|
||||
CalculatorTool,
|
||||
)
|
||||
|
||||
|
||||
# ==================== 通用接口契约测试 ====================
|
||||
|
||||
# (工具类, 期望 get_name 返回值, 必填参数键列表)
|
||||
TOOL_CASES = [
|
||||
(ReadFileTool, "read_file", ["filepath"]),
|
||||
(WriteFileTool, "write_file", ["filepath", "content"]),
|
||||
(ListFilesTool, "list_files", []),
|
||||
(SimpleWebCrawlerTool, "web", ["action", "url"]),
|
||||
(GetCurrentTimeTool, "get_current_time", []),
|
||||
(CalculatorTool, "calculator", ["a", "b", "operation"]),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tool_cls,expected_name,required_keys", TOOL_CASES)
|
||||
def test_tool_interface_contract(tool_cls, expected_name, required_keys):
|
||||
"""所有工具必须满足接口契约:name/描述非空、parameters 是合法 schema、必填字段在 properties 中。"""
|
||||
tool = tool_cls()
|
||||
assert tool.get_name() == expected_name
|
||||
|
||||
desc = tool.get_description()
|
||||
assert isinstance(desc, str) and len(desc) > 0
|
||||
|
||||
params = tool.get_parameters()
|
||||
assert isinstance(params, dict)
|
||||
assert params.get("type") == "object"
|
||||
props = params.get("properties", {})
|
||||
assert isinstance(props, dict) and len(props) > 0
|
||||
for k in required_keys:
|
||||
assert k in props, f"必填参数 {k} 未出现在 properties 中"
|
||||
# required 中声明的字段必须在 properties 中
|
||||
for k in params.get("required", []):
|
||||
assert k in props
|
||||
|
||||
|
||||
# ==================== GetCurrentTimeTool ====================
|
||||
|
||||
def test_get_current_time_full_format():
|
||||
tool = GetCurrentTimeTool()
|
||||
r = tool.execute(format="full")
|
||||
assert "当前时间" in r
|
||||
# 应包含 YYYY-MM-DD HH:MM:SS 形式
|
||||
assert datetime.now().strftime("%Y-%m-%d") in r
|
||||
|
||||
|
||||
def test_get_current_time_date_format():
|
||||
tool = GetCurrentTimeTool()
|
||||
r = tool.execute(format="date")
|
||||
assert datetime.now().strftime("%Y-%m-%d") in r
|
||||
|
||||
|
||||
def test_get_current_time_time_format():
|
||||
tool = GetCurrentTimeTool()
|
||||
r = tool.execute(format="time")
|
||||
assert datetime.now().strftime("%H:%M") in r
|
||||
|
||||
|
||||
def test_get_current_time_streaming_emits_complete():
|
||||
tool = GetCurrentTimeTool()
|
||||
q: queue.Queue = queue.Queue()
|
||||
tool.execute_streaming(q, format="date")
|
||||
msgs = []
|
||||
while not q.empty():
|
||||
msgs.append(q.get_nowait())
|
||||
# 至少应有一个 complete 事件
|
||||
types = [m[0] for m in msgs]
|
||||
assert "complete" in types
|
||||
# complete 事件内容非空
|
||||
complete_msgs = [m[1] for m in msgs if m[0] == "complete"]
|
||||
assert complete_msgs and len(complete_msgs[0]) > 0
|
||||
|
||||
|
||||
# ==================== CalculatorTool ====================
|
||||
|
||||
def test_calculator_add():
|
||||
r = CalculatorTool().execute(a=3, b=4, operation="add")
|
||||
assert "7" in r and "计算成功" in r
|
||||
|
||||
|
||||
def test_calculator_subtract():
|
||||
r = CalculatorTool().execute(a=10, b=4, operation="subtract")
|
||||
assert "6" in r
|
||||
|
||||
|
||||
def test_calculator_multiply():
|
||||
r = CalculatorTool().execute(a=6, b=7, operation="multiply")
|
||||
assert "42" in r
|
||||
|
||||
|
||||
def test_calculator_divide():
|
||||
r = CalculatorTool().execute(a=20, b=4, operation="divide")
|
||||
assert "5" in r
|
||||
|
||||
|
||||
def test_calculator_divide_by_zero():
|
||||
r = CalculatorTool().execute(a=1, b=0, operation="divide")
|
||||
assert "除数不能为0" in r
|
||||
|
||||
|
||||
def test_calculator_unknown_operation():
|
||||
r = CalculatorTool().execute(a=1, b=1, operation="power")
|
||||
assert "未知运算" in r
|
||||
|
||||
|
||||
# ==================== ReadFileTool ====================
|
||||
|
||||
def test_read_file_head_returns_command_prompt():
|
||||
"""head 模式应返回含 head 命令 + run_command 指引的 markdown。"""
|
||||
r = ReadFileTool().execute(filepath="aibox.py", head=True, max_lines=10)
|
||||
assert "head -n 10" in r
|
||||
assert "run_command" in r
|
||||
|
||||
|
||||
def test_read_file_tail_returns_command_prompt():
|
||||
r = ReadFileTool().execute(filepath="aibox.py", tail=True, max_lines=20)
|
||||
assert "tail -n 20" in r
|
||||
assert "run_command" in r
|
||||
|
||||
|
||||
def test_read_file_range_returns_sed_command():
|
||||
r = ReadFileTool().execute(filepath="aibox.py", start_line=5, end_line=15)
|
||||
assert "sed -n '5,15p'" in r
|
||||
assert "run_command" in r
|
||||
|
||||
|
||||
def test_read_file_max_lines_cap_at_500():
|
||||
"""max_lines 应被钳制到 500 防止超大范围。"""
|
||||
r = ReadFileTool().execute(filepath="aibox.py", max_lines=99999)
|
||||
assert "sed -n '1,500p'" in r
|
||||
|
||||
|
||||
def test_read_file_missing_filepath_returns_error():
|
||||
r = ReadFileTool().execute(filepath=None)
|
||||
assert "❌" in r and "filepath" in r
|
||||
|
||||
|
||||
def test_read_file_nonexistent_returns_error():
|
||||
r = ReadFileTool().execute(filepath="/no/such/file_xyz.txt", head=True)
|
||||
assert "❌" in r
|
||||
|
||||
|
||||
# ==================== WriteFileTool ====================
|
||||
|
||||
def test_write_file_overwrite_mode():
|
||||
r = WriteFileTool().execute(filepath="output.txt", content="hello")
|
||||
assert ">" in r and "run_command" in r
|
||||
# 覆盖模式应使用单 > 重定向
|
||||
assert "cat > " in r
|
||||
|
||||
|
||||
def test_write_file_append_mode():
|
||||
r = WriteFileTool().execute(filepath="output.txt", content="more", mode="a")
|
||||
assert "cat >> " in r
|
||||
assert "run_command" in r
|
||||
|
||||
|
||||
def test_write_file_heredoc_uses_eof_marker():
|
||||
r = WriteFileTool().execute(filepath="out.txt", content="data")
|
||||
assert "AIBOX_EOF" in r
|
||||
|
||||
|
||||
def test_write_file_missing_content_returns_error():
|
||||
r = WriteFileTool().execute(filepath="out.txt", content=None)
|
||||
assert "❌" in r and "content" in r
|
||||
|
||||
|
||||
def test_write_file_creates_parent_dir_command():
|
||||
"""写深路径时应在命令中包含 mkdir -p。"""
|
||||
r = WriteFileTool().execute(filepath="sub/dir/out.txt", content="x")
|
||||
assert "mkdir -p" in r
|
||||
|
||||
|
||||
# ==================== ListFilesTool ====================
|
||||
|
||||
def test_list_files_default_no_hidden_flag():
|
||||
r = ListFilesTool().execute(directory=".")
|
||||
assert "ls -l " in r
|
||||
assert "run_command" in r
|
||||
|
||||
|
||||
def test_list_files_show_hidden_flag():
|
||||
r = ListFilesTool().execute(directory=".", show_hidden=True)
|
||||
assert "ls -la " in r
|
||||
|
||||
|
||||
def test_list_files_nonexistent_dir_returns_error():
|
||||
r = ListFilesTool().execute(directory="/no/such/dir_xyz")
|
||||
assert "❌" in r
|
||||
|
||||
|
||||
def test_list_files_path_is_file_returns_error(tmp_path):
|
||||
f = tmp_path / "afile.txt"
|
||||
f.write_text("x")
|
||||
r = ListFilesTool().execute(directory=str(f))
|
||||
assert "❌" in r and "不是目录" in r
|
||||
|
||||
|
||||
# ==================== SimpleWebCrawlerTool ====================
|
||||
|
||||
def test_web_get_action_returns_curl_prompt():
|
||||
r = SimpleWebCrawlerTool().execute(action="get", url="https://example.com")
|
||||
assert "curl -sL" in r
|
||||
assert "run_command" in r
|
||||
assert "https://example.com" in r
|
||||
|
||||
|
||||
def test_web_title_action_uses_grep():
|
||||
r = SimpleWebCrawlerTool().execute(action="title", url="https://example.com")
|
||||
assert "curl" in r and "<title>" in r
|
||||
|
||||
|
||||
def test_web_links_action_uses_grep():
|
||||
r = SimpleWebCrawlerTool().execute(action="links", url="https://example.com")
|
||||
assert "href" in r and "curl" in r
|
||||
|
||||
|
||||
def test_web_text_action_uses_sed():
|
||||
r = SimpleWebCrawlerTool().execute(action="text", url="https://example.com")
|
||||
assert "sed" in r and "curl" in r
|
||||
|
||||
|
||||
def test_web_search_requires_keyword():
|
||||
r = SimpleWebCrawlerTool().execute(action="search", url="https://example.com")
|
||||
assert "❌" in r and "keyword" in r
|
||||
|
||||
|
||||
def test_web_search_with_keyword_uses_grep_c():
|
||||
r = SimpleWebCrawlerTool().execute(
|
||||
action="search", url="https://example.com", keyword="hello"
|
||||
)
|
||||
assert "grep -c" in r and "hello" in r
|
||||
|
||||
|
||||
def test_web_unknown_action_returns_error():
|
||||
r = SimpleWebCrawlerTool().execute(action="download", url="https://example.com")
|
||||
assert "❌" in r and "未知操作" in r
|
||||
|
||||
|
||||
# ==================== FileToolBase ====================
|
||||
|
||||
# FileToolBase 是抽象基类(继承自 Tool 的抽象方法未实现),无法直接实例化。
|
||||
# 用一个最小具体子类来测试基类提供的 normalize_path / check_path_safety 方法。
|
||||
class _ConcreteFileTool(FileToolBase):
|
||||
def get_name(self): return "_test"
|
||||
def get_description(self): return "test"
|
||||
def get_parameters(self): return {"type": "object", "properties": {}}
|
||||
def execute(self, **kwargs): return ""
|
||||
|
||||
|
||||
def test_filetoolbase_normalize_path_expands_user():
|
||||
base = _ConcreteFileTool()
|
||||
p = base.normalize_path("~/foo")
|
||||
assert "~" not in p
|
||||
assert os.path.isabs(p)
|
||||
|
||||
|
||||
def test_filetoolbase_check_path_safety_existing_file(tmp_path):
|
||||
f = tmp_path / "ok.txt"
|
||||
f.write_text("hi")
|
||||
base = _ConcreteFileTool()
|
||||
ok, abs_path, err = base.check_path_safety(str(f), check_exists=True)
|
||||
assert ok is True
|
||||
assert abs_path == str(f)
|
||||
assert err == ""
|
||||
|
||||
|
||||
def test_filetoolbase_check_path_safety_nonexistent_with_check(tmp_path):
|
||||
base = _ConcreteFileTool()
|
||||
ok, _, err = base.check_path_safety(str(tmp_path / "missing.txt"), check_exists=True)
|
||||
assert ok is False
|
||||
assert "文件不存在" in err
|
||||
|
||||
|
||||
def test_filetoolbase_check_path_safety_oversize_file(tmp_path):
|
||||
"""超大文件应被拒绝。"""
|
||||
f = tmp_path / "big.bin"
|
||||
# 写入略超 MAX_FILE_SIZE 的内容
|
||||
from common import constants
|
||||
f.write_bytes(b"x" * (constants.MAX_FILE_SIZE + 1))
|
||||
base = _ConcreteFileTool()
|
||||
ok, _, err = base.check_path_safety(str(f))
|
||||
assert ok is False
|
||||
assert "文件太大" in err
|
||||
|
||||
|
||||
def test_filetoolbase_is_subclass_of_tool():
|
||||
"""FileToolBase 应是 Tool 的子类,且三个文件工具都继承它。"""
|
||||
from tools import Tool
|
||||
assert issubclass(FileToolBase, Tool)
|
||||
assert issubclass(ReadFileTool, FileToolBase)
|
||||
assert issubclass(WriteFileTool, FileToolBase)
|
||||
assert issubclass(ListFilesTool, FileToolBase)
|
||||
@@ -0,0 +1,231 @@
|
||||
"""subsystems.db 单元测试
|
||||
|
||||
覆盖:
|
||||
- MemoryDatabase: 增删查 / FTS5 全文检索 / 批量删除
|
||||
- MemoDatabase: 增删查 / 待执行 / 重复任务 / 即时任务 / 统计
|
||||
每个测试用独立临时 sqlite 文件,互不污染。
|
||||
"""
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
import pytest
|
||||
|
||||
from subsystems import MemoryDatabase, MemoDatabase
|
||||
from common import constants
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mem_db(tmp_path):
|
||||
"""独立临时文件的 MemoryDatabase。"""
|
||||
db_path = str(tmp_path / "test_memory.db")
|
||||
return MemoryDatabase(db_path=db_path)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def memo_db(tmp_path):
|
||||
"""独立临时文件的 MemoDatabase。"""
|
||||
db_path = str(tmp_path / "test_memo.db")
|
||||
return MemoDatabase(db_path=db_path)
|
||||
|
||||
|
||||
# ==================== MemoryDatabase ====================
|
||||
|
||||
def test_memory_add_and_get(mem_db):
|
||||
mid = mem_db.add_memory(
|
||||
topic="测试主题", summary="测试摘要",
|
||||
content="详细内容", tags=["t1", "t2"]
|
||||
)
|
||||
assert isinstance(mid, int) and mid > 0
|
||||
|
||||
m = mem_db.get_memory_by_id(mid)
|
||||
assert m is not None
|
||||
assert m["topic"] == "测试主题"
|
||||
assert m["summary"] == "测试摘要"
|
||||
assert m["tags"] == ["t1", "t2"]
|
||||
|
||||
|
||||
def test_memory_search_by_keyword(mem_db):
|
||||
mem_db.add_memory(topic="Python learning", summary="learn Python basics", content="variables and functions")
|
||||
mem_db.add_memory(topic="Go learning", summary="learn Go basics", content="goroutine")
|
||||
|
||||
results = mem_db.search_memories(["Python"])
|
||||
assert len(results) == 1
|
||||
assert "Python" in results[0]["topic"]
|
||||
|
||||
|
||||
def test_memory_search_match_all(mem_db):
|
||||
mem_db.add_memory(topic="Python database", summary="sqlite usage", content="FTS5 search")
|
||||
mem_db.add_memory(topic="Python network", summary="requests usage", content="HTTP request")
|
||||
|
||||
# match_all=True:两个关键词都要命中
|
||||
results = mem_db.search_memories(["Python", "sqlite"], match_all=True)
|
||||
assert len(results) == 1
|
||||
assert "database" in results[0]["topic"]
|
||||
|
||||
# match_all=False:任一关键词命中
|
||||
results = mem_db.search_memories(["Python", "sqlite"], match_all=False)
|
||||
assert len(results) == 2
|
||||
|
||||
|
||||
def test_memory_search_empty_keywords_returns_empty(mem_db):
|
||||
assert mem_db.search_memories([]) == []
|
||||
|
||||
|
||||
def test_memory_delete(mem_db):
|
||||
mid = mem_db.add_memory(topic="待删除", summary="摘要")
|
||||
assert mem_db.delete_memory(mid) is True
|
||||
assert mem_db.get_memory_by_id(mid) is None
|
||||
# 再次删除应返回 False
|
||||
assert mem_db.delete_memory(mid) is False
|
||||
|
||||
|
||||
def test_memory_batch_delete(mem_db):
|
||||
ids = [mem_db.add_memory(topic=f"t{i}", summary=f"s{i}") for i in range(3)]
|
||||
deleted = mem_db.delete_memories_by_ids(ids)
|
||||
assert deleted == 3
|
||||
for mid in ids:
|
||||
assert mem_db.get_memory_by_id(mid) is None
|
||||
|
||||
|
||||
def test_memory_batch_delete_empty_list(mem_db):
|
||||
assert mem_db.delete_memories_by_ids([]) == 0
|
||||
|
||||
|
||||
def test_memory_get_recent(mem_db):
|
||||
for i in range(5):
|
||||
mem_db.add_memory(topic=f"recent_{i}", summary=f"s{i}")
|
||||
recent = mem_db.get_recent_memories(limit=3)
|
||||
assert len(recent) == 3
|
||||
|
||||
|
||||
def test_memory_nonexistent_id_returns_none(mem_db):
|
||||
assert mem_db.get_memory_by_id(99999) is None
|
||||
|
||||
|
||||
# ==================== MemoDatabase ====================
|
||||
|
||||
def test_memo_add_and_get(memo_db):
|
||||
reminder = (datetime.now() + timedelta(days=1)).strftime(constants.DATETIME_FORMAT)
|
||||
mid = memo_db.add_memo(title="测试任务", content="内容", reminder_time=reminder)
|
||||
assert isinstance(mid, int) and mid > 0
|
||||
|
||||
m = memo_db.get_memo(mid)
|
||||
assert m is not None
|
||||
assert m["title"] == "测试任务"
|
||||
assert m["is_completed"] == 0
|
||||
|
||||
|
||||
def test_memo_pending_filter_by_time(memo_db):
|
||||
"""未到期的任务不应出现在 pending 列表中。"""
|
||||
future = (datetime.now() + timedelta(days=1)).strftime(constants.DATETIME_FORMAT)
|
||||
past = (datetime.now() - timedelta(days=1)).strftime(constants.DATETIME_FORMAT)
|
||||
|
||||
memo_db.add_memo(title="未来任务", reminder_time=future)
|
||||
memo_db.add_memo(title="过期任务", reminder_time=past)
|
||||
|
||||
pending = memo_db.get_pending_memos()
|
||||
titles = [m["title"] for m in pending]
|
||||
assert "过期任务" in titles
|
||||
assert "未来任务" not in titles
|
||||
|
||||
|
||||
def test_memo_complete(memo_db):
|
||||
reminder = (datetime.now() + timedelta(days=1)).strftime(constants.DATETIME_FORMAT)
|
||||
mid = memo_db.add_memo(title="一次性任务", reminder_time=reminder)
|
||||
assert memo_db.complete_memo(mid) is True
|
||||
|
||||
m = memo_db.get_memo(mid)
|
||||
assert m["is_completed"] == 1
|
||||
|
||||
|
||||
def test_memo_complete_nonexistent_returns_false(memo_db):
|
||||
assert memo_db.complete_memo(99999) is False
|
||||
|
||||
|
||||
def test_memo_immediate(memo_db):
|
||||
"""即时任务应被 get_immediate_memos 返回。"""
|
||||
mid = memo_db.add_memo(title="即时任务", is_immediate=True)
|
||||
imm = memo_db.get_immediate_memos()
|
||||
assert any(m["id"] == mid for m in imm)
|
||||
|
||||
|
||||
def test_memo_mark_as_triggered_completes_one_shot(memo_db):
|
||||
"""一次性任务触发后应标记完成。"""
|
||||
past = (datetime.now() - timedelta(days=1)).strftime(constants.DATETIME_FORMAT)
|
||||
mid = memo_db.add_memo(title="一次性", reminder_time=past)
|
||||
nxt = memo_db.mark_as_triggered(mid)
|
||||
assert nxt is None # 一次性任务无下次提醒
|
||||
m = memo_db.get_memo(mid)
|
||||
assert m["is_completed"] == 1
|
||||
assert m["trigger_count"] == 1
|
||||
|
||||
|
||||
def test_memo_mark_as_triggered_daily_repeat(memo_db):
|
||||
"""daily 重复任务触发后应计算下次提醒时间。"""
|
||||
past = (datetime.now() - timedelta(days=1)).strftime(constants.DATETIME_FORMAT)
|
||||
mid = memo_db.add_memo(title="每日任务", reminder_time=past, repeat_type="daily")
|
||||
nxt = memo_db.mark_as_triggered(mid)
|
||||
assert nxt is not None
|
||||
# 下次应在未来
|
||||
next_dt = datetime.strptime(nxt, constants.DATETIME_FORMAT)
|
||||
assert next_dt > datetime.now() - timedelta(minutes=1)
|
||||
m = memo_db.get_memo(mid)
|
||||
assert m["is_completed"] == 0 # 重复任务不完成
|
||||
|
||||
|
||||
def test_memo_stats(memo_db):
|
||||
past = (datetime.now() - timedelta(days=1)).strftime(constants.DATETIME_FORMAT)
|
||||
future = (datetime.now() + timedelta(days=1)).strftime(constants.DATETIME_FORMAT)
|
||||
memo_db.add_memo(title="过期", reminder_time=past)
|
||||
memo_db.add_memo(title="未来", reminder_time=future)
|
||||
memo_db.add_memo(title="即时", is_immediate=True)
|
||||
|
||||
stats = memo_db.get_memo_stats()
|
||||
assert stats["total"] == 3
|
||||
assert stats["pending"] == 3
|
||||
assert stats["overdue"] == 1 # 只有过期且非即时的
|
||||
assert stats["immediate"] == 1
|
||||
|
||||
|
||||
def test_memo_get_all_excludes_completed_by_default(memo_db):
|
||||
reminder = (datetime.now() + timedelta(days=1)).strftime(constants.DATETIME_FORMAT)
|
||||
mid = memo_db.add_memo(title="待完成", reminder_time=reminder)
|
||||
memo_db.complete_memo(mid)
|
||||
|
||||
all_memos = memo_db.get_all_memos(include_completed=False)
|
||||
assert all(m["is_completed"] == 0 for m in all_memos)
|
||||
|
||||
all_with_completed = memo_db.get_all_memos(include_completed=True)
|
||||
assert any(m["id"] == mid for m in all_with_completed)
|
||||
|
||||
|
||||
def test_memo_get_pending_tasks_for_user(memo_db):
|
||||
past = (datetime.now() - timedelta(days=1)).strftime(constants.DATETIME_FORMAT)
|
||||
memo_db.add_memo(title="用户任务", reminder_time=past)
|
||||
|
||||
tasks = memo_db.get_pending_tasks_for_user()
|
||||
assert len(tasks) >= 1
|
||||
assert all("id" in t and "title" in t and "is_overdue" in t for t in tasks)
|
||||
assert any(t["title"] == "用户任务" and t["is_overdue"] for t in tasks)
|
||||
|
||||
|
||||
def test_memo_add_user_task_disabled(monkeypatch, memo_db):
|
||||
"""当 MEMO_ALLOW_USER_TASKS=False 时应抛异常。"""
|
||||
monkeypatch.setattr(type(constants), "MEMO_ALLOW_USER_TASKS",
|
||||
property(lambda self: False))
|
||||
with pytest.raises(Exception, match="禁用"):
|
||||
memo_db.add_user_task(title="x")
|
||||
|
||||
|
||||
def test_memo_custom_repeat_minutes(memo_db):
|
||||
"""custom 重复 + minutes 单位应正确计算下次。"""
|
||||
past = (datetime.now() - timedelta(minutes=1)).strftime(constants.DATETIME_FORMAT)
|
||||
mid = memo_db.add_memo(
|
||||
title="每5分钟", reminder_time=past,
|
||||
repeat_type="custom", repeat_interval=5, repeat_interval_unit="minutes"
|
||||
)
|
||||
nxt = memo_db.mark_as_triggered(mid)
|
||||
assert nxt is not None
|
||||
next_dt = datetime.strptime(nxt, constants.DATETIME_FORMAT)
|
||||
# 下次约在 5 分钟后
|
||||
delta = next_dt - datetime.now()
|
||||
assert timedelta(minutes=4) < delta < timedelta(minutes=6)
|
||||
@@ -0,0 +1,188 @@
|
||||
"""subsystems.executors 单元测试
|
||||
|
||||
覆盖:
|
||||
- CommandExecutor 抽象基类契约
|
||||
- SubprocessCommandExecutor 命令执行/流式/超时/取消/输出限制
|
||||
"""
|
||||
import queue
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from subsystems import CommandExecutor, SubprocessCommandExecutor
|
||||
from common import constants
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def use_tmp_workspace(tmp_path, monkeypatch):
|
||||
"""每个测试把 WORKSPACE_DIR 重定向到临时目录(并预先创建),避免默认 ~/.aibox/workspace 不存在。"""
|
||||
ws = tmp_path / "ws"
|
||||
ws.mkdir()
|
||||
monkeypatch.setattr(type(constants), "WORKSPACE_DIR",
|
||||
property(lambda self: str(ws)))
|
||||
|
||||
|
||||
# ==================== 抽象基类契约 ====================
|
||||
|
||||
def test_command_executor_is_abstract():
|
||||
"""CommandExecutor 是 ABC,不能直接实例化。"""
|
||||
with pytest.raises(TypeError):
|
||||
CommandExecutor()
|
||||
|
||||
|
||||
def test_subprocess_executor_is_subclass():
|
||||
assert issubclass(SubprocessCommandExecutor, CommandExecutor)
|
||||
|
||||
|
||||
def test_executor_is_enabled_reflects_constants():
|
||||
"""is_enabled 应反映 constants.COMMANDS_ENABLED。"""
|
||||
from common import constants
|
||||
ex = SubprocessCommandExecutor()
|
||||
assert ex.is_enabled() == constants.COMMANDS_ENABLED
|
||||
|
||||
|
||||
def test_executor_format_result_success():
|
||||
ex = SubprocessCommandExecutor()
|
||||
r = ex.format_result(True, "ok", {"k": "v"})
|
||||
assert r.startswith("✅ ok")
|
||||
assert "k: v" in r
|
||||
|
||||
|
||||
def test_executor_format_result_failure():
|
||||
ex = SubprocessCommandExecutor()
|
||||
r = ex.format_result(False, "fail")
|
||||
assert r.startswith("❌ fail")
|
||||
|
||||
|
||||
def test_executor_format_result_multiline_value():
|
||||
ex = SubprocessCommandExecutor()
|
||||
r = ex.format_result(True, "ok", {"输出": "line1\nline2"})
|
||||
assert "line1" in r and "line2" in r
|
||||
|
||||
|
||||
# ==================== 实际命令执行 ====================
|
||||
|
||||
def test_execute_echo_command():
|
||||
"""执行 echo 命令应返回成功结果。"""
|
||||
ex = SubprocessCommandExecutor()
|
||||
r = ex.execute("echo hello_test_123", timeout=5)
|
||||
assert "hello_test_123" in r
|
||||
assert "✅" in r
|
||||
|
||||
|
||||
def test_execute_command_returns_exit_code():
|
||||
ex = SubprocessCommandExecutor()
|
||||
r = ex.execute("exit 0", timeout=5)
|
||||
assert "返回码: 0" in r
|
||||
|
||||
|
||||
def test_execute_failing_command_returns_nonzero_code():
|
||||
ex = SubprocessCommandExecutor()
|
||||
r = ex.execute("exit 7", timeout=5)
|
||||
assert "返回码: 7" in r
|
||||
# 非零返回码应标记为失败
|
||||
assert "❌" in r
|
||||
|
||||
|
||||
def test_execute_timeout():
|
||||
"""超时应被检测并返回错误信息。"""
|
||||
ex = SubprocessCommandExecutor()
|
||||
# ping -t 在 Windows 上无限循环;用 ping -n 跨平台
|
||||
r = ex.execute("ping -n 100 127.0.0.1", timeout=2)
|
||||
assert "超时" in r
|
||||
|
||||
|
||||
def test_execute_unknown_command_returns_error():
|
||||
ex = SubprocessCommandExecutor()
|
||||
r = ex.execute("nonexistent_command_xyz_123", timeout=5)
|
||||
# shell=True 下会返回 127 或类似错误码,但不会抛异常
|
||||
assert "返回码" in r
|
||||
|
||||
|
||||
# ==================== 流式执行 ====================
|
||||
|
||||
def test_execute_streaming_emits_lines_and_complete():
|
||||
ex = SubprocessCommandExecutor()
|
||||
q: queue.Queue = queue.Queue()
|
||||
ex.execute_streaming("echo stream_line_test", q, timeout=5)
|
||||
|
||||
msgs = []
|
||||
while not q.empty():
|
||||
msgs.append(q.get_nowait())
|
||||
|
||||
types = [m[0] for m in msgs]
|
||||
assert "line" in types
|
||||
assert "complete" in types
|
||||
|
||||
complete_payload = [m[1] for m in msgs if m[0] == "complete"][0]
|
||||
assert "stream_line_test" in complete_payload
|
||||
|
||||
|
||||
def test_execute_streaming_emits_command_header():
|
||||
"""流式执行开头应输出命令/超时/工作目录等头部信息。"""
|
||||
ex = SubprocessCommandExecutor()
|
||||
q: queue.Queue = queue.Queue()
|
||||
ex.execute_streaming("echo hi", q, timeout=5)
|
||||
|
||||
lines = []
|
||||
while not q.empty():
|
||||
t, p = q.get_nowait()
|
||||
if t == "line":
|
||||
lines.append(p)
|
||||
|
||||
header = "".join(lines[:6])
|
||||
assert "执行命令" in header
|
||||
assert "超时" in header
|
||||
assert "工作目录" in header
|
||||
|
||||
|
||||
# ==================== 取消机制 ====================
|
||||
|
||||
def test_cancel_terminates_running_process():
|
||||
"""cancel() 在进程运行时调用应不抛异常并尝试 kill 进程。
|
||||
Windows 下 kill 后管道读取可能短暂阻塞,但 cancel 本身必须安全返回。
|
||||
线程设为 daemon,确保不会卡住测试进程退出。
|
||||
"""
|
||||
ex = SubprocessCommandExecutor()
|
||||
q: queue.Queue = queue.Queue()
|
||||
import threading
|
||||
t = threading.Thread(
|
||||
target=ex.execute_streaming,
|
||||
args=("ping -n 100 127.0.0.1", q),
|
||||
kwargs={"timeout": 60},
|
||||
daemon=True,
|
||||
)
|
||||
t.start()
|
||||
time.sleep(0.5) # 等进程启动
|
||||
# 关键断言:cancel 不抛异常(若抛异常测试即失败)
|
||||
ex.cancel()
|
||||
# 等待线程退出(Windows 下可能因管道阻塞超时,daemon 确保不卡死主进程)
|
||||
t.join(timeout=5)
|
||||
|
||||
|
||||
def test_cancel_when_no_process_running_is_noop():
|
||||
"""无进程运行时调用 cancel 应是空操作。"""
|
||||
ex = SubprocessCommandExecutor()
|
||||
ex.cancel() # 不应抛异常
|
||||
|
||||
|
||||
# ==================== 输出大小限制 ====================
|
||||
|
||||
def test_output_size_limit_truncates():
|
||||
"""输出超过 1MB 应被截断。"""
|
||||
ex = SubprocessCommandExecutor()
|
||||
# 生成约 2MB 输出
|
||||
r = ex.execute("python -c \"print('x'*2048576)\"", timeout=10)
|
||||
assert "截断" in r or "❌" in r
|
||||
|
||||
|
||||
# ==================== 禁用场景 ====================
|
||||
|
||||
def test_execute_when_disabled_returns_disabled_message(monkeypatch):
|
||||
"""当 COMMANDS_ENABLED=False 时应返回禁用提示。"""
|
||||
from common import constants
|
||||
monkeypatch.setattr(type(constants), "COMMANDS_ENABLED",
|
||||
property(lambda self: False))
|
||||
ex = SubprocessCommandExecutor()
|
||||
r = ex.execute("echo hi", timeout=5)
|
||||
assert "禁用" in r
|
||||
@@ -0,0 +1,190 @@
|
||||
"""subsystems.prompt_manager 单元测试
|
||||
|
||||
覆盖:
|
||||
- PromptManager 单例语义
|
||||
- 提示词加载/排序/拼接
|
||||
- identity/soul/status 读写
|
||||
- 从模板复制行为
|
||||
- reload 语义
|
||||
"""
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from subsystems import PromptManager
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def isolated_pm(tmp_path, monkeypatch):
|
||||
"""每个测试用独立 prompts 目录 + 重置单例,避免测试间状态污染。"""
|
||||
prompts_dir = tmp_path / "prompts"
|
||||
prompts_dir.mkdir()
|
||||
|
||||
# 重置单例状态
|
||||
PromptManager._instance = None
|
||||
PromptManager._prompts = []
|
||||
PromptManager._prompts_dir = None
|
||||
|
||||
# 让 ConfigManager 返回临时目录
|
||||
from common import ConfigManager
|
||||
original_get = ConfigManager.get
|
||||
|
||||
def mock_get(self, key, default=None):
|
||||
if key == "system.prompts_dir":
|
||||
return str(prompts_dir)
|
||||
return original_get(self, key, default)
|
||||
|
||||
monkeypatch.setattr(ConfigManager, "get", mock_get)
|
||||
|
||||
pm = PromptManager()
|
||||
yield pm
|
||||
|
||||
# 清理单例
|
||||
PromptManager._instance = None
|
||||
PromptManager._prompts = []
|
||||
PromptManager._prompts_dir = None
|
||||
|
||||
|
||||
# ==================== 单例 ====================
|
||||
|
||||
def test_prompt_manager_is_singleton(isolated_pm):
|
||||
pm2 = PromptManager()
|
||||
assert pm2 is isolated_pm
|
||||
|
||||
|
||||
def test_prompt_manager_resets_between_tests(isolated_pm):
|
||||
"""fixture 应重置单例,确保测试隔离。"""
|
||||
assert PromptManager._instance is isolated_pm
|
||||
|
||||
|
||||
# ==================== 加载与排序 ====================
|
||||
|
||||
def test_loads_template_files(isolated_pm):
|
||||
"""首次构造应从 aibox/prompts/ 模板复制核心文件。"""
|
||||
files = os.listdir(isolated_pm._prompts_dir)
|
||||
# 至少应有 00_IDENTITY / 99_SOUL / 100_STATUS
|
||||
assert "00_IDENTITY.md" in files
|
||||
assert "99_SOUL.md" in files
|
||||
assert "100_STATUS.md" in files
|
||||
|
||||
|
||||
def test_get_prompts_returns_nonempty_string(isolated_pm):
|
||||
s = isolated_pm.get_prompts()
|
||||
assert isinstance(s, str) and len(s) > 0
|
||||
|
||||
|
||||
def test_get_sorted_prompts_list_returns_list(isolated_pm):
|
||||
lst = isolated_pm.get_sorted_prompts_list()
|
||||
assert isinstance(lst, list)
|
||||
assert len(lst) > 0
|
||||
assert all("file" in p and "content" in p for p in lst)
|
||||
|
||||
|
||||
def test_prompt_sort_order_identity_first(isolated_pm):
|
||||
"""00_IDENTITY 应排在最前。"""
|
||||
lst = isolated_pm.get_sorted_prompts_list()
|
||||
assert lst[0]["file"] == "00_IDENTITY.md"
|
||||
|
||||
|
||||
def test_prompt_sort_order_status_last_or_near_last(isolated_pm):
|
||||
"""100_STATUS 应在 SOUL 之后。"""
|
||||
lst = isolated_pm.get_sorted_prompts_list()
|
||||
files = [p["file"] for p in lst]
|
||||
if "100_STATUS.md" in files and "99_SOUL.md" in files:
|
||||
assert files.index("100_STATUS.md") > files.index("99_SOUL.md")
|
||||
|
||||
|
||||
def test_skill_prompts_filtered_correctly(isolated_pm):
|
||||
"""N_ 开头的技能文件应被 get_skill_prompts 识别。"""
|
||||
skills = isolated_pm.get_skill_prompts()
|
||||
# 模板里至少有一个 1_xxx.md 技能文件
|
||||
assert len(skills) >= 0 # 至少不报错
|
||||
for p in skills:
|
||||
assert p["file"][0].isdigit()
|
||||
|
||||
|
||||
# ==================== 读写 identity/soul/status ====================
|
||||
|
||||
def test_update_and_get_identity(isolated_pm):
|
||||
ok = isolated_pm.update_identity("测试用户-张三")
|
||||
assert ok is True
|
||||
content = isolated_pm.get_identity()
|
||||
assert "测试用户-张三" in content
|
||||
assert "用户身份信息" in content
|
||||
|
||||
|
||||
def test_update_and_get_soul(isolated_pm):
|
||||
ok = isolated_pm.update_soul("我是测试智能体")
|
||||
assert ok is True
|
||||
content = isolated_pm.get_soul()
|
||||
assert "我是测试智能体" in content
|
||||
|
||||
|
||||
def test_update_and_get_status(isolated_pm):
|
||||
ok = isolated_pm.update_status("测试状态-初始化完成")
|
||||
assert ok is True
|
||||
content = isolated_pm.get_status()
|
||||
assert "测试状态-初始化完成" in content
|
||||
|
||||
|
||||
def test_get_identity_nonexistent_returns_fallback(monkeypatch, tmp_path):
|
||||
"""无 identity 文件时应返回兜底字符串。"""
|
||||
PromptManager._instance = None
|
||||
PromptManager._prompts = []
|
||||
empty_dir = tmp_path / "empty"
|
||||
empty_dir.mkdir()
|
||||
|
||||
from common import ConfigManager
|
||||
original_get = ConfigManager.get
|
||||
|
||||
def mock_get(self, key, default=None):
|
||||
if key == "system.prompts_dir":
|
||||
return str(empty_dir)
|
||||
return original_get(self, key, default)
|
||||
|
||||
monkeypatch.setattr(ConfigManager, "get", mock_get)
|
||||
pm = PromptManager()
|
||||
# 删除 identity 文件模拟不存在
|
||||
id_file = empty_dir / "00_IDENTITY.md"
|
||||
if id_file.exists():
|
||||
id_file.unlink()
|
||||
r = pm.get_identity()
|
||||
assert "尚不清楚" in r or "无法读取" in r
|
||||
|
||||
PromptManager._instance = None
|
||||
PromptManager._prompts = []
|
||||
PromptManager._prompts_dir = None
|
||||
|
||||
|
||||
# ==================== reload ====================
|
||||
|
||||
def test_reload_returns_count(isolated_pm):
|
||||
n = isolated_pm.reload()
|
||||
assert isinstance(n, int) and n >= 0
|
||||
|
||||
|
||||
def test_reload_picks_up_new_file(isolated_pm, tmp_path):
|
||||
"""reload 应能发现新增的提示词文件。"""
|
||||
before = isolated_pm.get_prompt_count()
|
||||
# 写入一个新的提示词文件
|
||||
new_file = os.path.join(isolated_pm._prompts_dir, "99_EXTRA_TEST.md")
|
||||
with open(new_file, "w", encoding="utf-8") as f:
|
||||
f.write("# 额外测试提示词\n内容")
|
||||
isolated_pm.reload()
|
||||
after = isolated_pm.get_prompt_count()
|
||||
assert after == before + 1
|
||||
|
||||
|
||||
# ==================== 模板不覆盖已存在文件 ====================
|
||||
|
||||
def test_existing_file_not_overwritten_by_template(isolated_pm):
|
||||
"""已存在的运行时文件不应被模板覆盖。"""
|
||||
# 修改 identity 文件
|
||||
isolated_pm.update_identity("我的自定义身份")
|
||||
custom_content = isolated_pm.get_identity()
|
||||
assert "我的自定义身份" in custom_content
|
||||
|
||||
# reload 触发 _ensure_from_template,但不应覆盖
|
||||
isolated_pm.reload()
|
||||
after = isolated_pm.get_identity()
|
||||
assert "我的自定义身份" in after
|
||||
@@ -0,0 +1,204 @@
|
||||
"""subsystems.skill_manager 单元测试
|
||||
|
||||
覆盖:
|
||||
- SkillManager 加载 Markdown 技能
|
||||
- 加载可执行技能(--help)
|
||||
- search / get_skill_content / get_all / reload / count
|
||||
每个测试用独立临时技能目录,互不污染。
|
||||
"""
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from subsystems import SkillManager
|
||||
from common import constants
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def isolated_skill_dir(tmp_path, monkeypatch):
|
||||
"""每个测试用独立技能目录 + monkeypatch constants.SKILLS_DIR。"""
|
||||
skills_dir = tmp_path / "skills"
|
||||
skills_dir.mkdir()
|
||||
monkeypatch.setattr(type(constants), "SKILLS_DIR",
|
||||
property(lambda self: str(skills_dir)))
|
||||
return skills_dir
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sm(isolated_skill_dir):
|
||||
"""构造 SkillManager(依赖 isolated_skill_dir fixture 注入目录)。"""
|
||||
return SkillManager()
|
||||
|
||||
|
||||
# ==================== 空目录 ====================
|
||||
|
||||
def test_empty_skills_dir(sm):
|
||||
assert sm.count() == 0
|
||||
assert sm.get_all() == []
|
||||
|
||||
|
||||
# ==================== Markdown 技能 ====================
|
||||
|
||||
def test_loads_markdown_skill(isolated_skill_dir, sm):
|
||||
"""Markdown 技能应被加载,简介取首行(# 标题去 # 前缀)。"""
|
||||
skill_file = isolated_skill_dir / "weather.md"
|
||||
skill_file.write_text("# 查询天气\n查询指定城市天气\n详细用法...", encoding="utf-8")
|
||||
|
||||
sm.reload()
|
||||
assert sm.count() == 1
|
||||
|
||||
s = sm.get_skill("weather")
|
||||
assert s is not None
|
||||
assert s["summary"] == "查询天气"
|
||||
assert s["is_markdown"] is True
|
||||
assert s["is_executable"] is False
|
||||
assert "查询指定城市天气" in s["content"]
|
||||
|
||||
|
||||
def test_markdown_skill_content_retrieval(isolated_skill_dir, sm):
|
||||
"""sm fixture 构造时目录为空,写入文件后需 reload 才能加载。"""
|
||||
(isolated_skill_dir / "notes.md").write_text("# 笔记技能\n笔记内容ABC", encoding="utf-8")
|
||||
sm.reload()
|
||||
content = sm.get_skill_content("notes")
|
||||
assert "笔记内容ABC" in content
|
||||
|
||||
|
||||
def test_markdown_skill_empty_file_skipped(isolated_skill_dir, sm):
|
||||
"""空文件应被跳过。"""
|
||||
(isolated_skill_dir / "empty.md").write_text("", encoding="utf-8")
|
||||
sm.reload()
|
||||
assert sm.count() == 0
|
||||
|
||||
|
||||
def test_markdown_skill_no_title_uses_filename(isolated_skill_dir, sm):
|
||||
"""无 # 标题时,首行普通文本作为简介。"""
|
||||
(isolated_skill_dir / "plain.md").write_text("这是普通首行\n第二行", encoding="utf-8")
|
||||
sm.reload()
|
||||
s = sm.get_skill("plain")
|
||||
assert s["summary"] == "这是普通首行"
|
||||
|
||||
|
||||
# ==================== 隐藏/下划线文件跳过 ====================
|
||||
|
||||
def test_hidden_file_skipped(isolated_skill_dir, sm):
|
||||
(isolated_skill_dir / ".hidden.md").write_text("# 隐藏", encoding="utf-8")
|
||||
(isolated_skill_dir / "_internal.md").write_text("# 内部", encoding="utf-8")
|
||||
sm.reload()
|
||||
assert sm.count() == 0
|
||||
|
||||
|
||||
# ==================== 搜索 ====================
|
||||
|
||||
def test_search_by_name(isolated_skill_dir, sm):
|
||||
(isolated_skill_dir / "weather.md").write_text("# 天气\n查天气", encoding="utf-8")
|
||||
(isolated_skill_dir / "calculator.md").write_text("# 计算器\n算数", encoding="utf-8")
|
||||
sm.reload()
|
||||
|
||||
results = sm.search("weather")
|
||||
assert len(results) == 1
|
||||
assert results[0]["name"] == "weather"
|
||||
|
||||
|
||||
def test_search_by_summary(isolated_skill_dir, sm):
|
||||
(isolated_skill_dir / "x.md").write_text("# 天气预报技能\n内容", encoding="utf-8")
|
||||
sm.reload()
|
||||
results = sm.search("天气")
|
||||
assert len(results) == 1
|
||||
|
||||
|
||||
def test_search_by_content(isolated_skill_dir, sm):
|
||||
(isolated_skill_dir / "x.md").write_text("# 标题\n独特关键词XYZ123", encoding="utf-8")
|
||||
sm.reload()
|
||||
results = sm.search("XYZ123")
|
||||
assert len(results) == 1
|
||||
|
||||
|
||||
def test_search_no_match(isolated_skill_dir, sm):
|
||||
(isolated_skill_dir / "x.md").write_text("# 标题\n内容", encoding="utf-8")
|
||||
sm.reload()
|
||||
assert sm.search("不存在的关键词") == []
|
||||
|
||||
|
||||
def test_search_case_insensitive(isolated_skill_dir, sm):
|
||||
(isolated_skill_dir / "Weather.md").write_text("# Weather Skill\ncontent", encoding="utf-8")
|
||||
sm.reload()
|
||||
assert len(sm.search("weather")) == 1
|
||||
assert len(sm.search("WEATHER")) == 1
|
||||
|
||||
|
||||
# ==================== get_all / count ====================
|
||||
|
||||
def test_get_all_returns_list(isolated_skill_dir, sm):
|
||||
(isolated_skill_dir / "a.md").write_text("# A\n内容A", encoding="utf-8")
|
||||
(isolated_skill_dir / "b.md").write_text("# B\n内容B", encoding="utf-8")
|
||||
sm.reload()
|
||||
|
||||
all_skills = sm.get_all()
|
||||
assert len(all_skills) == 2
|
||||
assert all("name" in s and "summary" in s and "content" in s for s in all_skills)
|
||||
assert sm.count() == 2
|
||||
|
||||
|
||||
# ==================== reload ====================
|
||||
|
||||
def test_reload_picks_up_new_file(isolated_skill_dir, sm):
|
||||
assert sm.count() == 0
|
||||
(isolated_skill_dir / "new.md").write_text("# 新技能\n内容", encoding="utf-8")
|
||||
n = sm.reload()
|
||||
assert n == 1
|
||||
assert sm.count() == 1
|
||||
|
||||
|
||||
def test_reload_drops_deleted_file(isolated_skill_dir, sm):
|
||||
f = isolated_skill_dir / "temp.md"
|
||||
f.write_text("# 临时\n内容", encoding="utf-8")
|
||||
sm.reload()
|
||||
assert sm.count() == 1
|
||||
|
||||
f.unlink()
|
||||
sm.reload()
|
||||
assert sm.count() == 0
|
||||
|
||||
|
||||
# ==================== 不存在的技能 ====================
|
||||
|
||||
def test_get_nonexistent_skill_returns_none(sm):
|
||||
assert sm.get_skill("不存在") is None
|
||||
|
||||
|
||||
def test_get_content_nonexistent_returns_none(sm):
|
||||
assert sm.get_skill_content("不存在") is None
|
||||
|
||||
|
||||
# ==================== 可执行技能(跨平台:Windows 用 .bat,Unix 用 .sh) ====================
|
||||
|
||||
def test_loads_executable_skill(isolated_skill_dir, sm):
|
||||
"""加载一个可执行技能文件(带 --help)。
|
||||
Windows 下用 .bat(SkillManager 的 _is_executable 识别 .bat),
|
||||
Unix 下用 .sh 并加可执行权限。
|
||||
"""
|
||||
import os as _os
|
||||
if _os.name == 'nt':
|
||||
# Windows:用 .bat 文件
|
||||
skill_file = isolated_skill_dir / "greet.bat"
|
||||
skill_file.write_text(
|
||||
"@echo off\r\n"
|
||||
"echo 打招呼技能 - 输入名字打招呼\r\n",
|
||||
encoding="utf-8"
|
||||
)
|
||||
else:
|
||||
# Unix:用 .sh 并加可执行权限
|
||||
skill_file = isolated_skill_dir / "greet.sh"
|
||||
skill_file.write_text(
|
||||
"#!/bin/bash\n"
|
||||
"echo '打招呼技能 - 输入名字打招呼'\n",
|
||||
encoding="utf-8"
|
||||
)
|
||||
_os.chmod(str(skill_file), 0o755)
|
||||
|
||||
sm.reload()
|
||||
s = sm.get_skill("greet")
|
||||
assert s is not None
|
||||
assert s["is_executable"] is True
|
||||
assert s["is_markdown"] is False
|
||||
assert "打招呼" in s["summary"]
|
||||
+24
-3
@@ -3,8 +3,9 @@ tools - aibox 工具系统基础模块(包)
|
||||
|
||||
抽离自 aibox.py 的工具基类与注册器,供 aibox.py 及后续迁移的
|
||||
工具子类导入。当前包含:
|
||||
- base.py : Tool 抽象基类
|
||||
- registry.py : ToolRegistry 工具注册器
|
||||
- base.py : Tool 抽象基类
|
||||
- registry.py : ToolRegistry 工具注册器
|
||||
- basic_tools.py : 基础工具集(A 批:文件/网页/时间/计算器,无外部子系统依赖)
|
||||
形成单向依赖:common <- tools <- aibox。
|
||||
|
||||
采用"重导出"策略:外部统一用 `from tools import Tool, ToolRegistry`,
|
||||
@@ -13,5 +14,25 @@ tools - aibox 工具系统基础模块(包)
|
||||
|
||||
from .base import Tool
|
||||
from .registry import ToolRegistry
|
||||
from .basic_tools import (
|
||||
FileToolBase,
|
||||
ReadFileTool,
|
||||
WriteFileTool,
|
||||
ListFilesTool,
|
||||
SimpleWebCrawlerTool,
|
||||
GetCurrentTimeTool,
|
||||
CalculatorTool,
|
||||
)
|
||||
|
||||
__all__ = ["Tool", "ToolRegistry"]
|
||||
__all__ = [
|
||||
"Tool",
|
||||
"ToolRegistry",
|
||||
# A 批基础工具
|
||||
"FileToolBase",
|
||||
"ReadFileTool",
|
||||
"WriteFileTool",
|
||||
"ListFilesTool",
|
||||
"SimpleWebCrawlerTool",
|
||||
"GetCurrentTimeTool",
|
||||
"CalculatorTool",
|
||||
]
|
||||
@@ -0,0 +1,463 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
tools.basic_tools - 基础工具集(A 批抽离)
|
||||
|
||||
抽离自 aibox.py 的无外部子系统依赖工具类,仅依赖 common 与标准库。
|
||||
包含:
|
||||
- FileToolBase : 文件操作工具基类
|
||||
- ReadFileTool : 读文件(返回命令行指引)
|
||||
- WriteFileTool : 写文件(返回命令行指引)
|
||||
- ListFilesTool : 列目录(返回命令行指引)
|
||||
- SimpleWebCrawlerTool : 网页访问(返回 curl 指引)
|
||||
- GetCurrentTimeTool : 获取当前时间
|
||||
- CalculatorTool : 数学计算
|
||||
|
||||
依赖方向:common <- tools.basic_tools,不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
import os
|
||||
import queue
|
||||
import shlex
|
||||
from datetime import datetime
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
from common import ConfigManager, constants
|
||||
from .base import Tool
|
||||
|
||||
|
||||
# ==================== 文件操作工具基类 ====================
|
||||
|
||||
class FileToolBase(Tool):
|
||||
"""文件操作工具基类"""
|
||||
|
||||
def __init__(self, logger=None):
|
||||
super().__init__(logger)
|
||||
|
||||
def normalize_path(self, filepath: str) -> str:
|
||||
return os.path.abspath(os.path.expanduser(filepath))
|
||||
|
||||
def check_path_safety(self, filepath: str, check_exists: bool = False) -> Tuple[bool, str, str]:
|
||||
"""检查路径安全性"""
|
||||
try:
|
||||
abs_path = self.normalize_path(filepath)
|
||||
|
||||
if check_exists and not os.path.exists(abs_path):
|
||||
return False, abs_path, f"❌ 文件不存在: {abs_path}"
|
||||
|
||||
# 检查文件大小
|
||||
if os.path.exists(abs_path) and os.path.isfile(abs_path):
|
||||
file_size = os.path.getsize(abs_path)
|
||||
max_size = constants.MAX_FILE_SIZE
|
||||
if file_size > max_size:
|
||||
return False, abs_path, f"❌ 文件太大 ({file_size/1024/1024:.1f}MB)"
|
||||
|
||||
return True, abs_path, ""
|
||||
|
||||
except Exception as e:
|
||||
return False, filepath, f"❌ 路径检查错误: {str(e)}"
|
||||
|
||||
|
||||
# ==================== 文件操作工具(修改版,支持流式输出) ====================
|
||||
|
||||
class ReadFileTool(FileToolBase):
|
||||
"""读取文件内容工具(支持流式输出)"""
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "read_file"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return "读取文件内容,支持指定起始行和结束行,或读取开头/结尾"
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"filepath": {"type": "string"},
|
||||
"encoding": {"type": "string", "default": "utf-8"},
|
||||
"start_line": {"type": "integer", "default": 1},
|
||||
"end_line": {"type": "integer", "default": 0},
|
||||
"max_lines": {"type": "integer", "default": 100},
|
||||
"head": {"type": "boolean", "default": False},
|
||||
"tail": {"type": "boolean", "default": False}
|
||||
},
|
||||
"required": ["filepath"]
|
||||
}
|
||||
|
||||
def execute(self, filepath: str, encoding: str = "utf-8",
|
||||
start_line: int = 1, end_line: int = 0,
|
||||
max_lines: int = 100, head: bool = False,
|
||||
tail: bool = False, **kwargs) -> str:
|
||||
if error := self.validate_params(["filepath"], filepath=filepath):
|
||||
return error
|
||||
|
||||
is_safe, abs_path, error = self.check_path_safety(filepath, check_exists=True)
|
||||
if not is_safe:
|
||||
return error
|
||||
|
||||
# 限制max_lines上限,防止请求超大范围
|
||||
max_lines = min(max_lines, 500)
|
||||
quoted = shlex.quote(abs_path)
|
||||
|
||||
# 构造等价的命令行读取命令
|
||||
if head:
|
||||
cmd = f"head -n {max_lines} {quoted}"
|
||||
elif tail:
|
||||
cmd = f"tail -n {max_lines} {quoted}"
|
||||
elif end_line > 0:
|
||||
cmd = f"sed -n '{start_line},{end_line}p' {quoted}"
|
||||
elif max_lines > 0:
|
||||
end = start_line + max_lines - 1
|
||||
cmd = f"sed -n '{start_line},{end}p' {quoted}"
|
||||
else:
|
||||
cmd = f"sed -n '{start_line},$p' {quoted}"
|
||||
|
||||
# 返回 markdown 提示词,指导智能体用命令行工具读取文件
|
||||
prompt = f"""请使用 `run_command` 工具执行以下命令来读取文件内容,不要直接返回文件内容:
|
||||
|
||||
```bash
|
||||
{cmd}
|
||||
```
|
||||
|
||||
- 若需统计总行数:`wc -l {quoted}`
|
||||
- 若需查看文件编码:`file -b --mime-encoding {quoted}`
|
||||
- 若单行超长,可在命令后追加 `| cut -c 1-10000` 截断到 10000 字符
|
||||
- 文件大小上限 {constants.MAX_FILE_SIZE} 字节,单次最多读取 500 行"""
|
||||
return prompt
|
||||
|
||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||
"""流式执行文件读取"""
|
||||
filepath = kwargs.get("filepath")
|
||||
|
||||
# 发送开始执行的消息
|
||||
output_queue.put(("line", f"\n📄 读取文件: {filepath}\n"))
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
# 执行操作
|
||||
result = self.execute(**kwargs)
|
||||
|
||||
# 发送结果
|
||||
output_queue.put(("line", result + "\n"))
|
||||
output_queue.put(("complete", result))
|
||||
|
||||
|
||||
class WriteFileTool(FileToolBase):
|
||||
"""写入文件工具(支持流式输出)"""
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "write_file"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return "文件写入工具 - 支持覆盖、追加"
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"filepath": {"type": "string"},
|
||||
"content": {"type": "string"},
|
||||
"mode": {"type": "string", "enum": ["w", "a"], "default": "w"},
|
||||
"encoding": {"type": "string", "default": "utf-8"}
|
||||
},
|
||||
"required": ["filepath", "content"]
|
||||
}
|
||||
|
||||
def execute(self, filepath: str = None, content: str = None, mode: str = "w",
|
||||
encoding: str = "utf-8", **kwargs) -> str:
|
||||
if error := self.validate_params(["filepath", "content"], filepath=filepath, content=content):
|
||||
return error
|
||||
|
||||
is_safe, abs_path, error = self.check_path_safety(filepath)
|
||||
if not is_safe:
|
||||
return error
|
||||
|
||||
quoted = shlex.quote(abs_path)
|
||||
directory = os.path.dirname(abs_path)
|
||||
dir_cmd = f"mkdir -p {shlex.quote(directory)}" if directory else ""
|
||||
|
||||
# 内容写入使用 heredoc,避免引号/特殊字符转义问题
|
||||
redirect = ">" if mode == "w" else ">>"
|
||||
write_cmd = f"cat {redirect} {quoted} <<'AIBOX_EOF'\n{content}\nAIBOX_EOF"
|
||||
if dir_cmd:
|
||||
write_cmd = f"{dir_cmd} && {write_cmd}"
|
||||
|
||||
# 返回 markdown 提示词,指导智能体用命令行工具写入文件
|
||||
prompt = f"""请使用 `run_command` 工具执行以下命令来写入文件,不要在此工具中直接写入:
|
||||
|
||||
```bash
|
||||
{write_cmd}
|
||||
```
|
||||
|
||||
- 如遇 heredoc 内容含特殊字符导致解析失败,可用 `printf '%s' '内容' > {quoted}` 覆盖写入
|
||||
- 若需检查文件大小:`wc -c {quoted}`
|
||||
- 文件大小上限 {constants.MAX_FILE_SIZE} 字节"""
|
||||
return prompt
|
||||
|
||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||
"""流式执行文件写入"""
|
||||
filepath = kwargs.get("filepath")
|
||||
|
||||
# 发送开始执行的消息
|
||||
output_queue.put(("line", f"\n📝 写入文件: {filepath}\n"))
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
# 执行操作
|
||||
result = self.execute(**kwargs)
|
||||
|
||||
# 发送结果
|
||||
output_queue.put(("line", result + "\n"))
|
||||
output_queue.put(("complete", result))
|
||||
|
||||
|
||||
class ListFilesTool(FileToolBase):
|
||||
"""列出目录内容工具(支持流式输出)"""
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "list_files"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return "列出指定目录下的文件和子目录"
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"directory": {"type": "string", "default": "."},
|
||||
"show_hidden": {"type": "boolean", "default": False}
|
||||
}
|
||||
}
|
||||
|
||||
def execute(self, directory: str = ".", show_hidden: bool = False, **kwargs) -> str:
|
||||
is_safe, abs_path, error = self.check_path_safety(directory)
|
||||
if not is_safe:
|
||||
return error
|
||||
|
||||
if not os.path.exists(abs_path):
|
||||
return f"❌ 目录不存在: {abs_path}"
|
||||
|
||||
if not os.path.isdir(abs_path):
|
||||
return f"❌ 路径不是目录: {abs_path}"
|
||||
|
||||
quoted = shlex.quote(abs_path)
|
||||
flag = "-la" if show_hidden else "-l"
|
||||
|
||||
# 返回 markdown 提示词,指导智能体用命令行工具列目录
|
||||
prompt = f"""请使用 `run_command` 工具执行以下命令来列出目录内容,不要直接在此工具中列出:
|
||||
|
||||
```bash
|
||||
ls {flag} {quoted}
|
||||
```
|
||||
|
||||
- 若需区分目录/文件并显示大小:`ls -lhF {quoted}`
|
||||
- 若需统计目录总大小:`du -sh {quoted}`
|
||||
- 若需递归列出:`find {quoted} -maxdepth 2`
|
||||
- 若需了解文件类型:`file {quoted}/*`"""
|
||||
return prompt
|
||||
|
||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||
"""流式执行列出目录"""
|
||||
directory = kwargs.get("directory", ".")
|
||||
|
||||
# 发送开始执行的消息
|
||||
output_queue.put(("line", f"\n📁 列出目录: {directory}\n"))
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
# 执行操作
|
||||
result = self.execute(**kwargs)
|
||||
|
||||
# 发送结果
|
||||
output_queue.put(("line", result + "\n"))
|
||||
output_queue.put(("complete", result))
|
||||
|
||||
|
||||
# ==================== 网页访问工具 ====================
|
||||
|
||||
class SimpleWebCrawlerTool(Tool):
|
||||
"""简化版网页爬虫工具(返回 markdown 提示词,指导改用 curl 命令行)"""
|
||||
|
||||
def __init__(self, logger=None):
|
||||
super().__init__(logger)
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "web"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return """简单易用的网页访问工具 - 一句话搞定!
|
||||
支持:获取网页、提取内容、搜索关键词、下载文件"""
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"action": {
|
||||
"type": "string",
|
||||
"enum": ["get", "text", "title", "links", "search"]
|
||||
},
|
||||
"url": {"type": "string"},
|
||||
"keyword": {"type": "string"}
|
||||
},
|
||||
"required": ["action", "url"]
|
||||
}
|
||||
|
||||
def _proxy_curl_args(self) -> str:
|
||||
"""根据当前代理配置构造 curl --socks5 参数"""
|
||||
proxies = ConfigManager().get_proxy()
|
||||
if not proxies:
|
||||
return ""
|
||||
proxy_url = proxies.get("https") or proxies.get("http") or ""
|
||||
host_port = proxy_url.split("://", 1)[-1]
|
||||
if "@" in host_port:
|
||||
host_port = host_port.split("@", 1)[-1]
|
||||
return f"--socks5 {shlex.quote(host_port)}"
|
||||
|
||||
def execute(self, action: str, url: str, **kwargs) -> str:
|
||||
if action not in ("get", "text", "title", "links", "search"):
|
||||
return f"❌ 未知操作: {action}"
|
||||
|
||||
quoted_url = shlex.quote(url)
|
||||
timeout = constants.WEB_TIMEOUT
|
||||
user_agent = shlex.quote(constants.WEB_USER_AGENT)
|
||||
proxy = self._proxy_curl_args()
|
||||
base = f"curl -sL --max-time {timeout} -A {user_agent} {proxy}".strip()
|
||||
|
||||
if action == "get":
|
||||
cmd = f"{base} -w '\\nHTTP_CODE:%{{http_code}}' {quoted_url}"
|
||||
extra = "如需查看响应头:`curl -sIL ...`;如需保存到文件:`curl -sL ... -o 文件名`"
|
||||
elif action == "text":
|
||||
cmd = f"{base} {quoted_url} | sed -r 's/<[^>]+>//g'"
|
||||
extra = "若系统装有 lynx,可用 `lynx -dump -stdin` 提取纯文本更准确"
|
||||
elif action == "title":
|
||||
cmd = f"{base} {quoted_url} | grep -oP '<title>\\K[^<]+'"
|
||||
extra = "无输出则表示未找到标题"
|
||||
elif action == "links":
|
||||
cmd = f"{base} {quoted_url} | grep -oP 'href=\"\\K[^\"]+'"
|
||||
extra = "输出所有链接 URL,可配合 `head` 限制数量"
|
||||
else: # search
|
||||
keyword = kwargs.get("keyword")
|
||||
if not keyword:
|
||||
return "❌ 需要提供 keyword 参数"
|
||||
cmd = f"{base} {quoted_url} | grep -c {shlex.quote(keyword)}"
|
||||
extra = f"返回匹配行数,大于 0 即表示找到关键词 '{keyword}'"
|
||||
|
||||
# 返回 markdown 提示词,指导智能体用 curl 完成网页访问
|
||||
prompt = f"""请使用 `run_command` 工具执行以下命令来完成网页操作,不要直接返回网页内容:
|
||||
|
||||
```bash
|
||||
{cmd}
|
||||
```
|
||||
|
||||
{extra}
|
||||
- 下载大小限制 {constants.WEB_MAX_DOWNLOAD_SIZE} 字节,可在命令后追加 `| head -c {constants.WEB_MAX_DOWNLOAD_SIZE}` 限制
|
||||
- 当前已按配置自动附带代理:{proxy or '无(如需代理请先调用 set_proxy 工具)'}"""
|
||||
return prompt
|
||||
|
||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||
"""流式执行网页操作"""
|
||||
action = kwargs.get("action")
|
||||
url = kwargs.get("url")
|
||||
|
||||
# 发送开始执行的消息
|
||||
output_queue.put(("line", f"\n🌐 执行网页操作: {action} - {url}\n"))
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
# 执行操作
|
||||
result = self.execute(**kwargs)
|
||||
|
||||
# 发送结果
|
||||
output_queue.put(("line", result + "\n"))
|
||||
output_queue.put(("complete", result))
|
||||
|
||||
|
||||
# ==================== 系统工具(修改版,支持流式输出) ====================
|
||||
|
||||
class GetCurrentTimeTool(Tool):
|
||||
"""获取当前时间工具(支持流式输出)"""
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "get_current_time"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return "获取当前日期和时间"
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"format": {"type": "string", "enum": ["full", "date", "time"], "default": "full"}
|
||||
}
|
||||
}
|
||||
|
||||
def execute(self, format: str = "full", **kwargs) -> str:
|
||||
now = datetime.now()
|
||||
|
||||
if format == "date":
|
||||
result = now.strftime("%Y-%m-%d")
|
||||
elif format == "time":
|
||||
result = now.strftime("%H:%M:%S")
|
||||
else:
|
||||
result = now.strftime(constants.DATETIME_FORMAT)
|
||||
|
||||
return self.format_result(True, "当前时间", {"时间": result})
|
||||
|
||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||
"""流式执行获取时间"""
|
||||
output_queue.put(("line", "\n🕐 获取当前时间\n"))
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
result = self.execute(**kwargs)
|
||||
|
||||
output_queue.put(("line", result + "\n"))
|
||||
output_queue.put(("complete", result))
|
||||
|
||||
|
||||
class CalculatorTool(Tool):
|
||||
"""计算器工具(支持流式输出)"""
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "calculator"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return "执行数学计算"
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"a": {"type": "number"},
|
||||
"b": {"type": "number"},
|
||||
"operation": {"type": "string", "enum": ["add", "subtract", "multiply", "divide"]}
|
||||
},
|
||||
"required": ["a", "b", "operation"]
|
||||
}
|
||||
|
||||
def execute(self, a: float, b: float, operation: str, **kwargs) -> str:
|
||||
operations = {
|
||||
"add": ("+", a + b),
|
||||
"subtract": ("-", a - b),
|
||||
"multiply": ("*", a * b),
|
||||
"divide": ("/", a / b if b != 0 else None)
|
||||
}
|
||||
|
||||
if operation not in operations:
|
||||
return f"❌ 未知运算: {operation}"
|
||||
|
||||
op_symbol, result = operations[operation]
|
||||
|
||||
if result is None:
|
||||
return "❌ 错误:除数不能为0"
|
||||
|
||||
data = {
|
||||
"表达式": f"{a} {op_symbol} {b}",
|
||||
"结果": result
|
||||
}
|
||||
|
||||
return self.format_result(True, "计算成功", data)
|
||||
|
||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||
"""流式执行计算"""
|
||||
output_queue.put(("line", "\n🧮 执行计算\n"))
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
result = self.execute(**kwargs)
|
||||
|
||||
output_queue.put(("line", result + "\n"))
|
||||
output_queue.put(("complete", result))
|
||||
Reference in new issue
Block a user