zikai/zWorkFlow
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(mcp): get_chat_images/get_voice_messages 加 offset/time_range

Browse Source 
\`get_chat_images\` 和 \`get_voice_messages\` 仅有 \`limit\`, 接口与
\`get_chat_history\` / \`search_messages\` (\`offset\` + \`start_time\` +
\`end_time\`) 不对齐:

1. 查不了"某段时间内的图片/语音"
2. 不支持分页, 单次取 \`limit=1000\` 一次性拉
3. LLM 用同样模式调不同工具时签名不一致, 容易出错

两个工具各加 3 个可选参数:
- \`offset: int = 0\`
- \`start_time: str = ""\`
- \`end_time: str = ""\`

复用上游已有的 \`_validate_pagination\` + \`_parse_time_range\` helpers。

- 输入校验失败立即报错 (offset 负数 / 时间格式错 / start > end)
- 每 shard 拉 \`limit + offset\` 张候选, 合并后全局 \`create_time DESC\`
  排序, 切片 \`[offset : offset + limit]\` 出本页
- 单 shard 凑得起本页, 避免某 shard 缺数据时本页变短
- header 显示 offset/limit 和时间范围 (传了才显示)

- 加 \`start_ts=None\` / \`end_ts=None\` 参数
- SQL 动态拼 \`create_time >= ?\` / \`<= ?\` clause
- 不传时序参数完全等价旧行为 (向后兼容)

- 同样 3 个参数 + \`_validate_pagination\` + \`_parse_time_range\`
- VoiceInfo 表 SQL 动态拼 \`chat_name_id = ? AND create_time ?...\`
- 多 shard 各取 \`limit + offset\` 后合并切片

新增 \`tests/test_chat_images_query_align.py\` 8 个 case:
- offset 负数报错
- start > end 报错
- candidate_limit = limit + offset (shard 调用确认)
- 时间参数正确解析为 unix 秒并透传
- offset=2, limit=2 切到全局排序后第 3-4 张
- header 包含时间范围
- header 包含 offset/limit
- 默认调用(不传新参)行为与旧接口一致

修改 \`tests/test_get_chat_images_multishard.py\` 的 fake_list 签名:
- 旧: \`(db_path, table_name, username, lim)\` 位置参
- 新: \`(db_path, table_name, username, limit=20, start_ts=None, end_ts=None)\`
- 既支持旧调用模式 (kwargs), 也兼容新签名

全量 \`pytest tests/\` 208/208 通过。

3 个可选参数全部带默认值 → 既有调用方零修改。

shard candidate=\`limit+offset\` 的成本: 大 offset 时单 shard 请求量
增大。但 image/voice 表每 chat 单 shard 一般 < 10K 条, 实际 cost 可
忽略。如果将来要做"翻 100 页"级深翻, 可以加 keyset pagination, 现
在 offset 模式与 \`get_chat_history\` 一致即可。

与 #103 / #104 触碰同一文件, 合并顺序无所谓 — 后合的 rebase 即可。
This commit is contained in:
Belugary
2026-05-14 12:42:28 +08:00
committed by ylytdeng
parent 6606122c86
commit 5bc275b81c
4 changed files with 171 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
decode_image.py
View File

@@ -450,17 +450,30 @@ class ImageResolver:
'size': os.path.getsize(final_path),
}
def list_chat_images(self, db_path, table_name, username, limit=20):
"""列出某个聊天中的所有图片消息"""
def list_chat_images(self, db_path, table_name, username, limit=20, start_ts=None, end_ts=None):
"""列出某个聊天中的所有图片消息
可选 start_ts / end_ts (unix 秒) 过滤时间范围。
"""
clauses = ['local_type = 3']
params = []
if start_ts is not None:
clauses.append('create_time >= ?')
params.append(start_ts)
if end_ts is not None:
clauses.append('create_time <= ?')
params.append(end_ts)
params.append(limit)
where_sql = ' AND '.join(clauses)
conn = sqlite3.connect(db_path)
try:
rows = conn.execute(f"""
SELECT local_id, create_time
FROM [{table_name}]
WHERE local_type = 3
WHERE {where_sql}
ORDER BY create_time DESC
LIMIT ?
""", (limit,)).fetchall()
""", params).fetchall()
except Exception as e:
conn.close()
return []

71
mcp_server.py
View File

@@ -2830,7 +2830,7 @@ def decode_transfer(chat_name: str, local_id: int, create_time: int = 0) -> str:
@mcp.tool()
def get_chat_images(chat_name: str, limit: int = 20) -> str:
def get_chat_images(chat_name: str, limit: int = 20, offset: int = 0, start_time: str = "", end_time: str = "") -> str:
"""列出某个聊天中的图片消息。
返回图片的时间、local_id、MD5、文件大小等信息。
@@ -2839,7 +2839,16 @@ def get_chat_images(chat_name: str, limit: int = 20) -> str:
Args:
chat_name: 聊天对象的名字、备注名或wxid
limit: 返回数量，默认20
offset: 分页偏移量，默认0
start_time: 起始时间，支持 YYYY-MM-DD / YYYY-MM-DD HH:MM / YYYY-MM-DD HH:MM:SS
end_time: 结束时间，支持 YYYY-MM-DD / YYYY-MM-DD HH:MM / YYYY-MM-DD HH:MM:SS
"""
try:
_validate_pagination(limit, offset)
start_ts, end_ts = _parse_time_range(start_time, end_time)
except ValueError as e:
return f"错误: {e}"
username = resolve_username(chat_name)
if not username:
return f"找不到聊天对象: {chat_name}"
@@ -2854,12 +2863,15 @@ def get_chat_images(chat_name: str, limit: int = 20) -> str:
if not shards:
return f"找不到 {display_name} 的消息记录"
# 每个 shard 取 limit 张, 合并后按 create_time DESC 全局排序, 取最新 limit 张。
# 单 shard 至少够本次返回, 避免一个 shard 凑不出 limit 时其他 shard 没机会贡献。
# 每个 shard 取 limit+offset 张候选, 合并后按 create_time DESC 全局排序, 切片
# [offset : offset+limit] 出本页。单 shard 至少凑得起本页, 避免某 shard 缺数据
# 时本页变短。
candidate_limit = limit + offset
all_images = []
for shard in shards:
shard_images = _image_resolver.list_chat_images(
shard['db_path'], shard['table_name'], username, limit
shard['db_path'], shard['table_name'], username,
limit=candidate_limit, start_ts=start_ts, end_ts=end_ts,
)
all_images.extend(shard_images)
@@ -2867,10 +2879,10 @@ def get_chat_images(chat_name: str, limit: int = 20) -> str:
return f"{display_name} 无图片消息"
all_images.sort(key=lambda img: img['create_time'], reverse=True)
images = all_images[:limit]
paged = all_images[offset:offset + limit]
lines = []
for img in images:
for img in paged:
time_str = datetime.fromtimestamp(img['create_time']).strftime('%Y-%m-%d %H:%M')
line = f"[{time_str}] local_id={img['local_id']}"
if img.get('md5'):
@@ -2882,7 +2894,10 @@ def get_chat_images(chat_name: str, limit: int = 20) -> str:
line += " (无资源信息)"
lines.append(line)
return f"{display_name} 的 {len(lines)} 张图片:\n\n" + "\n".join(lines) + _pagination_hint(len(lines), limit, 0)
header = f"{display_name} 的 {len(lines)} 张图片（offset={offset}, limit={limit}）"
if start_time or end_time:
header += f"\n时间范围: {start_time or '最早'} ~ {end_time or '最新'}"
return header + ":\n\n" + "\n".join(lines) + _pagination_hint(len(lines), limit, offset)
# ============ 语音解密 ============
@@ -2952,7 +2967,7 @@ def _silk_to_wav(voice_data, create_time, username, local_id):
@mcp.tool()
def get_voice_messages(chat_name: str, limit: int = 20) -> str:
def get_voice_messages(chat_name: str, limit: int = 20, offset: int = 0, start_time: str = "", end_time: str = "") -> str:
"""列出某个聊天中的语音消息。
返回语音的时间、local_id 和大小，可配合 decode_voice 工具解码。
@@ -2960,7 +2975,16 @@ def get_voice_messages(chat_name: str, limit: int = 20) -> str:
Args:
chat_name: 聊天对象的名字、备注名或wxid
limit: 返回数量，默认20
offset: 分页偏移量，默认0
start_time: 起始时间，支持 YYYY-MM-DD / YYYY-MM-DD HH:MM / YYYY-MM-DD HH:MM:SS
end_time: 结束时间，支持 YYYY-MM-DD / YYYY-MM-DD HH:MM / YYYY-MM-DD HH:MM:SS
"""
try:
_validate_pagination(limit, offset)
start_ts, end_ts = _parse_time_range(start_time, end_time)
except ValueError as e:
return f"错误: {e}"
username = resolve_username(chat_name)
if not username:
return f"找不到聊天对象: {chat_name}"
@@ -2971,31 +2995,48 @@ def get_voice_messages(chat_name: str, limit: int = 20) -> str:
if not MEDIA_DB_KEYS:
return "找不到 media DB"
# 从每个分片各取最多 limit 条后合并再截断：分片若有时间重叠也不会漏最新消息
# 每分片各取 limit+offset 条候选, 合并后全局排序切片 [offset:offset+limit] 出本页。
candidate_limit = limit + offset
clauses = ['chat_name_id = ?']
if start_ts is not None:
clauses.append('create_time >= ?')
if end_ts is not None:
clauses.append('create_time <= ?')
where_sql = ' AND '.join(clauses)
rows = []
for media_db in _iter_media_db_paths():
with closing(sqlite3.connect(media_db)) as conn:
chat_name_id = _get_chat_name_id(conn, username)
if chat_name_id is None:
continue
params = [chat_name_id]
if start_ts is not None:
params.append(start_ts)
if end_ts is not None:
params.append(end_ts)
params.append(candidate_limit)
rows.extend(conn.execute(
"SELECT local_id, create_time, length(voice_data) FROM VoiceInfo "
"WHERE chat_name_id = ? ORDER BY create_time DESC LIMIT ?",
(chat_name_id, limit),
f"SELECT local_id, create_time, length(voice_data) FROM VoiceInfo "
f"WHERE {where_sql} ORDER BY create_time DESC LIMIT ?",
params,
).fetchall())
if not rows:
return f"{display_name} 无语音消息"
rows.sort(key=lambda r: r[1], reverse=True)
rows = rows[:limit]
paged = rows[offset:offset + limit]
lines = []
for local_id, create_time, size in rows:
for local_id, create_time, size in paged:
time_str = datetime.fromtimestamp(create_time).strftime('%Y-%m-%d %H:%M')
lines.append(f"[{time_str}] local_id={local_id} {size/1024:.0f}KB")
return f"{display_name} 的 {len(lines)} 条语音消息:\n\n" + "\n".join(lines) + _pagination_hint(len(lines), limit, 0)
header = f"{display_name} 的 {len(lines)} 条语音消息（offset={offset}, limit={limit}）"
if start_time or end_time:
header += f"\n时间范围: {start_time or '最早'} ~ {end_time or '最新'}"
return header + ":\n\n" + "\n".join(lines) + _pagination_hint(len(lines), limit, offset)
@mcp.tool()

97
tests/test_chat_images_query_align.py Normal file
View File

@@ -0,0 +1,97 @@
"""测试 get_chat_images 新增的 offset / start_time / end_time 参数。"""
import os
import sys
from unittest.mock import patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import mcp_server
def _img(local_id, create_time, md5=None, size=None):
info = {'local_id': local_id, 'create_time': create_time, 'md5': md5}
if size is not None:
info['size'] = size
return info
def _run_with(shard_images_map, **kwargs):
"""Helper: stub collaborators and call get_chat_images with new kwargs."""
shards = [{'db_path': k, 'table_name': 'Msg_x'} for k in shard_images_map]
captured = {'calls': []}
def fake_list(db_path, table_name, username, limit=20, start_ts=None, end_ts=None):
captured['calls'].append({
'db_path': db_path, 'limit': limit, 'start_ts': start_ts, 'end_ts': end_ts,
})
return shard_images_map.get(db_path, [])
with patch.object(mcp_server, 'resolve_username', return_value='wxid_demo'), \
patch.object(mcp_server, 'get_contact_names', return_value={'wxid_demo': 'Demo'}), \
patch.object(mcp_server, '_find_msg_tables_for_user', return_value=shards), \
patch.object(mcp_server._image_resolver, 'list_chat_images', side_effect=fake_list):
return mcp_server.get_chat_images('Demo', **kwargs), captured
def test_invalid_offset_returns_error():
out, _ = _run_with({}, offset=-1)
assert '错误' in out
def test_invalid_time_range_returns_error():
"""start_time 晚于 end_time 应报错。"""
out, _ = _run_with({}, start_time='2026-05-10', end_time='2026-05-01')
assert '错误' in out
def test_candidate_limit_includes_offset():
"""每 shard 拉 limit+offset 张候选, 保证全局分页能切到正确的页。"""
_, captured = _run_with({'/a': [_img(1, 1000)]}, limit=5, offset=10)
assert captured['calls'][0]['limit'] == 15
def test_start_end_ts_forwarded_to_shard_query():
"""start_time / end_time 解析为 unix 秒后透传给 shard 查询。"""
_, captured = _run_with(
{'/a': []},
start_time='2026-05-01',
end_time='2026-05-31',
)
call = captured['calls'][0]
assert call['start_ts'] is not None
assert call['end_ts'] is not None
assert call['start_ts'] < call['end_ts']
def test_offset_slices_paged_window():
"""offset=2, limit=2 取全局排序后第 3-4 张图片。"""
shard_a = [_img(1, 1100), _img(2, 1000)]
shard_b = [_img(3, 1300), _img(4, 1200)]
out, _ = _run_with({'/a': shard_a, '/b': shard_b}, limit=2, offset=2)
# 全局排序后顺序: 1300, 1200, 1100, 1000 → 第 3-4 是 1100, 1000 → local_id 1, 2
assert 'local_id=1' in out
assert 'local_id=2' in out
assert 'local_id=3' not in out
assert 'local_id=4' not in out
def test_header_shows_time_range_when_given():
shard_a = [_img(1, 1000, md5='abc')]
out, _ = _run_with({'/a': shard_a}, start_time='2026-05-01')
assert '时间范围' in out
assert '2026-05-01' in out
def test_header_shows_offset_limit():
shard_a = [_img(1, 1000, md5='abc')]
out, _ = _run_with({'/a': shard_a}, limit=10, offset=20)
assert 'offset=20' in out
assert 'limit=10' in out
def test_default_behavior_unchanged():
"""不传新参数时行为与旧接口一致 — offset=0 切片就是 [:limit]。"""
shard_a = [_img(1, 1100, md5='a1'), _img(2, 1000, md5='a2')]
out, _ = _run_with({'/a': shard_a})
assert 'local_id=1' in out
assert 'local_id=2' in out

2
tests/test_get_chat_images_multishard.py
View File

@@ -33,7 +33,7 @@ class GetChatImagesMultiShardTests(unittest.TestCase):
def _run(self, shards, shard_images_map, limit=20):
"""Helper: stub the two collaborators and call the tool."""
def fake_list(db_path, table_name, username, lim):
def fake_list(db_path, table_name, username, limit=20, start_ts=None, end_ts=None):
return shard_images_map.get(db_path, [])
with patch.object(mcp_server, "_find_msg_tables_for_user",