import json
import os

trace_dir = '/Users/elksmmx/Desktop/agent 2.10/Agent/examples/find knowledge/.trace/7ff963bc-3106-47fc-9725-b712f3e0d7d5/messages'

files = sorted([f for f in os.listdir(trace_dir) if f.endswith('.json')])
print(f"总消息数: {len(files)}\n")

all_msgs = []
user_assistant = []

for f in files:
    with open(os.path.join(trace_dir, f)) as fp:
        d = json.load(fp)
        seq = d['sequence']
        role = d['role']
        all_msgs.append((seq, role))
        if role in ('user', 'assistant'):
            user_assistant.append(seq)

# 显示前60条
print("前60条消息的role分布:")
for seq, role in all_msgs[:60]:
    print(f"{seq:3d} {role:10s}")

print(f"\n统计:")
print(f"- 前60条中user/assistant消息: {len([s for s in user_assistant if s <= 60])}个")
print(f"- 前60条中tool消息: {len([s for s, r in all_msgs[:60] if r == 'tool'])}个")

# 计算修改前后的缓存点位置
print(f"\n修改前（包含tool）:")
cacheable_old = [(i, seq, role) for i, (seq, role) in enumerate(all_msgs[:60]) if role in ('user', 'assistant', 'tool')]
print(f"- 可缓存消息数: {len(cacheable_old)}")
if len(cacheable_old) >= 20:
    print(f"- 第20条缓存点: seq={cacheable_old[19][1]}, role={cacheable_old[19][2]}")
if len(cacheable_old) >= 40:
    print(f"- 第40条缓存点: seq={cacheable_old[39][1]}, role={cacheable_old[39][2]}")
if len(cacheable_old) >= 60:
    print(f"- 第60条缓存点: seq={cacheable_old[59][1]}, role={cacheable_old[59][2]}")

print(f"\n修改后（只有user/assistant）:")
cacheable_new = [(i, seq, role) for i, (seq, role) in enumerate(all_msgs[:60]) if role in ('user', 'assistant')]
print(f"- 可缓存消息数: {len(cacheable_new)}")
if len(cacheable_new) >= 20:
    print(f"- 第20条缓存点: seq={cacheable_new[19][1]}, role={cacheable_new[19][2]}")
if len(cacheable_new) >= 40:
    print(f"- 第40条缓存点: seq={cacheable_new[39][1]}, role={cacheable_new[39][2]}")
if len(cacheable_new) >= 60:
    print(f"- 第60条缓存点: seq={cacheable_new[59][1]}, role={cacheable_new[59][2]}")