test_time_decay_function.py•7.91 kB
#!/usr/bin/env python3
"""
时间衰减函数单元测试
测试时间感知检索排序的核心算法
"""
import math
from datetime import datetime, timedelta
from typing import List, Dict, Any
def apply_time_decay_ranking(results: List[Dict[str, Any]], tau_days: float = 30.0) -> List[Dict[str, Any]]:
"""应用时间衰减函数到检索结果，实现时间感知排序"""
try:
current_time = datetime.now()
for result in results:
metadata = result.get('metadata', {})
timestamp_str = metadata.get('timestamp')
if timestamp_str:
try:
# 解析时间戳
if timestamp_str.endswith('Z'):
# ISO格式带Z后缀
memory_time = datetime.fromisoformat(timestamp_str.replace('Z', '+00:00'))
elif '+' in timestamp_str or timestamp_str.count('-') > 2:
# ISO格式带时区信息
memory_time = datetime.fromisoformat(timestamp_str)
else:
# 简单ISO格式
memory_time = datetime.fromisoformat(timestamp_str)
# 计算时间差（天数）
time_delta = (current_time - memory_time.replace(tzinfo=None)).total_seconds() / (24 * 3600)
# 计算时间衰减因子：exp(-Δt/τ)
time_decay_factor = math.exp(-time_delta / tau_days)
# 获取当前分数
current_score = result.get('score', 0.5)
# 应用时间衰减：rank = cos_sim × exp(-Δt/τ)
time_aware_score = current_score * time_decay_factor
result['score'] = time_aware_score
result['time_decay_factor'] = time_decay_factor
result['days_ago'] = time_delta
print(f" ⏰ 记忆#{result.get('id')}: {time_delta:.1f}天前, "
f"衰减={time_decay_factor:.3f}, score={current_score:.3f}→{time_aware_score:.3f}")
except (ValueError, TypeError) as e:
# 时间戳解析失败，保持原分数
print(f" ⚠️ 记忆#{result.get('id')} 时间戳解析失败: {timestamp_str}, 错误: {e}")
result['time_decay_factor'] = 1.0
result['days_ago'] = float('inf')
else:
# 没有时间戳，设置默认衰减因子
result['time_decay_factor'] = 0.1 # 给没有时间戳的记忆较低权重
result['days_ago'] = float('inf')
print(f" ⚠️ 记忆#{result.get('id')} 缺少时间戳，使用默认衰减因子0.1")
# 按时间感知分数重新排序
results.sort(key=lambda x: x.get('score', 0), reverse=True)
return results
except Exception as e:
print(f"❌ 时间感知排序失败: {e}")
return results
def test_time_decay_function():
"""测试时间衰减函数"""
print("🧪 时间衰减函数测试")
print("=" * 40)
# 创建测试数据
test_results = [
{
'id': '1',
'content': '今天学习了Python编程',
'score': 0.8,
'metadata': {
'timestamp': datetime.now().isoformat()
}
},
{
'id': '2',
'content': '一周前学习了机器学习',
'score': 0.9, # 更高的原始分数
'metadata': {
'timestamp': (datetime.now() - timedelta(days=7)).isoformat()
}
},
{
'id': '3',
'content': '一个月前学习了深度学习',
'score': 0.85,
'metadata': {
'timestamp': (datetime.now() - timedelta(days=30)).isoformat()
}
},
{
'id': '4',
'content': '三个月前学习了AI基础',
'score': 0.95, # 最高的原始分数
'metadata': {
'timestamp': (datetime.now() - timedelta(days=90)).isoformat()
}
},
{
'id': '5',
'content': '没有时间戳的记忆',
'score': 0.7,
'metadata': {}
}
]
print("\n📊 原始结果（按原始分数排序）:")
original_results = sorted(test_results, key=lambda x: x['score'], reverse=True)
for i, result in enumerate(original_results, 1):
print(f" {i}. {result['content']} (分数: {result['score']:.3f})")
print("\n⏰ 应用时间感知排序 (τ=30天):")
time_aware_results = apply_time_decay_ranking(test_results.copy(), tau_days=30.0)
print("\n📈 时间感知排序结果:")
for i, result in enumerate(time_aware_results, 1):
days_ago = result.get('days_ago', float('inf'))
time_decay_factor = result.get('time_decay_factor', 1.0)
score = result.get('score', 0)
if days_ago != float('inf'):
print(f" {i}. {result['content']} (最终分数: {score:.3f}, {days_ago:.1f}天前, 时间权重: {time_decay_factor:.3f})")
else:
print(f" {i}. {result['content']} (最终分数: {score:.3f}, 无时间戳, 时间权重: {time_decay_factor:.3f})")
print("\n🔍 验证时间感知效果:")
# 检查最新的记忆是否排在前面
if time_aware_results:
first_result = time_aware_results[0]
first_days_ago = first_result.get('days_ago', float('inf'))
if first_days_ago < 1: # 最新的记忆
print("✅ 时间感知排序正常 - 最新记忆排在第一位")
else:
print(f"⚠️ 第一位是 {first_days_ago:.1f} 天前的记忆")
# 检查时间衰减是否正确应用
recent_found = False
old_found = False
for result in time_aware_results:
days_ago = result.get('days_ago', float('inf'))
if days_ago < 7:
recent_found = True
elif days_ago > 60:
old_found = True
if recent_found and old_found:
print("✅ 时间范围覆盖正常 - 包含新旧记忆")
return True
def test_different_tau_values():
"""测试不同τ值的影响"""
print("\n🎛️ 测试不同τ值的影响")
print("-" * 30)
# 创建简单测试数据
test_data = [
{
'id': '1',
'content': '今天的记忆',
'score': 0.5,
'metadata': {'timestamp': datetime.now().isoformat()}
},
{
'id': '2',
'content': '30天前的记忆',
'score': 0.5,
'metadata': {'timestamp': (datetime.now() - timedelta(days=30)).isoformat()}
}
]
tau_values = [7.0, 30.0, 90.0]
for tau in tau_values:
print(f"\nτ = {tau}天:")
results = apply_time_decay_ranking(test_data.copy(), tau_days=tau)
for result in results:
days_ago = result.get('days_ago', 0)
time_decay_factor = result.get('time_decay_factor', 1.0)
score = result.get('score', 0)
print(f" {result['content']}: {days_ago:.0f}天前, 权重={time_decay_factor:.3f}, 最终分数={score:.3f}")
if __name__ == "__main__":
print("🚀 开始时间衰减函数测试")
# 测试核心函数
success = test_time_decay_function()
# 测试不同参数
test_different_tau_values()
if success:
print("\n🎉 时间衰减函数测试通过!")
else:
print("\n💥 测试失败!")